Improve throttling in sync jobs

datadog_checks_base/changelog.d/17716.added

@@ -0,0 +1 @@
+Improve throttling in sync jobs of `DBMAsyncJob` class. Instead of putting a thread to sleep check if the job is too early to run.

datadog_checks_base/datadog_checks/base/utils/db/utils.py

@@ -121,13 +121,20 @@ def __init__(self, rate_limit_s):
         self.period_s = 1.0 / self.rate_limit_s if self.rate_limit_s > 0 else 0
         self.last_event = 0
 
-    def sleep(self):
+    def update_last_time_and_sleep(self):
         """
         Sleeps long enough to enforce the rate limit
         """
         elapsed_s = time.time() - self.last_event
         sleep_amount = max(self.period_s - elapsed_s, 0)
         time.sleep(sleep_amount)
+        self.update_last_time()
+
+    def shall_execute(self):
+        elapsed_s = time.time() - self.last_event
+        return elapsed_s >= self.period_s
+
+    def update_last_time(self):
         self.last_event = time.time()
 
 
@@ -298,7 +305,7 @@ def run_job_loop(self, tags):
         self._last_check_run = time.time()
         if self._run_sync or is_affirmative(os.environ.get('DBM_THREADED_JOB_RUN_SYNC', "false")):
             self._log.debug("Running threaded job synchronously. job=%s", self._job_name)
-            self._run_job_rate_limited()
+            self._run_sync_job_rate_limited()
         elif self._job_loop_future is None or not self._job_loop_future.running():
             self._job_loop_future = DBMAsyncJob.executor.submit(self._job_loop)
         else:
@@ -362,11 +369,26 @@ def _job_loop(self):
     def _set_rate_limit(self, rate_limit):
         if self._rate_limiter.rate_limit_s != rate_limit:
             self._rate_limiter = ConstantRateLimiter(rate_limit)
+
+    def _run_sync_job_rate_limited(self):
+        if self._rate_limiter.shall_execute():
+            try:
+                self._run_job_traced()
+            except:
+                raise
+            finally:
+                self._rate_limiter.update_last_time()
 
     def _run_job_rate_limited(self):
-        self._run_job_traced()
-        if not self._cancel_event.isSet():
-            self._rate_limiter.sleep()
+        try:
+            self._run_job_traced()
+        except:
+            raise
+        finally:
+            if not self._cancel_event.isSet():
+                self._rate_limiter.update_last_time_and_sleep()
+            else:
+                self._rate_limiter.update_last_time()
 
     @_traced_dbm_async_job_method
     def _run_job_traced(self):

datadog_checks_base/tests/base/utils/db/test_util.py

@@ -53,12 +53,23 @@ def test_constant_rate_limiter():
     start = time.time()
     sleep_count = 0
     while time.time() - start < test_duration_s:
-        ratelimiter.sleep()
+        ratelimiter.update_last_time_and_sleep()
         sleep_count += 1
     max_expected_count = rate_limit * test_duration_s
     assert max_expected_count - 1 <= sleep_count <= max_expected_count + 1
 
 
+def test_constant_rate_limiter_shell_execute():
+    rate_limit = 1
+    test_duration_s = 0.1
+    ratelimiter = ConstantRateLimiter(rate_limit)
+    assert ratelimiter.shall_execute()
+    ratelimiter.update_last_time()
+    assert not ratelimiter.shall_execute()
+    time.sleep(1)
+    assert ratelimiter.shall_execute()
+
+
 def test_ratelimiting_ttl_cache():
     ttl = 0.1
     cache = RateLimitingTTLCache(maxsize=5, ttl=ttl)