Merge pull request #266 from splitio/UniqueKeysTracker

Unique keys tracker

Author: mmelograno

splitio/engine/unique_keys_tracker.py

@@ -0,0 +1,89 @@
+import abc
+import threading
+import logging
+from splitio.engine.filters.bloom_filter import BloomFilter
+
+_LOGGER = logging.getLogger(__name__)
+
+class BaseUniqueKeysTracker(object, metaclass=abc.ABCMeta):
+    """Unique Keys Tracker interface."""
+
+    @abc.abstractmethod
+    def track(self, key, feature_name):
+        """
+        Return a boolean flag
+
+        """
+        pass
+
+    @abc.abstractmethod
+    def start(self):
+        """
+        No return value
+
+        """
+        pass
+
+    @abc.abstractmethod
+    def stop(self):
+        """
+        No return value
+
+        """
+        pass
+
+class UniqueKeysTracker(BaseUniqueKeysTracker):
+    """Unique Keys Tracker class."""
+
+    def __init__(self, cache_size=30000, max_bulk_size=5000, task_refresh_rate = 24):
+        self._cache_size = cache_size
+        self._max_bulk_size = max_bulk_size
+        self._task_refresh_rate = task_refresh_rate
+        self._filter = BloomFilter(cache_size)
+        self._lock = threading.RLock()
+        self._cache = {}
+        # TODO: initialize impressions sender adapter and task referesh rate in next PR
+
+    def track(self, key, feature_name):
+        """
+        Return a boolean flag
+
+        """
+        if self._filter.contains(feature_name+key):
+            return False
+
+        with self._lock:
+            self._add_or_update(feature_name, key)
+            self._filter.add(feature_name+key)
+
+        if len(self._cache[feature_name]) == self._cache_size:
+            _LOGGER.warn("MTK Cache size for Split [%s] has reach maximum unique keys [%d], flushing data now.", feature_name, self._cache_size)
+#            TODO: Flush the data and reset split cache in next PR
+        if self._get_dict_size() >= self._max_bulk_size:
+            _LOGGER.info("Bulk MTK cache size has reach maximum, flushing data now.")
+#            TODO: Flush the data and reset split cache in next PR
+
+        return True
+
+    def _get_dict_size(self):
+        total_size = 0
+        for key in self._cache:
+            total_size = total_size + len(self._cache[key])
+        return total_size
+
+    def _add_or_update(self, feature_name, key):
+        if feature_name not in self._cache:
+            self._cache[feature_name] = set()
+        self._cache[feature_name].add(key)
+
+    def start(self):
+        """
+        TODO: Add start posting impressions job in next PR
+
+        """
+
+    def stop(self):
+        """
+        TODO: Add stop posting impressions job in next PR
+
+        """

tests/engine/test_unique_keys_tracker.py

@@ -0,0 +1,51 @@
+"""BloomFilter unit tests."""
+
+import threading
+from splitio.engine.unique_keys_tracker import UniqueKeysTracker
+from splitio.engine.filters.bloom_filter import BloomFilter
+
+class UniqueKeysTrackerTests(object):
+    """StandardRecorderTests test cases."""
+
+    def test_adding_and_removing_keys(self, mocker):
+        tracker = UniqueKeysTracker()
+
+        assert(tracker._cache_size > 0)
+        assert(tracker._max_bulk_size > 0)
+        assert(tracker._task_refresh_rate > 0)
+        assert(isinstance(tracker._filter, BloomFilter))
+
+        key1 = 'key1'
+        key2 = 'key2'
+        key3 = 'key3'
+        split1= 'feature1'
+        split2= 'feature2'
+
+        assert(tracker.track(key1, split1))
+        assert(tracker.track(key3, split1))
+        assert(not tracker.track(key1, split1))
+        assert(tracker.track(key2, split2))
+
+        assert(tracker._filter.contains(split1+key1))
+        assert(not tracker._filter.contains(split1+key2))
+        assert(tracker._filter.contains(split2+key2))
+        assert(not tracker._filter.contains(split2+key1))
+        assert(key1 in tracker._cache[split1])
+        assert(key3 in tracker._cache[split1])
+        assert(key2 in tracker._cache[split2])
+        assert(not key3 in tracker._cache[split2])
+
+    def test_cache_size(self, mocker):
+        cache_size = 10
+        tracker = UniqueKeysTracker(cache_size)
+
+        split1= 'feature1'
+        for x in range(1, cache_size + 1):
+            tracker.track('key' + str(x), split1)
+        split2= 'feature2'
+        for x in range(1, int(cache_size / 2) + 1):
+            tracker.track('key' + str(x), split2)
+
+        assert(tracker._get_dict_size() == (cache_size + (cache_size / 2)))
+        assert(len(tracker._cache[split1]) == cache_size)
+        assert(len(tracker._cache[split2]) == cache_size / 2)