- add a data structure that can store aggregated data within

slots based on interval

Author: zzzeek

sqlalchemy_collectd/server/aggregator.py

@@ -12,3 +12,73 @@ def put(self, message):
     def outgoing(self):
         while self.queue:
             yield self.queue.pop()
+
+
+class TimeBucket(object):
+    """Store the last N seconds of time-stamped data within
+    interval-keyed buckets.
+
+    The idea is we can store and retrieve records that were
+    within the last N seconds only, or in the previous
+    2N-N seconds, or 3N-2N seconds, including that we can efficiently
+    clean up old ranges in O(1) time.
+
+    E.g. assume four buckets and interval of 100::
+
+        bucket[0] -> timestamp 50000-50100
+        bucket[1] -> timestamp 50101-50200
+        bucket[2] -> timestamp 50201-50300
+        bucket[3] -> timestamp 50301-50400
+
+    100 seconds later::
+
+        bucket[0] -> timestamp 50401-50500
+        bucket[1] -> timestamp 50101-50200
+        bucket[2] -> timestamp 50201-50300
+        bucket[3] -> timestamp 50301-50400
+
+    100 seconds later::
+
+        bucket[0] -> timestamp 50401-50500
+        bucket[1] -> timestamp 50501-50600
+        bucket[2] -> timestamp 50201-50300
+        bucket[3] -> timestamp 50301-50400
+
+    etc.
+
+    The object assumes if a new timestamp is coming in that is newer
+    than the current bucket, we go to the next bucket.   If the next bucket
+    has data from the old range it had B buckets ago, we empty it out first.
+
+    """
+    __slots__ = 'num_buckets', 'buckets', 'interval'
+
+    def __init__(self, num_buckets, interval):
+        self.num_buckets = num_buckets
+        self.buckets = [
+            {"slot": None, "data": {}} for i in range(num_buckets)
+        ]
+        self.interval = interval
+
+    def _get_bucket(self, timestamp):
+        slot = (timestamp // self.interval)
+        bucket_num = slot % self.num_buckets
+        bucket = self.buckets[bucket_num]
+        bucket_slot = bucket["slot"]
+        if bucket_slot is None:
+            bucket["slot"] = slot
+        elif bucket_slot < slot:
+            bucket["data"].clear()
+            bucket["slot"] = slot
+        elif bucket_slot > slot:
+            raise KeyError()
+        return bucket
+
+    def put(self, timestamp, key, data):
+        self._get_bucket(timestamp)['data'][key] = data
+
+    def get(self, current_time, key):
+        return self._get_bucket(current_time)['data'].get(key)
+
+    def get_data(self, current_time):
+        return self._get_bucket(current_time)['data']

sqlalchemy_collectd/server/tests/test_aggregator.py

@@ -0,0 +1,75 @@
+import mock
+import unittest
+import random
+
+from .. import aggregator
+
+
+class TimeBucketTest(unittest.TestCase):
+    def _generate(self, interval):
+        current = random.randint(729404, 930285)
+
+        for i in range(200):
+            yield current
+            current += random.randint(0, interval // 2)
+
+    def test_put(self):
+        agg = aggregator.TimeBucket(4, 10)
+        agg.put(50530, "key", "value_50530")
+
+        self.assertEqual(agg.get(50532, "key"), "value_50530")
+
+        agg.put(50534, "key", "value_50534")
+        self.assertEqual(agg.get(50535, "key"), "value_50534")
+
+        self.assertEqual(agg.get(50538, "key"), "value_50534")
+
+        # bucket has expired with this timestamp
+        self.assertEqual(agg.get(50542, "key"), None)
+
+        # but we can still see it w/ recent timestamp
+        self.assertEqual(agg.get(50539, "key"), "value_50534")
+
+        # current bucket
+        agg.put(50545, "key", "value_50545")
+
+        # bump
+        agg.put(50556, "key", "value_50556")
+
+        # still see old value
+        self.assertEqual(agg.get(50539, "key"), "value_50534")
+
+        # bump
+        agg.put(50562, "key", "value_50562")
+
+        # still see old value
+        self.assertEqual(agg.get(50539, "key"), "value_50534")
+        self.assertEqual(agg.get_data(50539)["key"], "value_50534")
+
+        # bump
+        agg.put(50574, "key", "value_50574")
+
+        # now it's gone
+        self.assertRaises(
+            KeyError,
+            agg.get_data, 50539
+        )
+
+    def test_series(self):
+        agg = aggregator.TimeBucket(4, 10)
+        previous_time = None
+        for round_, time in enumerate(self._generate(10)):
+            if round_ % 4 == 1:
+                value = agg.get(time, "key")
+                if value is None:
+                    assert previous_time // 10 != time // 10
+                else:
+                    assert previous_time // 10 == time // 10
+                    self.assertEqual(
+                        max([bucket["data"].get("key", '')
+                            for bucket in agg.buckets]),
+                        value
+                    )
+            else:
+                agg.put(time, "key", "value_%s" % time)
+            previous_time = time

sqlalchemy_collectd/tests/test_protocol.py

@@ -1,7 +1,7 @@
 import mock
 import unittest
 
-from sqlalchemy_collectd import protocol
+from .. import protocol
 
 
 class CollectDProtocolTest(unittest.TestCase):