diff --git a/datadog_checks_base/changelog.d/24442.added b/datadog_checks_base/changelog.d/24442.added new file mode 100644 index 0000000000000..0f61031827c47 --- /dev/null +++ b/datadog_checks_base/changelog.d/24442.added @@ -0,0 +1 @@ +Add an async job registry to `DatabaseCheck` so DBM integrations can register `DBMAsyncJob`s additively and run, cancel, or shut them all down through a single entry point (`register_async_job`, `run_async_jobs`, `cancel_async_jobs`, `shutdown_async_jobs`), along with a `DBMAsyncJob.shutdown()` hook for releasing lifetime-scoped resources on unschedule. diff --git a/datadog_checks_base/datadog_checks/base/checks/db.py b/datadog_checks_base/datadog_checks/base/checks/db.py index 6c87ad9113ee9..d604bbfba8212 100644 --- a/datadog_checks_base/datadog_checks/base/checks/db.py +++ b/datadog_checks_base/datadog_checks/base/checks/db.py @@ -4,14 +4,22 @@ from abc import abstractmethod from string import Template +from typing import TYPE_CHECKING, Dict, List from datadog_checks.base.agent import datadog_agent from datadog_checks.base.utils.db.utils import TagManager from . import AgentCheck +if TYPE_CHECKING: + from datadog_checks.base.utils.db.utils import DBMAsyncJob + class DatabaseCheck(AgentCheck): + """ + Base class for Database Monitoring (DBM) integrations. + """ + #: Authoritative DBM platform identifier for this integration. #: Subclasses should set this explicitly; it is the value surfaced by #: :attr:`dbms` and used across DBM payloads, metric name prefixes and async jobs. @@ -23,6 +31,49 @@ def __init__(self, *args, **kwargs): self._database_identifier = None self._dbms_fallback_warning_logged = False self.tag_manager = TagManager() + #: Async jobs owned by this check, keyed by job name, populated via + #: :meth:`register_async_job`. + self._async_job_registry: Dict[str, "DBMAsyncJob"] = {} + + def register_async_job(self, job: "DBMAsyncJob") -> "DBMAsyncJob": + """ + Register ``job`` under its ``job_name`` so the check manages its lifecycle, and return it + unchanged. + + Registering a job whose name matches an already-registered job replaces it. Raises + ``ValueError`` if the job has no name. + """ + if job.job_name is None: + raise ValueError("Cannot register an async job without a job_name") + self._async_job_registry[job.job_name] = job + return job + + def run_async_jobs(self, tags: List[str]) -> None: + """Run each registered job's loop, forwarding ``tags`` to every job.""" + for job in self._async_job_registry.values(): + job.run_job_loop(tags) + + def cancel_async_jobs(self) -> None: + """ + Signal every registered job to stop, without waiting for loops to finish or releasing + resources. + + Safe to call while ``check()`` is running. Follow with :meth:`shutdown_async_jobs` to wait + for the loops and release resources. + """ + for job in self._async_job_registry.values(): + job.cancel() + + def shutdown_async_jobs(self) -> None: + """ + Wait for every registered job's loop to finish (:meth:`~DBMAsyncJob.wait_for_completion`) + and run its teardown (:meth:`~DBMAsyncJob.shutdown`). + + Must not run concurrently with ``check()``. + """ + for job in self._async_job_registry.values(): + job.wait_for_completion() + job.shutdown() def database_monitoring_query_sample(self, raw_event: str): self.event_platform_event(raw_event, "dbm-samples") diff --git a/datadog_checks_base/datadog_checks/base/utils/db/utils.py b/datadog_checks_base/datadog_checks/base/utils/db/utils.py index 8bb52064b1310..74e705a785d02 100644 --- a/datadog_checks_base/datadog_checks/base/utils/db/utils.py +++ b/datadog_checks_base/datadog_checks/base/utils/db/utils.py @@ -345,12 +345,26 @@ def __init__( if self._features is None: self._features = [None] + @property + def job_name(self) -> Optional[str]: + """The job's name""" + return self._job_name + def cancel(self): """ Send a signal to cancel the job loop asynchronously. """ self._cancel_event.set() + def wait_for_completion(self) -> None: + """ + Block until the job loop has finished running, then clear its future. No-op if the loop is + not running. Typically called after :meth:`cancel` to wait for the loop to stop. + """ + if self._job_loop_future: + self._job_loop_future.result() + self._job_loop_future = None + def run_job_loop(self, tags): """ :param tags: @@ -467,6 +481,8 @@ def _job_loop(self): ) finally: self._log.info("[%s] Shutting down job loop", self._job_tags_str) + # Runs on every loop exit, including the inactivity stop above, after which the loop may + # restart on the next check run. For one-time teardown on unschedule, override shutdown(). if self._shutdown_callback: self._shutdown_callback() @@ -498,6 +514,19 @@ def _run_job_traced(self): def run_job(self): raise NotImplementedError() + def shutdown(self) -> None: + """ + Release resources the job holds for its whole lifetime, such as dedicated DB connections or + clients. + + Called once when the owning check is unscheduled, after the loop has stopped. The default + is a no-op; override to close long-lived resources, and keep the implementation idempotent. + + Unlike ``shutdown_callback``, which runs on every loop exit and may be followed by a + restart, this runs only during final teardown. + """ + pass + @contextlib.contextmanager def tracked_query(check, operation, tags=None): diff --git a/datadog_checks_base/tests/base/checks/test_database_check.py b/datadog_checks_base/tests/base/checks/test_database_check.py index 95937d5fbd635..8e142f2819409 100644 --- a/datadog_checks_base/tests/base/checks/test_database_check.py +++ b/datadog_checks_base/tests/base/checks/test_database_check.py @@ -1,12 +1,14 @@ # (C) Datadog, Inc. 2026-present # All rights reserved # Licensed under a 3-clause BSD style license (see LICENSE) +from concurrent.futures.thread import ThreadPoolExecutor from unittest import mock import pytest from datadog_checks.base.checks.db import DatabaseCheck from datadog_checks.base.stubs.datadog_agent import datadog_agent +from datadog_checks.base.utils.db.utils import DBMAsyncJob class FakeDatabaseCheck(DatabaseCheck): @@ -23,6 +25,43 @@ def cloud_metadata(self): return {} +class RegistryTestJob(DBMAsyncJob): + """Minimal DBMAsyncJob used to exercise the DatabaseCheck async job registry.""" + + def __init__(self, check, enabled=True, job_name="test-job"): + super().__init__( + check, + enabled=enabled, + dbms="test-dbms", + rate_limit=10, + max_sleep_chunk_s=0.1, + job_name=job_name, + ) + self.shutdown_calls = 0 + + def shutdown(self): + self.shutdown_calls += 1 + + def run_job(self): + pass + + +@pytest.fixture +def registry_check(): + check = FakeDatabaseCheck("test", {}, [{}]) + yield check + # Stop any registered jobs so their loops don't outlive the test. + check.cancel_async_jobs() + check.shutdown_async_jobs() + + +@pytest.fixture(autouse=True) +def stop_orphaned_threads(): + # Recreate the shared executor per test so job loops don't leak across tests. + DBMAsyncJob.executor.shutdown(wait=True) + DBMAsyncJob.executor = ThreadPoolExecutor() + + def test_agent_hostname_resolves_once_and_caches(): check = FakeDatabaseCheck("test", {}, [{}]) # The hostname comes from an FFI call, so it should only be resolved once and cached. @@ -89,3 +128,64 @@ def database_identifier_params(self): if tags_after is not None: check.tag_manager.set_tags_from_list(tags_after, replace=True) assert check.database_identifier == expected + + +@pytest.mark.parametrize("register_twice", [False, True], ids=["single", "duplicate_instance"]) +def test_register_async_job_adds_and_dedupes(registry_check, register_twice): + job = RegistryTestJob(registry_check) + assert registry_check.register_async_job(job) is job + if register_twice: + assert registry_check.register_async_job(job) is job + assert registry_check._async_job_registry == {"test-job": job} + + +def test_register_async_job_replaces_job_with_same_name(registry_check): + first = RegistryTestJob(registry_check, job_name="query-metrics") + second = RegistryTestJob(registry_check, job_name="query-metrics") + + registry_check.register_async_job(first) + registry_check.register_async_job(second) + + assert registry_check._async_job_registry == {"query-metrics": second} + + +def test_register_async_job_requires_job_name(registry_check): + with pytest.raises(ValueError): + registry_check.register_async_job(RegistryTestJob(registry_check, job_name=None)) + + +@pytest.mark.parametrize("enabled", [True, False], ids=["enabled", "disabled"]) +def test_run_async_jobs_starts_only_enabled_jobs(registry_check, enabled): + job = registry_check.register_async_job(RegistryTestJob(registry_check, enabled=enabled)) + + registry_check.run_async_jobs([]) + + # Only enabled jobs get a running loop; disabled jobs are skipped by run_job_loop. + assert (job._job_loop_future is not None) == enabled + + +def test_cancel_async_jobs_signals_without_touching_futures(registry_check): + job = registry_check.register_async_job(RegistryTestJob(registry_check)) + registry_check.run_async_jobs([]) + assert job._job_loop_future is not None + + registry_check.cancel_async_jobs() + + # cancel_async_jobs only sets the cancel event; the future stays in place for shutdown to await. + assert job._cancel_event.is_set() + assert job._job_loop_future is not None + + +@pytest.mark.parametrize("started", [True, False], ids=["loop_started", "loop_not_started"]) +def test_shutdown_async_jobs_tears_down_and_calls_shutdown(registry_check, started): + job = registry_check.register_async_job(RegistryTestJob(registry_check)) + if started: + registry_check.run_async_jobs([]) + assert job._job_loop_future is not None + registry_check.cancel_async_jobs() + + registry_check.shutdown_async_jobs() + + # 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