prepare storages to handle traffic type validation

Author: mredolatti

splitio/storage/__init__.py

@@ -83,6 +83,19 @@ def get_all_splits(self):
         """
         pass
 
+    @abc.abstractmethod
+    def is_valid_traffic_type(self, traffic_type_name):
+        """
+        Return whether the traffic type exists in at least one split in cache.
+
+        :param traffic_type_name: Traffic type to validate.
+        :type traffic_type_name: str
+
+        :return: True if the traffic type is valid. False otherwise.
+        :rtype: bool
+        """
+        pass
+
     def get_segment_names(self):
         """
         Return a set of all segments referenced by splits in storage.

splitio/storage/inmemmory.py

@@ -3,6 +3,8 @@
 
 import logging
 import threading
+from collections import Counter
+
 from six.moves import queue
 from splitio.models.segments import Segment
 from splitio.storage import SplitStorage, SegmentStorage, ImpressionStorage, EventStorage, \
@@ -18,6 +20,7 @@ def __init__(self):
         self._lock = threading.RLock()
         self._splits = {}
         self._change_number = -1
+        self._traffic_types = Counter()
 
     def get(self, split_name):
         """
@@ -40,6 +43,7 @@ def put(self, split):
         """
         with self._lock:
             self._splits[split.name] = split
+            self._increase_traffic_type_count(split.traffic_type_name)
 
     def remove(self, split_name):
         """
@@ -52,13 +56,15 @@ def remove(self, split_name):
         :rtype: bool
         """
         with self._lock:
-            try:
-                self._splits.pop(split_name)
-                return True
-            except KeyError:
+            split = self._splits.get(split_name)
+            if not split:
                 self._logger.warning("Tried to delete nonexistant split %s. Skipping", split_name)
                 return False
 
+            self._splits.pop(split_name)
+            self._decrease_traffic_type_count(split.traffic_type_name)
+            return True
+
     def get_change_number(self):
         """
         Retrieve latest split change number.
@@ -98,6 +104,38 @@ def get_all_splits(self):
         with self._lock:
             return list(self._splits.values())
 
+    def is_valid_traffic_type(self, traffic_type_name):
+        """
+        Return whether the traffic type exists in at least one split in cache.
+
+        :param traffic_type_name: Traffic type to validate.
+        :type traffic_type_name: str
+
+        :return: True if the traffic type is valid. False otherwise.
+        :rtype: bool
+        """
+        with self._lock:
+            return traffic_type_name in self._traffic_types
+
+    def _increase_traffic_type_count(self, traffic_type_name):
+        """
+        Increase by one the count for a specific traffic type name.
+
+        :param traffic_type_name: Traffic type to increase the count.
+        :type traffic_type_name: str
+        """
+        self._traffic_types.update([traffic_type_name])
+
+    def _decrease_traffic_type_count(self, traffic_type_name):
+        """
+        Decrease by one the count for a specific traffic type name.
+
+        :param traffic_type_name: Traffic type to decrease the count.
+        :type traffic_type_name: str
+        """
+        self._traffic_types.subtract([traffic_type_name])
+        self._traffic_types += Counter()
+
 
 class InMemorySegmentStorage(SegmentStorage):
     """In-memory implementation of a segment storage."""

splitio/storage/redis.py

@@ -16,6 +16,7 @@ class RedisSplitStorage(SplitStorage):
 
     _SPLIT_KEY = 'SPLITIO.split.{split_name}'
     _SPLIT_TILL_KEY = 'SPLITIO.splits.till'
+    _TRAFFIC_TYPE_KEY = 'SPLITIO.trafficType.{traffic_type_name}'
 
     def __init__(self, redis_client):
         """
@@ -39,6 +40,18 @@ def _get_key(self, split_name):
         """
         return self._SPLIT_KEY.format(split_name=split_name)
 
+    def _get_traffic_type_key(self, traffic_type_name):
+        """
+        Use the provided split_name to build the appropriate redis key.
+
+        :param split_name: Name of the split to interact with in redis.
+        :type split_name: str
+
+        :return: Redis key.
+        :rtype: str.
+        """
+        return self._TRAFFIC_TYPE_KEY.format(traffic_type_name=traffic_type_name)
+
     def get(self, split_name):
         """
         Retrieve a split.
@@ -57,6 +70,25 @@ def get(self, split_name):
             self._logger.debug('Error: ', exc_info=True)
             return None
 
+    def is_valid_traffic_type(self, traffic_type_name):
+        """
+        Return whether the traffic type exists in at least one split in cache.
+
+        :param traffic_type_name: Traffic type to validate.
+        :type traffic_type_name: str
+
+        :return: True if the traffic type is valid. False otherwise.
+        :rtype: bool
+        """
+        try:
+            raw = self._redis.get(self._get_traffic_type_key(traffic_type_name))
+            count = json.loads(raw) if raw else 0
+            return count > 0
+        except RedisAdapterException:
+            self._logger.error('Error fetching split from storage')
+            self._logger.debug('Error: ', exc_info=True)
+            return False
+
     def put(self, split):
         """
         Store a split.

splitio/storage/uwsgi.py

@@ -19,8 +19,10 @@ class UWSGISplitStorage(SplitStorage):
 
     _KEY_TEMPLATE = 'split.{suffix}'
     _KEY_TILL = 'splits.till'
-    _KEY_FEATURE_LIST_LOCK = 'splits.list.lock'
     _KEY_FEATURE_LIST = 'splits.list'
+    _KEY_FEATURE_LIST_LOCK = 'splits.list.lock'
+    _KEY_TRAFFIC_TYPES = 'splits.traffic_types'
+    _KEY_TRAFFIC_TYPES_LOCK = 'splits.traffic_types.lock'
     _OVERWRITE_LOCK_SECONDS = 5
 
     def __init__(self, uwsgi_entrypoint):
@@ -64,21 +66,8 @@ def put(self, split):
             0,
             _SPLITIO_SPLITS_CACHE_NAMESPACE
         )
-
-        with UWSGILock(self._uwsgi, self._KEY_FEATURE_LIST_LOCK):
-            try:
-                current = set(json.loads(
-                    self._uwsgi.cache_get(self._KEY_FEATURE_LIST, _SPLITIO_MISC_NAMESPACE)
-                ))
-            except TypeError:
-                current = set()
-            current.add(split.name)
-            self._uwsgi.cache_update(
-                self._KEY_FEATURE_LIST,
-                json.dumps(list(current)),
-                0,
-                _SPLITIO_MISC_NAMESPACE
-            )
+        self._add_split_to_list(split.name)
+        self._increase_traffic_type_count(split.traffic_type_name)
 
     def remove(self, split_name):
         """
@@ -90,31 +79,23 @@ def remove(self, split_name):
         :return: True if the split was found and removed. False otherwise.
         :rtype: bool
         """
-        with UWSGILock(self._uwsgi, self._KEY_FEATURE_LIST_LOCK):
-            try:
-                current = set(json.loads(
-                    self._uwsgi.cache_get(self._KEY_FEATURE_LIST, _SPLITIO_MISC_NAMESPACE)
-                ))
-                current.remove(split_name)
-                self._uwsgi.cache_update(
-                    self._KEY_FEATURE_LIST,
-                    json.dumps(list(current)),
-                    0,
-                    _SPLITIO_MISC_NAMESPACE
-                )
-            except TypeError:
-                # Split list not found, no need to delete anything
-                pass
-            except KeyError:
-                # Split not found in list. nothing to do.
-                pass
+        # We need to fetch the split to get the traffic type name prior to deleting.
+        fetched = self.get(split_name)
+        if fetched is None:
+            self._logger.warning(
+                "Tried to remove feature \"%s\" not present in cache. Ignoring.", split_name
+            )
+            return
 
         result = self._uwsgi.cache_del(
             self._KEY_TEMPLATE.format(suffix=split_name),
             _SPLITIO_SPLITS_CACHE_NAMESPACE
         )
         if not result is False:
-            self._logger.warning("Trying to retrieve nonexistant split %s. Ignoring.", split_name)
+            self._logger.warning("Trying to delete nonexistant split %s. Ignoring.", split_name)
+
+        self._remove_split_from_list(split_name)
+        self._decrease_traffic_type_count(fetched.traffic_type_name)
         return result
 
     def get_change_number(self):
@@ -162,6 +143,116 @@ def get_all_splits(self):
         """
         return [self.get(split_name) for split_name in self.get_split_names()]
 
+    def is_valid_traffic_type(self, traffic_type_name):
+        """
+        Return whether the traffic type exists in at least one split in cache.
+
+        :param traffic_type_name: Traffic type to validate.
+        :type traffic_type_name: str
+
+        :return: True if the traffic type is valid. False otherwise.
+        :rtype: bool
+        """
+        try:
+            tts = json.loads(
+                self._uwsgi.cache_get(self._KEY_TRAFFIC_TYPES, _SPLITIO_MISC_NAMESPACE)
+            )
+            return traffic_type_name in tts
+        except TypeError:
+            return False
+
+    def _add_split_to_list(self, split_name):
+        """
+        Add a specific split to the list we keep track of.
+
+        :param split_name: Name of the split to add.
+        :type split_name: str
+        """
+        with UWSGILock(self._uwsgi, self._KEY_FEATURE_LIST_LOCK):
+            try:
+                current = set(json.loads(
+                    self._uwsgi.cache_get(self._KEY_FEATURE_LIST, _SPLITIO_MISC_NAMESPACE)
+                ))
+            except TypeError:
+                current = set()
+            current.add(split_name)
+            self._uwsgi.cache_update(
+                self._KEY_FEATURE_LIST,
+                json.dumps(list(current)),
+                0,
+                _SPLITIO_MISC_NAMESPACE
+            )
+
+    def _remove_split_from_list(self, split_name):
+        """
+        Remove a specific split from the list we keep track of.
+
+        :param split_name: Name of the split to remove.
+        :type split_name: str
+        """
+        with UWSGILock(self._uwsgi, self._KEY_FEATURE_LIST_LOCK):
+            try:
+                current = set(json.loads(
+                    self._uwsgi.cache_get(self._KEY_FEATURE_LIST, _SPLITIO_MISC_NAMESPACE)
+                ))
+                current.remove(split_name)
+                self._uwsgi.cache_update(
+                    self._KEY_FEATURE_LIST,
+                    json.dumps(list(current)),
+                    0,
+                    _SPLITIO_MISC_NAMESPACE
+                )
+            except TypeError:
+                # Split list not found, no need to delete anything
+                pass
+            except KeyError:
+                # Split not found in list. nothing to do.
+                pass
+
+    def _increase_traffic_type_count(self, traffic_type_name):
+        """
+        Increase by 1 the count for a specific traffic type.
+
+        :param traffic_type_name: Traffic type name to increase count.
+        :type traffic_type_name: str
+        """
+        with UWSGILock(self._uwsgi, self._KEY_TRAFFIC_TYPES_LOCK):
+            try:
+                tts = json.loads(
+                    self._uwsgi.cache_get(self._KEY_TRAFFIC_TYPES, _SPLITIO_MISC_NAMESPACE)
+                )
+                tts[traffic_type_name] = tts.get(traffic_type_name, 0) + 1
+
+            except TypeError:
+                tts = {traffic_type_name: 1}
+
+            self._uwsgi.cache_update(
+                self._KEY_TRAFFIC_TYPES, json.dumps(tts), 0, _SPLITIO_MISC_NAMESPACE
+            )
+
+    def _decrease_traffic_type_count(self, traffic_type_name):
+        """
+        Decreaase by 1 the count for a specific traffic type.
+
+        :param traffic_type_name: Traffic type name to decrease count.
+        :type traffic_type_name: str
+        """
+        with UWSGILock(self._uwsgi, self._KEY_TRAFFIC_TYPES_LOCK):
+            try:
+                tts = json.loads(
+                    self._uwsgi.cache_get(self._KEY_TRAFFIC_TYPES, _SPLITIO_MISC_NAMESPACE)
+                )
+                tts[traffic_type_name] = tts.get(traffic_type_name, 0) - 1
+                if tts[traffic_type_name] <= 0:
+                    del tts[traffic_type_name]
+            except TypeError:
+                # Traffic type list not present. nothing to do here.
+                return
+
+            self._uwsgi.cache_update(
+                self._KEY_TRAFFIC_TYPES, json.dumps(tts), 0, _SPLITIO_MISC_NAMESPACE
+            )
+
 
 class UWSGISegmentStorage(SegmentStorage):
     """UWSGI-Cache based implementation of a split storage."""