Implement max size logic for local storage for Azure exporters (#876)

Author: lzchen

contrib/opencensus-ext-azure/opencensus/ext/azure/common/__init__.py

@@ -97,7 +97,7 @@ def __init__(self, *args, **kwargs):
         minimum_retry_interval=60,  # minimum retry interval in seconds
         proxy=None,
         storage_maintenance_period=60,
-        storage_max_size=100*1024*1024,
+        storage_max_size=50*1024*1024,  # 50MiB
         storage_path=os.path.join(
             os.path.expanduser('~'),
             '.opencensus',

contrib/opencensus-ext-azure/opencensus/ext/azure/common/storage.py

@@ -1,10 +1,13 @@
 import datetime
 import json
+import logging
 import os
 import random
 
 from opencensus.common.schedule import PeriodicTask
 
+logger = logging.getLogger(__name__)
+
 
 def _fmt(timestamp):
     return timestamp.strftime('%Y-%m-%dT%H%M%S.%f')
@@ -77,7 +80,7 @@ class LocalFileStorage(object):
     def __init__(
             self,
             path,
-            max_size=100*1024*1024,  # 100MB
+            max_size=50*1024*1024,  # 50MiB
             maintenance_period=60,  # 1 minute
             retention_period=7*24*60*60,  # 7 days
             write_timeout=60,  # 1 minute
@@ -162,6 +165,8 @@ def get(self):
         return None
 
     def put(self, data, lease_period=0, silent=False):
+        if not self._check_storage_size():
+            return None
         blob = LocalFileBlob(os.path.join(
             self.path,
             '{}-{}.blob'.format(
@@ -170,3 +175,28 @@ def put(self, data, lease_period=0, silent=False):
             ),
         ))
         return blob.put(data, lease_period=lease_period, silent=silent)
+
+    def _check_storage_size(self):
+        size = 0
+        for dirpath, dirnames, filenames in os.walk(self.path):
+            for f in filenames:
+                fp = os.path.join(dirpath, f)
+                # skip if it is symbolic link
+                if not os.path.islink(fp):
+                    try:
+                        size += os.path.getsize(fp)
+                    except OSError:
+                        logger.error(
+                            "Path %s does not exist or is inaccessible.", fp
+                        )
+                        continue
+                    if size >= self.max_size:
+                        logger.warning(
+                            "Persistent storage max capacity has been "
+                            "reached. Currently at %fKB. Telemetry will be "
+                            "lost. Please consider increasing the value of "
+                            "'storage_max_size' in exporter config.",
+                            format(size/1024)
+                        )
+                        return False
+        return True

contrib/opencensus-ext-azure/tests/test_storage.py

@@ -25,7 +25,7 @@
     _seconds,
 )
 
-TEST_FOLDER = os.path.abspath('.test')
+TEST_FOLDER = os.path.abspath('.test.storage')
 
 
 def setUpModule():
@@ -116,6 +116,48 @@ def test_put(self):
                 self.assertIsNone(stor.put(input, silent=True))
                 self.assertRaises(Exception, lambda: stor.put(input))
 
+    def test_put_max_size(self):
+        input = (1, 2, 3)
+        with LocalFileStorage(os.path.join(TEST_FOLDER, 'asd')) as stor:
+            size_mock = mock.Mock()
+            size_mock.return_value = False
+            stor._check_storage_size = size_mock
+            stor.put(input)
+            self.assertEqual(stor.get(), None)
+
+    def test_check_storage_size_full(self):
+        input = (1, 2, 3)
+        with LocalFileStorage(os.path.join(TEST_FOLDER, 'asd2'), 1) as stor:
+            stor.put(input)
+            self.assertFalse(stor._check_storage_size())
+
+    def test_check_storage_size_not_full(self):
+        input = (1, 2, 3)
+        with LocalFileStorage(os.path.join(TEST_FOLDER, 'asd3'), 1000) as stor:
+            stor.put(input)
+            self.assertTrue(stor._check_storage_size())
+
+    def test_check_storage_size_no_files(self):
+        with LocalFileStorage(os.path.join(TEST_FOLDER, 'asd3'), 1000) as stor:
+            self.assertTrue(stor._check_storage_size())
+
+    def test_check_storage_size_links(self):
+        input = (1, 2, 3)
+        with LocalFileStorage(os.path.join(TEST_FOLDER, 'asd4'), 1000) as stor:
+            stor.put(input)
+            with mock.patch('os.path.islink') as os_mock:
+                os_mock.return_value = True
+            self.assertTrue(stor._check_storage_size())
+
+    def test_check_storage_size_error(self):
+        input = (1, 2, 3)
+        with LocalFileStorage(os.path.join(TEST_FOLDER, 'asd5'), 1) as stor:
+            with mock.patch('os.path.getsize', side_effect=throw(OSError)):
+                stor.put(input)
+                with mock.patch('os.path.islink') as os_mock:
+                    os_mock.return_value = True
+                self.assertTrue(stor._check_storage_size())
+
     def test_maintanence_routine(self):
         with mock.patch('os.makedirs') as m:
             m.return_value = None