Fixed adding telemetry latency to bucket

Author: chillaq

splitio/storage/pluggable.py

@@ -6,7 +6,7 @@
 
 from splitio.models import splits, segments
 from splitio.models.impressions import Impression
-from splitio.models.telemetry import MethodExceptions, MethodLatencies, TelemetryConfig, MAX_TAGS
+from splitio.models.telemetry import MethodExceptions, MethodLatencies, TelemetryConfig, MAX_TAGS, get_latency_bucket_index
 from splitio.storage import SplitStorage, SegmentStorage, ImpressionStorage, EventStorage, TelemetryStorage
 
 _LOGGER = logging.getLogger(__name__)
@@ -780,18 +780,10 @@ def record_latency(self, method, latency):
         :param latency: latency
         :type latency: int64
         """
-        self._method_latencies.add_latency(method, latency)
-        latencies = self._method_latencies.pop_all()['methodLatencies']
-        values = latencies[method.value]
-        total_keys = 0
-        bucket_number = 0
-        for bucket in values:
-            if bucket > 0:
-                latency_key = self._telemetry_latencies_key + '::' + self._sdk_metadata + '/' + method.value + '/' + str(bucket_number)
-                result = self._pluggable_adapter.increment(latency_key, bucket)
-                self.expire_keys(latency_key, self._TELEMETRY_KEY_DEFAULT_TTL, 1, result)
-                total_keys += 1
-            bucket_number = bucket_number + 0
+        bucket = get_latency_bucket_index(latency)
+        latency_key = self._telemetry_latencies_key + '::' + self._sdk_metadata + '/' + method.value + '/' + str(bucket)
+        result = self._pluggable_adapter.increment(latency_key, 1)
+        self.expire_keys(latency_key, self._TELEMETRY_KEY_DEFAULT_TTL, 1, result)
 
     def record_exception(self, method):
         """

splitio/storage/redis.py

@@ -5,7 +5,7 @@
 
 from splitio.models.impressions import Impression
 from splitio.models import splits, segments
-from splitio.models.telemetry import MethodExceptions, MethodLatencies, TelemetryConfig
+from splitio.models.telemetry import MethodExceptions, MethodLatencies, TelemetryConfig, get_latency_bucket_index
 from splitio.storage import SplitStorage, SegmentStorage, ImpressionStorage, EventStorage, \
     ImpressionPipelinedStorage, TelemetryStorage
 from splitio.storage.adapters.redis import RedisAdapterException
@@ -660,17 +660,9 @@ def add_latency_to_pipe(self, method, latency, pipe):
         :param pipe: Redis pipe.
         :type pipe: redis.pipe
         """
-        self._method_latencies.add_latency(method, latency)
-        latencies = self._method_latencies.pop_all()['methodLatencies']
-        values = latencies[method.value]
-        total_keys = 0
-        bucket_number = 0
-        for bucket in values:
-            if bucket > 0:
-                pipe.hincrby(self._TELEMETRY_LATENCIES_KEY, self._sdk_metadata.sdk_version + '/' + self._sdk_metadata.instance_name + '/' + self._sdk_metadata.instance_ip + '/' +
-                        method.value + '/' + str(bucket_number), bucket)
-                total_keys += 1
-            bucket_number = bucket_number + 0
+        bucket = get_latency_bucket_index(latency)
+        pipe.hincrby(self._TELEMETRY_LATENCIES_KEY, self._sdk_metadata.sdk_version + '/' + self._sdk_metadata.instance_name + '/' + self._sdk_metadata.instance_ip + '/' +
+            method.value + '/' + str(bucket), 1)
 
     def record_latency(self, method, latency):
         """

tests/storage/test_pluggable.py

@@ -599,11 +599,30 @@ def expire_keys_mock(*args, **kwargs):
             assert(args[1] == self.pluggable_telemetry_storage._TELEMETRY_KEY_DEFAULT_TTL)
             assert(args[2] == 1)
             assert(args[3] == 1)
-
         self.pluggable_telemetry_storage.expire_keys = expire_keys_mock
+        # should increment bucket 0
         self.pluggable_telemetry_storage.record_latency(MethodExceptionsAndLatencies.TREATMENT, 10)
         assert(self.mock_adapter._keys[self.pluggable_telemetry_storage._telemetry_latencies_key + '::python-1.1.1/hostname/ip/treatment/0'] == 1)
 
+        def expire_keys_mock2(*args, **kwargs):
+            assert(args[0] == self.pluggable_telemetry_storage._telemetry_latencies_key + '::python-1.1.1/hostname/ip/treatment/3')
+            assert(args[1] == self.pluggable_telemetry_storage._TELEMETRY_KEY_DEFAULT_TTL)
+            assert(args[2] == 1)
+            assert(args[3] == 1)
+        self.pluggable_telemetry_storage.expire_keys = expire_keys_mock2
+        # should increment bucket 3
+        self.pluggable_telemetry_storage.record_latency(MethodExceptionsAndLatencies.TREATMENT, 2260)
+
+        def expire_keys_mock3(*args, **kwargs):
+            assert(args[0] == self.pluggable_telemetry_storage._telemetry_latencies_key + '::python-1.1.1/hostname/ip/treatment/3')
+            assert(args[1] == self.pluggable_telemetry_storage._TELEMETRY_KEY_DEFAULT_TTL)
+            assert(args[2] == 1)
+            assert(args[3] == 2)
+        self.pluggable_telemetry_storage.expire_keys = expire_keys_mock3
+        # should increment bucket 3
+        self.pluggable_telemetry_storage.record_latency(MethodExceptionsAndLatencies.TREATMENT, 3280)
+        assert(self.mock_adapter._keys[self.pluggable_telemetry_storage._telemetry_latencies_key + '::python-1.1.1/hostname/ip/treatment/3'] == 2)
+
     def test_record_exception(self):
         def expire_keys_mock(*args, **kwargs):
             assert(args[0] == self.pluggable_telemetry_storage._telemetry_exceptions_key + '::python-1.1.1/hostname/ip/treatment')

tests/storage/test_redis.py

@@ -430,18 +430,27 @@ def test_record_active_and_redundant_factories(self, mocker):
         assert (redis_telemetry._tel_config._redundant_factory_count == redundant_factory_count)
 
     def test_add_latency_to_pipe(self, mocker):
-        def _mocked_hincrby(*args, **kwargs):
-            assert(args[1] == RedisTelemetryStorage._TELEMETRY_LATENCIES_KEY)
-            assert(args[2][-11:] == 'treatment/0')
-            assert(args[3] == 1)
-
         adapter = build({})
         metadata = SdkMetadata('python-1.1.1', 'hostname', 'ip')
         redis_telemetry = RedisTelemetryStorage(adapter, metadata)
         pipe = adapter._decorated.pipeline()
+
+        def _mocked_hincrby(*args, **kwargs):
+            assert(args[1] == RedisTelemetryStorage._TELEMETRY_LATENCIES_KEY)
+            assert(args[2][-11:] == 'treatment/0')
+            assert(args[3] == 1)
+        # should increment bucket 0
         with mock.patch('redis.client.Pipeline.hincrby', _mocked_hincrby):
             redis_telemetry.add_latency_to_pipe(MethodExceptionsAndLatencies.TREATMENT, 20, pipe)
 
+        def _mocked_hincrby2(*args, **kwargs):
+            assert(args[1] == RedisTelemetryStorage._TELEMETRY_LATENCIES_KEY)
+            assert(args[2][-11:] == 'treatment/3')
+            assert(args[3] == 1)
+        # should increment bucket 3
+        with mock.patch('redis.client.Pipeline.hincrby', _mocked_hincrby2):
+            redis_telemetry.add_latency_to_pipe(MethodExceptionsAndLatencies.TREATMENT, 2260, pipe)
+
     def test_record_exception(self, mocker):
         def _mocked_hincrby(*args, **kwargs):
             assert(args[1] == RedisTelemetryStorage._TELEMETRY_EXCEPTIONS_KEY)