Merge pull request #566 from splitio/rbs-fix-segment-initial-fetch

chillaq · web-flow · commit 3fea6cb0920b · 2025-05-19T08:29:43.000-07:00
Rbs fix segment initial fetch
diff --git a/splitio/api/splits.py b/splitio/api/splits.py
@@ -37,11 +37,20 @@ def __init__(self, client, sdk_key, sdk_metadata, telemetry_runtime_producer):
         self._spec_version = SPEC_VERSION
         self._last_proxy_check_timestamp = 0
         self.clear_storage = False
+        self._old_spec_since = None
 
-    def _check_last_proxy_check_timestamp(self):
+    def _check_last_proxy_check_timestamp(self, since):
         if self._spec_version == _SPEC_1_1 and ((utctime_ms() - self._last_proxy_check_timestamp) >= _PROXY_CHECK_INTERVAL_MILLISECONDS_SS):
             _LOGGER.info("Switching to new Feature flag spec (%s) and fetching.", SPEC_VERSION);
             self._spec_version = SPEC_VERSION
+            self._old_spec_since = since
+    
+    def _check_old_spec_since(self, change_number):
+        if self._spec_version == _SPEC_1_1 and self._old_spec_since is not None:
+            since = self._old_spec_since
+            self._old_spec_since = None
+            return since
+        return change_number
     
 
 class SplitsAPI(SplitsAPIBase):  # pylint: disable=too-few-public-methods
@@ -77,7 +86,9 @@ def fetch_splits(self, change_number, rbs_change_number, fetch_options):
         :rtype: dict
         """
         try:
-            self._check_last_proxy_check_timestamp()
+            self._check_last_proxy_check_timestamp(change_number)
+            change_number = self._check_old_spec_since(change_number)
+
             query, extra_headers = build_fetch(change_number, fetch_options, self._metadata, rbs_change_number)
             response = self._client.get(
                 'sdk',
@@ -145,7 +156,9 @@ async def fetch_splits(self, change_number, rbs_change_number, fetch_options):
         :rtype: dict
         """
         try:
-            self._check_last_proxy_check_timestamp()            
+            self._check_last_proxy_check_timestamp(change_number)
+            change_number = self._check_old_spec_since(change_number)
+            
             query, extra_headers = build_fetch(change_number, fetch_options, self._metadata, rbs_change_number)
             response = await self._client.get(
                 'sdk',
diff --git a/splitio/client/factory.py b/splitio/client/factory.py
@@ -564,7 +564,7 @@ def _build_in_memory_factory(api_key, cfg, sdk_url=None, events_url=None,  # pyl
 
     synchronizers = SplitSynchronizers(
         SplitSynchronizer(apis['splits'], storages['splits'], storages['rule_based_segments']),
-        SegmentSynchronizer(apis['segments'], storages['splits'], storages['segments']),
+        SegmentSynchronizer(apis['segments'], storages['splits'], storages['segments'], storages['rule_based_segments']),
         ImpressionSynchronizer(apis['impressions'], storages['impressions'],
                                cfg['impressionsBulkSize']),
         EventSynchronizer(apis['events'], storages['events'], cfg['eventsBulkSize']),
@@ -693,7 +693,7 @@ async def _build_in_memory_factory_async(api_key, cfg, sdk_url=None, events_url=
 
     synchronizers = SplitSynchronizers(
         SplitSynchronizerAsync(apis['splits'], storages['splits'], storages['rule_based_segments']),
-        SegmentSynchronizerAsync(apis['segments'], storages['splits'], storages['segments']),
+        SegmentSynchronizerAsync(apis['segments'], storages['splits'], storages['segments'], storages['rule_based_segments']),
         ImpressionSynchronizerAsync(apis['impressions'], storages['impressions'],
                                cfg['impressionsBulkSize']),
         EventSynchronizerAsync(apis['events'], storages['events'], cfg['eventsBulkSize']),
diff --git a/splitio/models/rule_based_segments.py b/splitio/models/rule_based_segments.py
@@ -152,6 +152,14 @@ def get_excluded_segments(self):
         """Return excluded segments"""
         return self._segments
 
+    def get_excluded_standard_segments(self):
+        """Return excluded segments"""
+        to_return = []
+        for segment in self._segments:
+            if segment.type == SegmentType.STANDARD:
+                to_return.append(segment.name)
+        return to_return
+
     def to_json(self):
         """Return a JSON representation of this object."""
         return {
diff --git a/splitio/storage/inmemmory.py b/splitio/storage/inmemmory.py
@@ -233,7 +233,7 @@ def contains(self, segment_names):
         
     def fetch_many(self, segment_names):
         return {rb_segment_name: self.get(rb_segment_name) for rb_segment_name in segment_names}
-
+        
 class InMemoryRuleBasedSegmentStorageAsync(RuleBasedSegmentsStorage):
     """InMemory implementation of a feature flag storage base."""    
     def __init__(self):
diff --git a/splitio/sync/segment.py b/splitio/sync/segment.py
@@ -10,6 +10,7 @@
 from splitio.util.backoff import Backoff
 from splitio.optional.loaders import asyncio, aiofiles
 from splitio.sync import util
+from splitio.util.storage_helper import get_standard_segment_names_in_rbs_storage, get_standard_segment_names_in_rbs_storage_async
 from splitio.optional.loaders import asyncio
 
 _LOGGER = logging.getLogger(__name__)
@@ -22,7 +23,7 @@
 
 
 class SegmentSynchronizer(object):
-    def __init__(self, segment_api, feature_flag_storage, segment_storage):
+    def __init__(self, segment_api, feature_flag_storage, segment_storage, rule_based_segment_storage):
         """
         Class constructor.
 
@@ -39,6 +40,7 @@ def __init__(self, segment_api, feature_flag_storage, segment_storage):
         self._api = segment_api
         self._feature_flag_storage = feature_flag_storage
         self._segment_storage = segment_storage
+        self._rule_based_segment_storage = rule_based_segment_storage
         self._worker_pool = workerpool.WorkerPool(_MAX_WORKERS, self.synchronize_segment)
         self._worker_pool.start()
         self._backoff = Backoff(
@@ -181,9 +183,12 @@ def synchronize_segments(self, segment_names = None, dont_wait = False):
         :rtype: bool
         """
         if segment_names is None:
-            segment_names = self._feature_flag_storage.get_segment_names()
+            segment_names = set(self._feature_flag_storage.get_segment_names())
+            segment_names.update(get_standard_segment_names_in_rbs_storage(self._rule_based_segment_storage))
 
         for segment_name in segment_names:
+            _LOGGER.debug("Adding segment name to sync worker")
+            _LOGGER.debug(segment_name)
             self._worker_pool.submit_work(segment_name)
         if (dont_wait):
             return True
@@ -204,7 +209,7 @@ def segment_exist_in_storage(self, segment_name):
 
 
 class SegmentSynchronizerAsync(object):
-    def __init__(self, segment_api, feature_flag_storage, segment_storage):
+    def __init__(self, segment_api, feature_flag_storage, segment_storage, rule_based_segment_storage):
         """
         Class constructor.
 
@@ -221,6 +226,7 @@ def __init__(self, segment_api, feature_flag_storage, segment_storage):
         self._api = segment_api
         self._feature_flag_storage = feature_flag_storage
         self._segment_storage = segment_storage
+        self._rule_based_segment_storage = rule_based_segment_storage
         self._worker_pool = workerpool.WorkerPoolAsync(_MAX_WORKERS, self.synchronize_segment)
         self._worker_pool.start()
         self._backoff = Backoff(
@@ -364,7 +370,8 @@ async def synchronize_segments(self, segment_names = None, dont_wait = False):
         :rtype: bool
         """
         if segment_names is None:
-            segment_names = await self._feature_flag_storage.get_segment_names()
+            segment_names = set(await self._feature_flag_storage.get_segment_names())
+            segment_names.update(await get_standard_segment_names_in_rbs_storage_async(self._rule_based_segment_storage))
 
         self._jobs = await self._worker_pool.submit_work(segment_names)
         if (dont_wait):
diff --git a/splitio/sync/split.py b/splitio/sync/split.py
@@ -139,7 +139,7 @@ def _fetch_until(self, fetch_options, till=None, rbs_till=None):
             rbs_segment_list = update_rule_based_segment_storage(self._rule_based_segment_storage, fetched_rule_based_segments, feature_flag_changes.get('rbs')['t'], self._api.clear_storage)
             
             fetched_feature_flags = [(splits.from_raw(feature_flag)) for feature_flag in feature_flag_changes.get('ff').get('d', [])]
-            segment_list = update_feature_flag_storage(self._feature_flag_storage, fetched_feature_flags, feature_flag_changes.get('ff')['t'], self._api.clear_storage)
+            segment_list.update(update_feature_flag_storage(self._feature_flag_storage, fetched_feature_flags, feature_flag_changes.get('ff')['t'], self._api.clear_storage))
             segment_list.update(rbs_segment_list)
             
             if feature_flag_changes.get('ff')['t'] == feature_flag_changes.get('ff')['s'] and feature_flag_changes.get('rbs')['t'] == feature_flag_changes.get('rbs')['s']:
diff --git a/splitio/util/storage_helper.py b/splitio/util/storage_helper.py
@@ -1,6 +1,7 @@
 """Storage Helper."""
 import logging
 from splitio.models import splits
+from splitio.models import rule_based_segments
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -58,17 +59,29 @@ def update_rule_based_segment_storage(rule_based_segment_storage, rule_based_seg
     for rule_based_segment in rule_based_segments:
         if rule_based_segment.status == splits.Status.ACTIVE:
             to_add.append(rule_based_segment)
-            segment_list.update(set(_get_segment_names(rule_based_segment.excluded.get_excluded_segments())))
+            segment_list.update(set(rule_based_segment.excluded.get_excluded_standard_segments()))
             segment_list.update(rule_based_segment.get_condition_segment_names())
         else:
             if rule_based_segment_storage.get(rule_based_segment.name) is not None:
                 to_delete.append(rule_based_segment.name)
 
     rule_based_segment_storage.update(to_add, to_delete, change_number)
     return segment_list
+    
+def get_standard_segment_names_in_rbs_storage(rule_based_segment_storage):
+    """
+    Retrieve a list of all standard segments names.
 
-def _get_segment_names(excluded_segments):
-    return [excluded_segment.name for excluded_segment in excluded_segments]
+    :return: Set of segment names.
+    :rtype: Set(str)
+    """
+    segment_list = set()
+    for rb_segment in rule_based_segment_storage.get_segment_names():
+        rb_segment_obj = rule_based_segment_storage.get(rb_segment)
+        segment_list.update(set(rb_segment_obj.excluded.get_excluded_standard_segments()))
+        segment_list.update(rb_segment_obj.get_condition_segment_names())
+        
+    return segment_list
     
 async def update_feature_flag_storage_async(feature_flag_storage, feature_flags, change_number, clear_storage=False):
     """
@@ -124,7 +137,7 @@ async def update_rule_based_segment_storage_async(rule_based_segment_storage, ru
     for rule_based_segment in rule_based_segments:
         if rule_based_segment.status == splits.Status.ACTIVE:
             to_add.append(rule_based_segment)
-            segment_list.update(set(_get_segment_names(rule_based_segment.excluded.get_excluded_segments())))
+            segment_list.update(set(rule_based_segment.excluded.get_excluded_standard_segments()))
             segment_list.update(rule_based_segment.get_condition_segment_names())
         else:
             if await rule_based_segment_storage.get(rule_based_segment.name) is not None:
@@ -133,6 +146,22 @@ async def update_rule_based_segment_storage_async(rule_based_segment_storage, ru
     await rule_based_segment_storage.update(to_add, to_delete, change_number)
     return segment_list
 
+async def get_standard_segment_names_in_rbs_storage_async(rule_based_segment_storage):
+    """
+    Retrieve a list of all standard segments names.
+
+    :return: Set of segment names.
+    :rtype: Set(str)
+    """
+    segment_list = set()
+    segment_names = await rule_based_segment_storage.get_segment_names()
+    for rb_segment in segment_names:
+        rb_segment_obj = await rule_based_segment_storage.get(rb_segment)
+        segment_list.update(set(rb_segment_obj.excluded.get_excluded_standard_segments()))
+        segment_list.update(rb_segment_obj.get_condition_segment_names())
+        
+    return segment_list
+
 def get_valid_flag_sets(flag_sets, flag_set_filter):
     """
     Check each flag set in given array, return it if exist in a given config flag set array, if config array is empty return all
diff --git a/tests/api/test_splits_api.py b/tests/api/test_splits_api.py
@@ -122,6 +122,41 @@ def get(sdk, splitChanges, sdk_key, extra_headers, query):
         assert self.query[2] == {'s': '1.3', 'since': 123, 'rbSince': -1}
         assert response == {"ff": {"d": [], "s": 123, "t": 456}, "rbs": {"d": [], "s": 123, "t": -1}}
         assert split_api.clear_storage
+
+    def test_using_old_spec_since(self, mocker):
+        """Test using old_spec_since variable."""
+        httpclient = mocker.Mock(spec=client.HttpClient)
+        self.counter = 0
+        self.query = []
+        def get(sdk, splitChanges, sdk_key, extra_headers, query):
+            self.counter += 1
+            self.query.append(query)
+            if self.counter == 1:
+                return client.HttpResponse(400, 'error', {})
+            if self.counter == 2:
+                return client.HttpResponse(200, '{"splits": [], "since": 123, "till": 456}', {})
+            if self.counter == 3:
+                return client.HttpResponse(400, 'error', {})
+            if self.counter == 4:
+                return client.HttpResponse(200, '{"splits": [], "since": 456, "till": 456}', {})
+            
+        httpclient.is_sdk_endpoint_overridden.return_value = True
+        httpclient.get = get
+        split_api = splits.SplitsAPI(httpclient, 'some_api_key', SdkMetadata('1.0', 'some', '1.2.3.4'), mocker.Mock())
+        response = split_api.fetch_splits(123, -1, FetchOptions(False, None, None, None))
+        assert response == {"ff": {"d": [], "s": 123, "t": 456}, "rbs": {"d": [], "s": -1, "t": -1}}
+        assert self.query == [{'s': '1.3', 'since': 123, 'rbSince': -1}, {'s': '1.1', 'since': 123}]
+        assert not split_api.clear_storage
+
+        time.sleep(1)
+        splits._PROXY_CHECK_INTERVAL_MILLISECONDS_SS = 10
+        
+        response = split_api.fetch_splits(456, -1, FetchOptions(False, None, None, None))
+        time.sleep(1)        
+        splits._PROXY_CHECK_INTERVAL_MILLISECONDS_SS = 1000000
+        assert self.query[2] == {'s': '1.3', 'since': 456, 'rbSince': -1}
+        assert self.query[3] == {'s': '1.1', 'since': 456}
+        assert response == {"ff": {"d": [], "s": 456, "t": 456}, "rbs": {"d": [], "s": -1, "t": -1}}
         
 class SplitAPIAsyncTests(object):
     """Split async API test cases."""
@@ -253,9 +288,45 @@ async def get(sdk, splitChanges, sdk_key, extra_headers, query):
         assert self.query == [{'s': '1.3', 'since': 123, 'rbSince': -1}, {'s': '1.1', 'since': 123}]
         assert not split_api.clear_storage
 
-        time.sleep(1)
+        time.sleep(1) 
         splits._PROXY_CHECK_INTERVAL_MILLISECONDS_SS = 10
         response = await split_api.fetch_splits(123, -1, FetchOptions(False, None, None, None))
         assert self.query[2] == {'s': '1.3', 'since': 123, 'rbSince': -1}
         assert response == {"ff": {"d": [], "s": 123, "t": 456}, "rbs": {"d": [], "s": 123, "t": -1}}
         assert split_api.clear_storage
+
+    @pytest.mark.asyncio
+    async def test_using_old_spec_since(self, mocker):
+        """Test using old_spec_since variable."""
+        httpclient = mocker.Mock(spec=client.HttpClient)
+        self.counter = 0
+        self.query = []
+        async def get(sdk, splitChanges, sdk_key, extra_headers, query):
+            self.counter += 1
+            self.query.append(query)
+            if self.counter == 1:
+                return client.HttpResponse(400, 'error', {})
+            if self.counter == 2:
+                return client.HttpResponse(200, '{"splits": [], "since": 123, "till": 456}', {})
+            if self.counter == 3:
+                return client.HttpResponse(400, 'error', {})
+            if self.counter == 4:
+                return client.HttpResponse(200, '{"splits": [], "since": 456, "till": 456}', {})
+            
+        httpclient.is_sdk_endpoint_overridden.return_value = True
+        httpclient.get = get
+        split_api = splits.SplitsAPIAsync(httpclient, 'some_api_key', SdkMetadata('1.0', 'some', '1.2.3.4'), mocker.Mock())
+        response = await split_api.fetch_splits(123, -1, FetchOptions(False, None, None, None))
+        assert response == {"ff": {"d": [], "s": 123, "t": 456}, "rbs": {"d": [], "s": -1, "t": -1}}
+        assert self.query == [{'s': '1.3', 'since': 123, 'rbSince': -1}, {'s': '1.1', 'since': 123}]
+        assert not split_api.clear_storage
+
+        time.sleep(1)
+        splits._PROXY_CHECK_INTERVAL_MILLISECONDS_SS = 10
+        
+        response = await split_api.fetch_splits(456, -1, FetchOptions(False, None, None, None))
+        time.sleep(1)        
+        splits._PROXY_CHECK_INTERVAL_MILLISECONDS_SS = 1000000
+        assert self.query[2] == {'s': '1.3', 'since': 456, 'rbSince': -1}
+        assert self.query[3] == {'s': '1.1', 'since': 456}
+        assert response == {"ff": {"d": [], "s": 456, "t": 456}, "rbs": {"d": [], "s": -1, "t": -1}}
diff --git a/tests/sync/test_segments_synchronizer.py b/tests/sync/test_segments_synchronizer.py
diff --git a/tests/sync/test_synchronizer.py b/tests/sync/test_synchronizer.py
diff --git a/tests/tasks/test_segment_sync.py b/tests/tasks/test_segment_sync.py
diff --git a/tests/util/test_storage_helper.py b/tests/util/test_storage_helper.py
