[DBMON-5571] Regularly check for cancel event between intervals (DataDog#21150)

* Regularly check for cancel event between intervals

* Add type hinting

* Add changelog

* Ignore internal dd.mongo.async_job.cancel

* Avoid canceling the check too quickly

Author: eric-weaver

datadog_checks_base/changelog.d/21150.fixed

@@ -0,0 +1 @@
+Regularly check for cancel event between DBMAsyncJob check intervals

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

@@ -107,21 +107,33 @@ class ConstantRateLimiter:
     Basic rate limiter that sleeps long enough to ensure the rate limit is not exceeded. Not thread safe.
     """
 
-    def __init__(self, rate_limit_s):
+    def __init__(self, rate_limit_s, max_sleep_chunk_s=5):
         """
         :param rate_limit_s: rate limit in seconds
+        :param max_sleep_chunk_s: maximum size of each sleep chunk while waiting for the next period
         """
         self.rate_limit_s = max(rate_limit_s, 0)
         self.period_s = 1.0 / self.rate_limit_s if self.rate_limit_s > 0 else 0
         self.last_event = 0
+        self.max_sleep_chunk_s = max(0, max_sleep_chunk_s)
 
-    def update_last_time_and_sleep(self):
+    def update_last_time_and_sleep(self, cancel_event: Optional[threading.Event] = None):
         """
         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)
+        if self.period_s <= 0:
+            self.update_last_time()
+            return
+
+        deadline = self.last_event + self.period_s
+        while True:
+            now = time.time()
+            remaining = deadline - now
+            if remaining <= 0:
+                break
+            if cancel_event is not None and cancel_event.is_set():
+                break
+            time.sleep(min(remaining, self.max_sleep_chunk_s if self.max_sleep_chunk_s > 0 else remaining))
         self.update_last_time()
 
     def shall_execute(self):
@@ -275,6 +287,7 @@ def __init__(
         min_collection_interval=15,
         dbms="TODO",
         rate_limit=1,
+        max_sleep_chunk_s=1,
         run_sync=False,
         enabled=True,
         expected_db_exceptions=(),
@@ -295,7 +308,8 @@ def __init__(
         self._last_check_run = 0
         self._shutdown_callback = shutdown_callback
         self._dbms = dbms
-        self._rate_limiter = ConstantRateLimiter(rate_limit)
+        self._rate_limiter = ConstantRateLimiter(rate_limit, max_sleep_chunk_s=max_sleep_chunk_s)
+        self._max_sleep_chunk_s = max_sleep_chunk_s
         self._run_sync = run_sync
         self._enabled = enabled
         self._expected_db_exceptions = expected_db_exceptions
@@ -387,7 +401,7 @@ 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)
+            self._rate_limiter = ConstantRateLimiter(rate_limit, max_sleep_chunk_s=self._max_sleep_chunk_s)
 
     def _run_sync_job_rate_limited(self):
         if self._rate_limiter.shall_execute():
@@ -401,7 +415,7 @@ def _run_job_rate_limited(self):
             raise
         finally:
             if not self._cancel_event.is_set():
-                self._rate_limiter.update_last_time_and_sleep()
+                self._rate_limiter.update_last_time_and_sleep(cancel_event=self._cancel_event)
             else:
                 self._rate_limiter.update_last_time()
 

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

@@ -242,7 +242,14 @@ def _mock_obfuscate_sql(query, options=None):
 
 class JobForTesting(DBMAsyncJob):
     def __init__(
-        self, check, run_sync=False, enabled=True, rate_limit=10, min_collection_interval=15, job_execution_time=0
+        self,
+        check,
+        run_sync=False,
+        enabled=True,
+        rate_limit=10,
+        min_collection_interval=15,
+        job_execution_time=0,
+        max_sleep_chunk_s=5,
     ):
         super(JobForTesting, self).__init__(
             check,
@@ -253,6 +260,7 @@ def __init__(
             config_host="test-host",
             dbms="test-dbms",
             rate_limit=rate_limit,
+            max_sleep_chunk_s=max_sleep_chunk_s,
             job_name="test-job",
             shutdown_callback=self.test_shutdown,
         )
@@ -307,6 +315,20 @@ def test_dbm_async_job_cancel(aggregator):
     aggregator.assert_metric("dbm.async_job_test.shutdown")
 
 
+def test_dbm_async_job_cancel_returns_early_on_long_sleep():
+    # Configure a very low rate so the sleep interval would be ~60s without cancellation
+    job = JobForTesting(AgentCheck(), rate_limit=1 / 60.0, max_sleep_chunk_s=0.1)
+    job.run_job_loop([])
+    # Allow the thread to start and enter the sleep window
+    time.sleep(0.2)
+    start = time.time()
+    job.cancel()
+    # Should finish well before the full ~10s timeout and 60s rate-limiter interval
+    job._job_loop_future.result(timeout=10)
+    elapsed = time.time() - start
+    assert elapsed < 10, "Job did not cancel before the full sleep interval"
+
+
 def test_dbm_async_job_run_sync(aggregator):
     job = JobForTesting(AgentCheck(), run_sync=True)
     job.run_job_loop([])

mongo/tests/test_integration.py

@@ -107,6 +107,7 @@ def test_integration_mongos(instance_integration_cluster, aggregator, check, dd_
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
             'dd.mongo.operation.time',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -200,6 +201,7 @@ def test_integration_replicaset_primary_in_shard(instance_integration, aggregato
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -299,6 +301,7 @@ def test_integration_replicaset_secondary_in_shard(instance_integration, aggrega
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -355,6 +358,7 @@ def test_integration_replicaset_arbiter_in_shard(instance_integration, aggregato
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -419,6 +423,7 @@ def test_integration_configsvr_primary(instance_integration, aggregator, check,
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -516,6 +521,7 @@ def test_integration_configsvr_secondary(instance_integration, aggregator, check
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -583,6 +589,7 @@ def test_integration_replicaset_primary(instance_integration, aggregator, check,
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -688,6 +695,7 @@ def test_integration_replicaset_primary_config(instance_integration, aggregator,
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -797,6 +805,7 @@ def test_integration_replicaset_secondary(
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -852,6 +861,7 @@ def test_integration_replicaset_arbiter(instance_integration, aggregator, check,
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -908,6 +918,7 @@ def test_standalone(instance_integration, aggregator, check, dd_run_check):
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -974,6 +985,7 @@ def test_db_names_with_nonexistent_database(check, instance_integration, aggrega
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )
@@ -1009,6 +1021,7 @@ def test_db_names_missing_existent_database(check, instance_integration, aggrega
             'dd.custom.mongo.count',
             'dd.custom.mongo.query_a.amount',
             'dd.custom.mongo.query_a.el',
+            'dd.mongo.async_job.cancel',
         ],
         check_submission_type=True,
     )

mongo/tests/test_integration_shard.py

@@ -27,7 +27,9 @@ def test_mongo_arbiter(aggregator, check, instance_arbiter, dd_run_check):
         if metric_name in METRIC_VAL_CHECKS:
             metric = aggregator.metrics(metric_name)[0]
             assert METRIC_VAL_CHECKS[metric_name](metric.value)
-    aggregator.assert_metrics_using_metadata(get_metadata_metrics(), check_submission_type=True)
+    aggregator.assert_metrics_using_metadata(
+        get_metadata_metrics(), check_submission_type=True, exclude=['dd.mongo.async_job.cancel']
+    )
 
     expected_metrics = {
         'mongodb.replset.health': 1.0,
@@ -77,4 +79,6 @@ def test_mongo_replset(instance_shard, aggregator, check, dd_run_check):
         'mongodb.replset.optime_lag',
         tags=replset_common_tags + ['replset_state:secondary', 'member:shard01b:27019', 'replset_me:shard01a:27018'],
     )
-    aggregator.assert_metrics_using_metadata(get_metadata_metrics(), check_submission_type=True)
+    aggregator.assert_metrics_using_metadata(
+        get_metadata_metrics(), check_submission_type=True, exclude=['dd.mongo.async_job.cancel']
+    )

mongo/tests/test_integration_standalone.py

@@ -42,7 +42,9 @@ def test_mongo_authdb(aggregator, check, instance_authdb, dd_run_check):
         if metric_name in METRIC_VAL_CHECKS:
             metric = aggregator.metrics(metric_name)[0]
             assert METRIC_VAL_CHECKS[metric_name](metric.value)
-    aggregator.assert_metrics_using_metadata(get_metadata_metrics(), check_submission_type=True)
+    aggregator.assert_metrics_using_metadata(
+        get_metadata_metrics(), check_submission_type=True, exclude=['dd.mongo.async_job.cancel']
+    )
 
 
 @pytest.mark.parametrize(
@@ -63,7 +65,9 @@ def test_mongo_db_test(aggregator, check, instance_user, dd_run_check):
         if metric_name in METRIC_VAL_CHECKS:
             metric = aggregator.metrics(metric_name)[0]
             assert METRIC_VAL_CHECKS[metric_name](metric.value)
-    aggregator.assert_metrics_using_metadata(get_metadata_metrics(), check_submission_type=True)
+    aggregator.assert_metrics_using_metadata(
+        get_metadata_metrics(), check_submission_type=True, exclude=['dd.mongo.async_job.cancel']
+    )
 
 
 def test_mongo_old_config(aggregator, check, instance, dd_run_check):
@@ -77,7 +81,9 @@ def test_mongo_old_config(aggregator, check, instance, dd_run_check):
         if metric_name in METRIC_VAL_CHECKS_OLD:
             metric = aggregator.metrics(metric_name)[0]
             assert METRIC_VAL_CHECKS_OLD[metric_name](metric.value)
-    aggregator.assert_metrics_using_metadata(get_metadata_metrics(), check_submission_type=True)
+    aggregator.assert_metrics_using_metadata(
+        get_metadata_metrics(), check_submission_type=True, exclude=['dd.mongo.async_job.cancel']
+    )
 
 
 def test_mongo_dbstats_tag(aggregator, check, instance_dbstats_tag_dbname, dd_run_check):
@@ -91,7 +97,9 @@ def test_mongo_dbstats_tag(aggregator, check, instance_dbstats_tag_dbname, dd_ru
         if metric_name in METRIC_VAL_CHECKS:
             metric = aggregator.metrics(metric_name)[0]
             assert METRIC_VAL_CHECKS[metric_name](metric.value)
-    aggregator.assert_metrics_using_metadata(get_metadata_metrics(), check_submission_type=True)
+    aggregator.assert_metrics_using_metadata(
+        get_metadata_metrics(), check_submission_type=True, exclude=['dd.mongo.async_job.cancel']
+    )
 
     expected_metrics = {
         'mongodb.stats.avgobjsize': None,

sqlserver/tests/test_statements.py

@@ -382,16 +382,16 @@ def _obfuscate_sql(sql_query, options=None):
     # 3) emit the query metrics based on the diff of current and last state
     with mock.patch.object(datadog_agent, 'obfuscate_sql', passthrough=True) as mock_agent:
         mock_agent.side_effect = _obfuscate_sql
-        dd_run_check(check)
+        dd_run_check(check, cancel=False)
         for _ in range(0, exe_count):
             for params in param_groups:
                 bob_conn.execute_with_retries(query, params, database=database)
-        dd_run_check(check)
+        dd_run_check(check, cancel=False)
         aggregator.reset()
         for _ in range(0, exe_count):
             for params in param_groups:
                 bob_conn.execute_with_retries(query, params, database=database)
-        dd_run_check(check)
+        dd_run_check(check, cancel=False)
 
     _conn_key_prefix = "dbm-"
     with check.connection.open_managed_default_connection(key_prefix=_conn_key_prefix):