sshban: periodically remove expired items from ExpiringCounter

1 files changed, 18 insertions, 0 deletions

diff --git a/sshban.py b/sshban.py
index 926ec95..102f4a3 100755
--- a/sshban.py
+++ b/sshban.py
@@ -55,12 +55,16 @@ class ExpiringCounter:
     default_timeout = None
 
     _data = None
+    _gc_counter = None
 
     def __init__(self, default_timeout):
         self._data = {}
         self.default_timeout = default_timeout
+        self._gc_counter = 0
 
     def __str__(self):
+        self._gc()
+
         now = self._now()
         ret = ''
         for key, (ts, count) in self._data.items():
@@ -86,6 +90,16 @@ class ExpiringCounter:
     def _now(self):
         return time.clock_gettime(time.CLOCK_BOOTTIME)
 
+    def _gc(self):
+        to_remove = []
+        now = self._now()
+        for key, (ts, count) in self._data.items():
+            if now - ts > self.default_timeout:
+                to_remove.append(key)
+
+        for key in to_remove:
+            del self._data[key]
+
     def inc(self, key, count = 1):
         now = self._now()
 
@@ -97,6 +111,10 @@ class ExpiringCounter:
         elif key in self:
             del self[key]
 
+        self._gc_counter += 1
+        if (self._gc_counter & ((1 << 10) - 1)) == 0:
+            self._gc()
+
         return newval
 
     def dec(self, key, count = 1):