feat: pause/start dcs (#474)

exposed API for people to control polling DCS

Author: weihao-statsig

statsig/spec_updater.py

@@ -46,6 +46,9 @@ def __init__(
         self._statsig_metadata = statsig_metadata
         self._background_download_configs = None
         self._background_download_id_lists = None
+        # DCS polling: set => polling enabled, clear => paused
+        self._dcs_polling_enabled_event = threading.Event()
+        self._dcs_polling_enabled_event.set()
         self._config_sync_strategies = self._get_sync_dcs_strategies()
         self._dcs_process_lock = threading.Lock()
         if options.out_of_sync_threshold_in_s is not None:
@@ -409,7 +412,7 @@ def sync_config_spec():
         self._background_download_configs = spawn_background_thread(
             "bg_download_config_specs",
             self._sync,
-            (sync_config_spec, interval, fast_start),
+            (sync_config_spec, interval, fast_start, self._dcs_polling_enabled_event),
             self._error_boundary,
         )
 
@@ -422,18 +425,33 @@ def _spawn_bg_poll_id_lists(self):
             self._error_boundary,
         )
 
-    def _sync(self, sync_func, interval, fast_start=False):
-        if fast_start:
+    def _sync(
+        self,
+        sync_func,
+        interval,
+        fast_start=False,
+        enabled_event: Optional[threading.Event] = None,
+    ):
+        if fast_start and (enabled_event is None or enabled_event.is_set()):
             sync_func()
 
         while True:
             try:
-                if self._shutdown_event.wait(interval):
-                    break
+                # Wait until next interval, unless shutdown is requested.
+                if self._shutdown_event.wait(timeout=interval):
+                    return
+                if enabled_event is not None and not enabled_event.is_set():
+                    continue
                 sync_func()
             except Exception as e:
                 self._error_boundary.log_exception("_sync", e)
 
+    def pause_polling_dcs(self):
+        self._dcs_polling_enabled_event.clear()
+
+    def start_polling_dcs(self):
+        self._dcs_polling_enabled_event.set()
+
     def _get_sync_dcs_strategies(self) -> List[DataSource]:
         try:
             if self._options.config_sync_sources is not None:

statsig/statsig.py

@@ -327,6 +327,20 @@ def shutdown():
     __instance.shutdown()
 
 
+def pause_polling_dcs():
+    """
+    Pauses background polling for Download Config Specs (DCS).
+    """
+    __instance.pause_polling_dcs()
+
+
+def start_polling_dcs():
+    """
+    Resumes background polling for Download Config Specs (DCS).
+    """
+    __instance.start_polling_dcs()
+
+
 def get_instance():
     """
     Returns the Statsig instance

statsig/statsig_server.py

@@ -415,6 +415,24 @@ def task():
 
         self._errorBoundary.swallow("shutdown", task)
 
+    def pause_polling_dcs(self):
+        def task():
+            if not self._initialized:
+                raise StatsigRuntimeError("Must call initialize before pausing DCS polling")
+            if self._spec_store is not None:
+                self._spec_store.spec_updater.pause_polling_dcs()
+
+        self._errorBoundary.swallow("pause_polling_dcs", task)
+
+    def start_polling_dcs(self):
+        def task():
+            if not self._initialized:
+                raise StatsigRuntimeError("Must call initialize before starting DCS polling")
+            if self._spec_store is not None:
+                self._spec_store.spec_updater.start_polling_dcs()
+
+        self._errorBoundary.swallow("start_polling_dcs", task)
+
     def override_gate(self, gate: str, value: bool, user_id: Optional[str] = None):
         self._errorBoundary.swallow(
             "override_gate", lambda: self._evaluator.override_gate(gate, value, user_id)

tests/test_pause_polling_dcs.py

@@ -0,0 +1,77 @@
+import json
+import os
+import time
+import unittest
+from unittest.mock import patch
+
+from network_stub import NetworkStub
+from statsig import StatsigOptions, statsig
+
+_network_stub = NetworkStub("http://test-pause-polling-dcs")
+
+with open(os.path.join(os.path.abspath(os.path.dirname(__file__)), '../testdata/download_config_specs.json')) as r:
+    CONFIG_SPECS_RESPONSE = r.read()
+PARSED_CONFIG_SPEC = json.loads(CONFIG_SPECS_RESPONSE)
+
+
+@patch('requests.Session.request', side_effect=_network_stub.mock)
+class TestPausePollingDcs(unittest.TestCase):
+    @classmethod
+    @patch('requests.Session.request', side_effect=_network_stub.mock)
+    def setUpClass(cls, mock_proxy):
+        cls.dcs_called = False
+        cls.dcs_call_count = 0
+
+    def setUp(self):
+        self.__class__.dcs_called = False
+        self.__class__.dcs_call_count = 0
+
+        def dcs_cb(url, **kwargs):
+            self.__class__.dcs_called = True
+            self.__class__.dcs_call_count += 1
+            return PARSED_CONFIG_SPEC
+
+        _network_stub.stub_request_with_value("get_id_lists", 200, {})
+        _network_stub.stub_request_with_function("download_config_specs/.*", 200, dcs_cb)
+
+    def tearDown(self):
+        statsig.shutdown()
+        _network_stub.reset()
+
+    def test_pause_and_resume_dcs_polling(self, request_mock):
+        options = StatsigOptions(api=_network_stub.host, rulesets_sync_interval=1)
+        statsig.initialize("secret-key", options)
+
+        paused_wait_s = 3.0
+        resume_timeout_s = 5.0
+
+        def wait_for_new_dcs_calls(prev_count: int, timeout_s: float) -> bool:
+            deadline = time.time() + timeout_s
+            while time.time() < deadline:
+                if self.__class__.dcs_call_count > prev_count:
+                    return True
+                time.sleep(0.05)
+            return False
+
+        # ignore the initialize-time DCS fetch
+        self.__class__.dcs_called = False
+        self.__class__.dcs_call_count = 0
+
+        # OFF -> ON -> OFF -> ON
+        statsig.pause_polling_dcs()
+        prev = self.__class__.dcs_call_count
+        time.sleep(paused_wait_s)
+        self.assertEqual(prev, self.__class__.dcs_call_count)
+
+        statsig.start_polling_dcs()
+        prev = self.__class__.dcs_call_count
+        self.assertTrue(wait_for_new_dcs_calls(prev, timeout_s=resume_timeout_s))
+
+        statsig.pause_polling_dcs()
+        prev = self.__class__.dcs_call_count
+        time.sleep(paused_wait_s)
+        self.assertEqual(prev, self.__class__.dcs_call_count)
+
+        statsig.start_polling_dcs()
+        prev = self.__class__.dcs_call_count
+        self.assertTrue(wait_for_new_dcs_calls(prev, timeout_s=resume_timeout_s))