Browse Source

better memory management for counters

master
Yves G 3 years ago
parent
commit
4b42b69b49
1 changed files with 90 additions and 37 deletions
  1. +90
    -37
      pyruse/counter.py

+ 90
- 37
pyruse/counter.py View File

@ -1,53 +1,106 @@
# pyruse is intended as a replacement to both fail2ban and epylog
# Copyright © 2017 Y. Gablin
# Copyright © 2017–2018 Y. Gablin
# Full licensing information in the LICENSE file, or gnu.org/licences/gpl-3.0.txt if the file is missing.
import datetime
class Counter():
_counters = {}
class _GraceAndTicks():
def __init__(self):
self.grace = None
self.ticks = []
class _CounterData():
def __init__(self):
self._keyVals = {}
def clean(self, refDT):
for k in list(self._keyVals.keys()):
v = self._keyVals[k]
if v.grace and v.grace <= refDT:
v.grace = None
if v.grace is None:
# None values (∞) are at the end of the list
for i in range(0, len(v.ticks)):
if v.ticks[i] and v.ticks[i] <= refDT:
continue
v.ticks = v.ticks[i:]
break
else:
del self._keyVals[k]
def graceActive(self, counterKey, refDT):
self.clean(refDT)
return counterKey in self._keyVals and self._keyVals[counterKey].grace
def augment(self, counterKey, refDT, until):
self.clean(refDT)
if counterKey in self._keyVals:
v = self._keyVals[counterKey]
if v.grace:
return 0
else:
v = _GraceAndTicks()
self._keyVals[counterKey] = v
l = len(v.ticks)
# chances are that until is greater than the last item
for i in range(l, 0, -1):
if until is None or (v.ticks[i - 1] and v.ticks[i - 1] < until):
v.ticks.insert(i, until)
break
else:
v.ticks.insert(0, until)
return l + 1
def lower(self, counterKey, refDT):
self.clean(refDT)
v = self._keyVals.get(counterKey, None)
if v is None or v.grace:
return 0
l = len(v.ticks)
if l == 1:
del self._keyVals[counterKey]
return 0
def _getCounter(self, counterName):
if counterName not in Counter._counters:
Counter._counters[counterName] = {}
return Counter._counters[counterName]
v.ticks.pop()
return l - 1
def _cleanTicks(self, ticks, refDT):
for tick in list(ticks):
if tick and tick < refDT:
ticks.remove(tick)
def reset(self, counterKey, refDT, graceUntil):
self.clean(refDT)
if graceUntil:
v = _GraceAndTicks()
v.grace = graceUntil
self._keyVals[counterKey] = v
elif counterKey in self._keyVals:
del self._keyVals[counterKey]
class Counter():
_counters = {}
def __init__(self, counter):
self.counterName = counter
if counter not in Counter._counters:
Counter._counters[counter] = _CounterData()
self.counter = Counter._counters[counter]
def augment(self, counterKey, duration = None):
counter = self._getCounter(self.counterName)
if counterKey not in counter:
counter[counterKey] = ([], None) # [countUntil,…], graceTermDate
(ticks, grace) = counter[counterKey]
now = datetime.datetime.utcnow()
self._cleanTicks(ticks, now)
if grace and grace < now:
grace = None
if not grace:
ticks.append(now + duration if duration else None)
counter[counterKey] = (ticks, grace)
return len(ticks)
if self.counter.graceActive(counterKey, now):
return 0
else:
return self.counter.augment(counterKey, now, now + duration if duration else None)
def lower(self, counterKey):
counter = self._getCounter(self.counterName)
if counterKey not in counter:
return 0
(ticks, _void) = counter[counterKey]
now = datetime.datetime.utcnow()
self._cleanTicks(ticks, now)
if len(ticks) > 0:
ticks.pop()
return len(ticks)
if self.counter.graceActive(counterKey, now):
return 0
else:
return self.counter.lower(counterKey, now)
def reset(self, counterKey, graceDuration = None):
counter = self._getCounter(self.counterName)
if graceDuration:
now = datetime.datetime.utcnow()
counter[counterKey] = ([], now + graceDuration)
elif counterKey in counter:
del(counter[counterKey])
now = datetime.datetime.utcnow()
self.counter.reset(
counterKey, now,
now + graceDuration if graceDuration else None
)

Loading…
Cancel
Save