From d4cc2c20b88e32d73d18f40396ee07d4f762d56c Mon Sep 17 00:00:00 2001
From: Leighton <lechen@microsoft.com>
Date: Thu, 2 Apr 2020 09:07:00 -0700
Subject: [PATCH] max storage

---
 azure_monitor/src/azure_monitor/options.py |  2 +-
 azure_monitor/src/azure_monitor/storage.py | 31 +++++++++++++++-
 azure_monitor/tests/test_storage.py        | 42 ++++++++++++++++++++++
 3 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/azure_monitor/src/azure_monitor/options.py b/azure_monitor/src/azure_monitor/options.py
index a8e6c91..739ad37 100644
--- a/azure_monitor/src/azure_monitor/options.py
+++ b/azure_monitor/src/azure_monitor/options.py
@@ -50,7 +50,7 @@ def __init__(
         connection_string: str = None,
         instrumentation_key: str = None,
         storage_maintenance_period: int = 60,
-        storage_max_size: int = 100 * 1024 * 1024,
+        storage_max_size: int = 50 * 1024 * 1024,
         storage_path: str = None,
         storage_retention_period: int = 7 * 24 * 60 * 60,
         timeout: int = 10.0,  # networking timeout in seconds
diff --git a/azure_monitor/src/azure_monitor/storage.py b/azure_monitor/src/azure_monitor/storage.py
index f8a265a..d01c1d8 100644
--- a/azure_monitor/src/azure_monitor/storage.py
+++ b/azure_monitor/src/azure_monitor/storage.py
@@ -3,11 +3,14 @@
 
 import datetime
 import json
+import logging
 import os
 import random
 
 from azure_monitor.utils import PeriodicTask
 
+logger = logging.getLogger(__name__)
+
 
 def _fmt(timestamp):
     return timestamp.strftime("%Y-%m-%dT%H%M%S.%f")
@@ -81,7 +84,7 @@ class LocalFileStorage:
     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
@@ -167,6 +170,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,
@@ -179,3 +184,27 @@ 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
diff --git a/azure_monitor/tests/test_storage.py b/azure_monitor/tests/test_storage.py
index a24b49a..c4258b0 100644
--- a/azure_monitor/tests/test_storage.py
+++ b/azure_monitor/tests/test_storage.py
@@ -108,6 +108,48 @@ def test_put(self):
                 self.assertIsNone(stor.put(test_input, silent=True))
                 self.assertRaises(Exception, lambda: stor.put(test_input))
 
+    def test_put_max_size(self):
+        test_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(test_input)
+            self.assertEqual(stor.get(), None)
+
+    def test_check_storage_size_full(self):
+        test_input = (1, 2, 3)
+        with LocalFileStorage(os.path.join(TEST_FOLDER, 'asd2'), 1) as stor:
+            stor.put(test_input)
+            self.assertFalse(stor._check_storage_size())
+
+    def test_check_storage_size_not_full(self):
+        test_input = (1, 2, 3)
+        with LocalFileStorage(os.path.join(TEST_FOLDER, 'asd3'), 1000) as stor:
+            stor.put(test_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):
+        test_input = (1, 2, 3)
+        with LocalFileStorage(os.path.join(TEST_FOLDER, 'asd4'), 1000) as stor:
+            stor.put(test_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):
+        test_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(test_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