implement env var based node and file limits on cortex and axon

synapse/axon.py

@@ -1,3 +1,4 @@
+import os
 import asyncio
 import hashlib
 import logging
@@ -71,7 +72,6 @@ async def put(self):
         await self._save()
         return
 
-
 class AxonHttpHasV1(s_httpapi.Handler):
 
     async def get(self, sha256):
@@ -80,7 +80,6 @@ async def get(self, sha256):
         resp = await self.cell.has(s_common.uhex(sha256))
         return self.sendRestRetn(resp)
 
-
 class AxonHttpDownloadV1(s_httpapi.Handler):
 
     async def get(self, sha256):
@@ -242,13 +241,39 @@ async def __anit__(self, dirn, conf=None):  # type: ignore
         self.axonmetrics.setdefault('size:bytes', 0)
         self.axonmetrics.setdefault('file:count', 0)
 
+        self._initAxonLimits()
         self.addHealthFunc(self._axonHealth)
 
         # modularize blob storage
         await self._initBlobStor()
 
         self._initAxonHttpApi()
 
+    def _reqBelowLimit(self):
+
+        if (self.maxbytes is not None and
+            self.maxbytes <= self.axonmetrics.get('size:bytes')):
+            mesg = f'Axon is at size:bytes limit: {self.maxbytes}'
+            raise s_exc.HitLimit(mesg=mesg)
+
+        if (self.maxcount is not None and
+            self.maxcount <= self.axonmetrics.get('file:count')):
+            mesg = f'Axon is at file:count limit: {self.maxcount}'
+            raise s_exc.HitLimit(mesg=mesg)
+
+    def _initAxonLimits(self):
+
+        self.maxbytes = None
+        self.maxcount = None
+
+        maxcount = os.getenv('AXON_MAX_COUNT')
+        if maxcount is not None:
+            self.maxcount = int(maxcount, 0)
+
+        maxbytes = os.getenv('AXON_MAX_BYTES')
+        if maxbytes is not None:
+            self.maxbytes = int(maxbytes, 0)
+
     async def _axonHealth(self, health):
         health.update('axon', 'nominal', '', data=await self.metrics())
 
@@ -310,6 +335,7 @@ async def metrics(self):
 
     async def save(self, sha256, genr):
 
+        self._reqBelowLimit()
         byts = self.axonslab.get(sha256, db=self.sizes)
         if byts is not None:
             return int.from_bytes(byts, 'big')

synapse/exc.py

@@ -222,6 +222,7 @@ class SlabAlreadyOpen(SynErr): pass
 class SpawnExit(SynErr): pass
 class FeatureNotSupported(SynErr): pass
 
+class HitLimit(SynErr): pass
 class ReadOnlyLayer(SynErr): pass
 class ReadOnlyProp(SynErr): pass
 class RecursionLimitHit(SynErr): pass

synapse/tests/test_axon.py

@@ -344,3 +344,21 @@ async def test_axon_perms(self):
                 await user.addRule((True, ('axon', 'has',)))
                 await user.addRule((True, ('axon', 'upload',)))
                 await self.runAxonTestBase(prox)
+
+    async def test_axon_limits(self):
+
+        with self.setTstEnvars(AXON_MAX_COUNT=10):
+            async with self.getTestAxon() as axon:
+                for i in range(10):
+                    await axon.put(s_common.buid())
+
+                with self.raises(s_exc.HitLimit):
+                    await axon.put(s_common.buid())
+
+        with self.setTstEnvars(AXON_MAX_BYTES=320):
+            async with self.getTestAxon() as axon:
+                for i in range(10):
+                    await axon.put(s_common.buid())
+
+                with self.raises(s_exc.HitLimit):
+                    await axon.put(s_common.buid())