[DBMON-6882] Add cancellation lifecycle to DatabaseCheck

datadog_checks_base/changelog.d/24844.added

@@ -0,0 +1 @@
+Add a cancellation lifecycle to ``DatabaseCheck`` so DBM integrations get thread-safe ``cancel()`` handling with a ``shutdown()`` teardown hook.

datadog_checks_base/datadog_checks/base/checks/db.py

@@ -2,6 +2,7 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
+import threading
 from abc import abstractmethod
 from string import Template
 from typing import TYPE_CHECKING, Dict, List
@@ -34,6 +35,12 @@ def __init__(self, *args, **kwargs):
         #: Async jobs owned by this check, keyed by job name, populated via
         #: :meth:`register_async_job`.
         self._async_job_registry: Dict[str, "DBMAsyncJob"] = {}
+        # Guards the cancellation state below, which `cancel()` and `run()` read and write from
+        # different threads.
+        self._cancel_lock = threading.Lock()
+        self._is_running = False
+        self._cancelled = False
+        self._finalized = False
 
     def register_async_job(self, job: "DBMAsyncJob") -> "DBMAsyncJob":
         """
@@ -49,7 +56,15 @@ def register_async_job(self, job: "DBMAsyncJob") -> "DBMAsyncJob":
         return job
 
     def run_async_jobs(self, tags: List[str]) -> None:
-        """Run each registered job's loop, forwarding ``tags`` to every job."""
+        """
+        Run each registered job's loop, forwarding ``tags`` to every job.
+
+        No-op once the check has been cancelled, so a ``check()`` that was in flight when the
+        cancel arrived does not restart the loops :meth:`cancel_async_jobs` has just stopped.
+        """
+        if self._cancelled:
+            self.log.debug("Not running async jobs, check has been cancelled")
+            return
         for job in self._async_job_registry.values():
             job.run_job_loop(tags)
 
@@ -75,6 +90,101 @@ def shutdown_async_jobs(self) -> None:
             job.wait_for_completion()
             job.shutdown()
 
+    @property
+    def is_cancelled(self) -> bool:
+        """
+        Whether :meth:`cancel` has been signaled.
+
+        ``check()`` implementations should consult this before starting work that would outlive the
+        run, and long-running collection loops should poll it so they stop promptly.
+        """
+        return self._cancelled
+
+    def run(self) -> str:
+        """
+        Run the check, recording whether it is in flight so :meth:`cancel` knows whether it may
+        tear the check down right away.
+
+        Returns an empty error report without running the check once it has been cancelled. When a
+        cancel arrives mid-run, the deferred teardown happens here, after ``check()`` returns.
+        """
+        with self._cancel_lock:
+            if self._cancelled:
+                self.log.debug("run() skipped, check already cancelled")
+                return ''
+            self._is_running = True
+        try:
+            return super().run()
+        finally:
+            with self._cancel_lock:
+                self._is_running = False
+                needs_finalize = self._cancelled
+            if needs_finalize:
+                self.log.debug("Cancel was signaled during the run, finalizing now that run() is complete")
+                self._finalize()
+
+    def cancel(self) -> None:
+        """
+        Signal that the check is being unscheduled.
+
+        The Agent may call this from another thread while ``check()`` is running, so this method
+        does no destructive work itself: closing a connection or dropping state that ``check()``
+        still depends on can crash the underlying client library when the run resumes. It only
+        signals the async jobs and records the cancellation, deferring teardown to
+        :meth:`_finalize`, which runs here when the check is idle and in :meth:`run` otherwise.
+
+        Integrations release their own resources by overriding :meth:`shutdown`, not this method.
+        """
+        self.log.debug("Marking check as cancelled")
+        self.cancel_async_jobs()
+        with self._cancel_lock:
+            self._cancelled = True
+            needs_finalize = not self._is_running
+        if needs_finalize:
+            self.log.debug("cancel() finalizing immediately, check is idle")
+            self._finalize()
+        else:
+            self.log.debug("cancel() deferred finalize, check is still running")
+
+    def _finalize(self) -> None:
+        """
+        Tear the check down: stop the async jobs, let the integration release its resources, then
+        drop the state that keeps the check alive.
+
+        Runs at most once, and never concurrently with ``check()`` — :meth:`cancel` and :meth:`run`
+        between them guarantee that.
+        """
+        with self._cancel_lock:
+            if self._finalized:
+                return
+            self._finalized = True
+        self.log.debug("Finalizing check: stopping async jobs and releasing resources")
+        self.shutdown_async_jobs()
+        self.shutdown()
+        # Dropping these breaks the reference cycles that would otherwise keep the check, and
+        # everything it holds, from being reclaimed once the Agent lets go of it. The jobs are
+        # stopped by this point, so releasing them here is safe.
+        self._async_job_registry.clear()
+        self.check_initializations.clear()
+        self._diagnosis = None
+        self.log.debug("Check cleanup complete")
+        # Must come last: the logging adapter reads back through this attribute for checks whose
+        # check_id was never resolved, so anything logged after this would fail.
+        self.log.check = None
+
+    def shutdown(self) -> None:
+        """
+        Release the resources this check holds for its whole lifetime, such as connections,
+        connection pools and clients.
+
+        Called once by :meth:`_finalize` during teardown, after the registered async jobs have
+        stopped and never while ``check()`` is running. Stopping those jobs is handled separately
+        by :meth:`shutdown_async_jobs`; this hook covers only what the check itself owns. To
+        unschedule a check, call :meth:`cancel` rather than this method.
+
+        The default is a no-op.
+        """
+
     def database_monitoring_query_sample(self, raw_event: str):
         self.event_platform_event(raw_event, "dbm-samples")
 

datadog_checks_base/tests/base/checks/test_database_check.py

@@ -1,6 +1,9 @@
 # (C) Datadog, Inc. 2026-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
+import gc
+import threading
+import weakref
 from concurrent.futures.thread import ThreadPoolExecutor
 from unittest import mock
 
@@ -10,6 +13,10 @@
 from datadog_checks.base.stubs.datadog_agent import datadog_agent
 from datadog_checks.base.utils.db.utils import DBMAsyncJob
 
+# Upper bound for waits on another thread; the tests only ever wait for a signal that is
+# already on its way.
+WAIT_TIMEOUT = 5
+
 
 class FakeDatabaseCheck(DatabaseCheck):
     @property
@@ -46,6 +53,27 @@ def run_job(self):
         pass
 
 
+class LifecycleCheck(FakeDatabaseCheck):
+    """DatabaseCheck that records lifecycle calls and can hold check() open."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.check_calls = 0
+        self.shutdown_calls = 0
+        self.in_check = threading.Event()
+        # Cleared by tests that need a cancel to land while check() is still executing.
+        self.release_check = threading.Event()
+        self.release_check.set()
+
+    def check(self, _):
+        self.check_calls += 1
+        self.in_check.set()
+        assert self.release_check.wait(timeout=WAIT_TIMEOUT)
+
+    def shutdown(self):
+        self.shutdown_calls += 1
+
+
 @pytest.fixture
 def registry_check():
     check = FakeDatabaseCheck("test", {}, [{}])
@@ -189,3 +217,124 @@ def test_shutdown_async_jobs_tears_down_and_calls_shutdown(registry_check, start
     # The future is cleared and shutdown runs once, whether or not a loop was started.
     assert job._job_loop_future is None
     assert job.shutdown_calls == 1
+
+
+def test_run_without_cancel_leaves_check_usable():
+    check = LifecycleCheck("test", {}, [{}])
+
+    assert check.run() == ''
+
+    # An uncancelled run reports no error, leaves no in-flight state, and tears nothing down.
+    assert check.check_calls == 1
+    assert not check._is_running
+    assert not check.is_cancelled
+    assert check.shutdown_calls == 0
+
+
+def test_cancel_when_idle_finalizes_immediately():
+    check = LifecycleCheck("test", {}, [{}])
+    job = check.register_async_job(RegistryTestJob(check))
+    check.run_async_jobs([])
+
+    check.cancel()
+
+    # With no run in flight there is nothing to wait for, so the whole teardown runs inline.
+    assert check.is_cancelled
+    assert job._cancel_event.is_set()
+    assert job._job_loop_future is None
+    assert job.shutdown_calls == 1
+    assert check.shutdown_calls == 1
+
+
+def test_cancel_during_run_defers_finalize_until_run_completes():
+    """Teardown must wait for check() to return.
+
+    Releasing resources under a running check() means the check resumes against a closed
+    connection, which crashes the Agent inside the database client library rather than raising.
+    """
+    check = LifecycleCheck("test", {}, [{}])
+    check.release_check.clear()
+    run_result = []
+    run_thread = threading.Thread(target=lambda: run_result.append(check.run()))
+    run_thread.start()
+    assert check.in_check.wait(timeout=WAIT_TIMEOUT)
+
+    check.cancel()
+
+    # The cancel is recorded, but check() is still executing so nothing may be released yet.
+    assert check.is_cancelled
+    assert check.shutdown_calls == 0
+
+    check.release_check.set()
+    run_thread.join(timeout=WAIT_TIMEOUT)
+
+    assert not run_thread.is_alive()
+    assert run_result == ['']
+    assert check.shutdown_calls == 1
+    assert not check._is_running
+
+
+def test_run_after_cancel_skips_the_check():
+    check = LifecycleCheck("test", {}, [{}])
+    check.cancel()
+
+    assert check.run() == ''
+    assert check.check_calls == 0
+
+
+@pytest.mark.parametrize("in_flight", [False, True], ids=["idle", "during_run"])
+def test_cancel_is_idempotent(in_flight):
+    check = LifecycleCheck("test", {}, [{}])
+    job = check.register_async_job(RegistryTestJob(check))
+    if in_flight:
+        check.release_check.clear()
+        run_thread = threading.Thread(target=check.run)
+        run_thread.start()
+        assert check.in_check.wait(timeout=WAIT_TIMEOUT)
+
+    check.cancel()
+    check.cancel()
+
+    if in_flight:
+        check.release_check.set()
+        run_thread.join(timeout=WAIT_TIMEOUT)
+        assert not run_thread.is_alive()
+    # Teardown runs once however many times the cancel is signaled, and whichever path runs it.
+    assert check.shutdown_calls == 1
+    assert job.shutdown_calls == 1
+
+
+def test_run_async_jobs_does_not_restart_jobs_after_cancel():
+    check = LifecycleCheck("test", {}, [{}])
+    job = check.register_async_job(RegistryTestJob(check))
+    check.cancel()
+
+    check.run_async_jobs([])
+
+    assert job._job_loop_future is None
+
+
+def test_check_is_reclaimed_after_cancel():
+    """Verify cancel() breaks every reference cycle, so refcounting alone reclaims the check.
+
+    If this fails, find which attribute on the reported referrer points back at the check and
+    clear it in _finalize(), or in the relevant shutdown() method.
+    """
+    check = LifecycleCheck("test", {}, [{}])
+    check.register_async_job(RegistryTestJob(check))
+    check.run_async_jobs([])
+    ref = weakref.ref(check)
+
+    check.cancel()
+
+    gc.collect()
+    gc.disable()
+    try:
+        del check
+        leaked = ref()
+        if leaked is not None:
+            referrers = [type(referrer).__name__ for referrer in gc.get_referrers(leaked)]
+            del leaked
+            pytest.fail(f"check still alive after cancel() and del -- pinned by: {referrers}")
+    finally:
+        gc.enable()