Move postgres async jobs onto the DatabaseCheck job registry (DataDog#24824)

* Move postgres async jobs onto the DatabaseCheck job registry

Replace the hand-rolled _async_jobs property, _cancel_async_jobs helper and
manual future/teardown loop with register_async_job, run_async_jobs,
cancel_async_jobs and shutdown_async_jobs from DatabaseCheck.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Add changelog

Co-authored-by: Cursor <cursoragent@cursor.com>

* Build postgres async jobs where they are registered

Only construct the jobs the configuration enables, instead of building all
four in __init__ and registering a subset. Unregistered jobs previously kept a
back-reference to the check that shutdown_async_jobs() never cleared.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Note that the job attributes only exist for tests

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: eric-weaver

postgres/changelog.d/24824.fixed

@@ -0,0 +1 @@
+Manage async job lifecycles through the ``DatabaseCheck`` job registry.

postgres/datadog_checks/postgres/data_observability.py

@@ -60,7 +60,7 @@ def __init__(self, check: PostgreSql, config: InstanceConfig):
         # Filter bad queries on check construction.
         self._queries, self._schedulers = self._filter_valid_queries(self._do_config.queries or ())
 
-    def _shutdown(self):
+    def shutdown(self) -> None:
         self._check = None
 
     @property

postgres/datadog_checks/postgres/metadata.py

@@ -137,7 +137,7 @@ def __init__(self, check: PostgreSql, config: InstanceConfig):
         self._tags_no_db = None
         self.tags = None
 
-    def _shutdown(self):
+    def shutdown(self) -> None:
         self._check = None
         self._schema_collector = None
         self._column_statistics_collector = None

postgres/datadog_checks/postgres/postgres.py

@@ -164,11 +164,13 @@ def __init__(self, name, init_config, instances):
             token_provider=self.build_token_provider(),
         )
         self.metrics_cache = PostgresMetricsCache(self._config)
-        # Initialize statement metrics collector before server version is known.
-        self.statement_metrics = PostgresStatementMetrics(self, self._config)
-        self.statement_samples = PostgresStatementSamples(self, self._config)
-        self.metadata_samples = PostgresMetadata(self, self._config)
-        self.data_observability = PostgresDataObservability(self, self._config)
+        # Only tests read these; the registry owns the jobs. Remove them once DatabaseCheck
+        # exposes a public job accessor. Jobs the configuration does not enable stay None.
+        self.statement_metrics = None
+        self.statement_samples = None
+        self.metadata_samples = None
+        self.data_observability = None
+        self._register_async_jobs()
         self._relations_manager = RelationsManager(self._config.relations, self._config.max_relations)
         self._clean_state()
         self._query_manager = QueryManager(self, lambda _: None, queries=[])  # query executor is set later
@@ -510,7 +512,7 @@ def cancel(self):
         cancel().
         """
         self.log.debug("Marking check as cancelled")
-        self._cancel_async_jobs()
+        self.cancel_async_jobs()
         needs_finalize = False
         with self._cancel_lock:
             self._cancelled = True
@@ -522,31 +524,22 @@ def cancel(self):
         else:
             self.log.debug("cancel() deferred finalize, check is still running")
 
-    @property
-    def _async_jobs(self):
-        """Return the async jobs active for this check's configuration."""
-        jobs = []
+    def _register_async_jobs(self):
+        """Build and register the async jobs enabled by this check's configuration."""
         if self._config.dbm:
-            jobs.extend([self.statement_metrics, self.statement_samples, self.metadata_samples])
-        elif self._config.data_observability.enabled:
-            jobs.append(self.metadata_samples)
+            # Built before the server version is known; _initialize_statement_metrics replaces it
+            # with the collector that suits the version.
+            self.statement_metrics = self.register_async_job(PostgresStatementMetrics(self, self._config))
+            self.statement_samples = self.register_async_job(PostgresStatementSamples(self, self._config))
+        if self._config.dbm or self._config.data_observability.enabled:
+            self.metadata_samples = self.register_async_job(PostgresMetadata(self, self._config))
         if self._config.data_observability.enabled:
-            jobs.append(self.data_observability)
-        return jobs
-
-    def _cancel_async_jobs(self):
-        """Signal async jobs to stop. Safe to call while check() is running."""
-        for job in self._async_jobs:
-            job.cancel()
+            self.data_observability = self.register_async_job(PostgresDataObservability(self, self._config))
 
     def _finalize(self):
         """Tear down check state. Must not run while check() is executing."""
         self.log.debug("Finalizing check: closing connections and clearing state")
-        for job in self._async_jobs:
-            if job._job_loop_future:
-                job._job_loop_future.result()
-                job._job_loop_future = None
-            job._shutdown()
+        self.shutdown_async_jobs()
         self._clean_state()
         self.check_initializations.clear()
         # TODO: move diagnosis cleanup into AgentCheck.cancel() in the base class
@@ -635,6 +628,8 @@ def load_version(self):
         self.set_metadata('version', self.raw_version)
 
     def _initialize_statement_metrics(self):
+        if not self._config.dbm:
+            return
         custom_pgss_view = self._config.pg_stat_statements_view != 'pg_stat_statements'
         if self._config.query_metrics.incremental_query_metrics and self.version < V10:
             self.log.warning(
@@ -652,11 +647,14 @@ def _initialize_statement_metrics(self):
 
         if self._config.query_metrics.incremental_query_metrics and self.version >= V10 and not custom_pgss_view:
             self.log.info("Using incremental query metrics collector")
-            self.statement_metrics = PostgresStatementMetricsV2(self, self._config)
+            collector = PostgresStatementMetricsV2(self, self._config)
         else:
             if not self._config.query_metrics.incremental_query_metrics:
                 self.log.info("Using legacy query metrics collector (full pg_stat_statements load)")
-            self.statement_metrics = PostgresStatementMetrics(self, self._config)
+            collector = PostgresStatementMetrics(self, self._config)
+        # Both collectors use the same job name, so registering replaces the instance built in
+        # _register_async_jobs.
+        self.statement_metrics = self.register_async_job(collector)
 
     def initialize_is_aurora(self):
         if self.is_aurora is None:
@@ -1242,14 +1240,7 @@ def check(self, _):
             if not self._config.only_custom_queries:
                 self._collect_stats(tags)
                 if not self._cancelled:
-                    if self._config.dbm:
-                        self.statement_metrics.run_job_loop(tags)
-                        self.statement_samples.run_job_loop(tags)
-                        self.metadata_samples.run_job_loop(tags)
-                    elif self._config.data_observability.enabled:
-                        self.metadata_samples.run_job_loop(tags)
-                    if self._config.data_observability.enabled:
-                        self.data_observability.run_job_loop(tags)
+                    self.run_async_jobs(tags)
                 if self._config.collect_wal_metrics is True:
                     # collect wal metrics for pg < 10 only when explicitly enabled
                     # (requires local filesystem access to the WAL directory)

postgres/datadog_checks/postgres/statement_samples.py

@@ -231,7 +231,7 @@ def __init__(self, check: PostgreSql, config: InstanceConfig):
         self._time_since_last_activity_event = 0
         self._pg_stat_activity_cols = None
 
-    def _shutdown(self):
+    def shutdown(self) -> None:
         self._check = None
         self._explain_parameterized_queries = None
         self._collection_strategy_cache = None

postgres/datadog_checks/postgres/statements.py

@@ -187,7 +187,7 @@ def __init__(self, check, config: InstanceConfig):
             ttl=60 * 60 / config.query_metrics.full_statement_text_samples_per_hour_per_query,
         )
 
-    def _shutdown(self):
+    def shutdown(self) -> None:
         self._check = None
         self._full_statement_text_cache = None
         self._state = None

postgres/datadog_checks/postgres/statements_v2.py

@@ -134,7 +134,7 @@ def __init__(self, check, config: InstanceConfig):
         )
         self._stat_column_cache: list[str] = []
 
-    def _shutdown(self):
+    def shutdown(self) -> None:
         self._check = None
         self._full_statement_text_cache = None
         self._delta_detector = None

postgres/pyproject.toml

@@ -28,7 +28,7 @@ classifiers = [
     "Private :: Do Not Upload",
 ]
 dependencies = [
-    "datadog-checks-base>=37.42.0",
+    "datadog-checks-base>=38.0.0",
 ]
 dynamic = [
     "version",

postgres/tests/test_unit.py

@@ -13,6 +13,8 @@
 
 from datadog_checks.postgres import PostgreSql, util
 from datadog_checks.postgres.schemas import PostgresSchemaCollector
+from datadog_checks.postgres.statements import PostgresStatementMetrics
+from datadog_checks.postgres.statements_v2 import PostgresStatementMetricsV2
 
 pytestmark = pytest.mark.unit
 
@@ -409,7 +411,7 @@ def test_check_gc_after_cancel(pg_instance):
     2. Find which attribute on that object points back to the check (usually
        ``self.check`` or ``self._check``).
     3. Null that attribute in ``cancel()`` or add it to the relevant
-       ``_shutdown()`` method.
+       ``shutdown()`` method.
     4. If the referrer is a closure or ``functools.partial``, find the
        registration site and null or clear the container that holds it.
     """
@@ -517,6 +519,59 @@ def test_run_after_cancel_returns_immediately(pg_instance):
     assert result == ''
 
 
+@pytest.mark.parametrize(
+    'dbm, data_observability_enabled, expected_jobs',
+    [
+        (False, False, []),
+        (True, False, ['query-metrics', 'query-samples', 'database-metadata']),
+        (False, True, ['database-metadata', 'data-observability']),
+        (True, True, ['query-metrics', 'query-samples', 'database-metadata', 'data-observability']),
+    ],
+)
+def test_async_job_registry_matches_config(pg_instance, dbm, data_observability_enabled, expected_jobs):
+    """Only the jobs enabled by the instance config are built and registered."""
+    pg_instance['dbm'] = dbm
+    pg_instance['data_observability'] = {'enabled': data_observability_enabled}
+
+    check = PostgreSql('postgres', {}, [pg_instance])
+
+    registered = check._async_job_registry
+    assert list(registered) == expected_jobs
+    # Each attribute holds the registered job, or None when the config does not enable it.
+    assert check.statement_metrics is registered.get('query-metrics')
+    assert check.statement_samples is registered.get('query-samples')
+    assert check.metadata_samples is registered.get('database-metadata')
+    assert check.data_observability is registered.get('data-observability')
+
+
+def test_initialize_statement_metrics_replaces_registered_job(pg_instance):
+    """The incremental collector replaces the placeholder registered before the version was known."""
+    pg_instance['dbm'] = True
+    pg_instance['query_metrics'] = {'enabled': True, 'incremental_query_metrics': True}
+
+    check = PostgreSql('postgres', {}, [pg_instance])
+    assert isinstance(check._async_job_registry['query-metrics'], PostgresStatementMetrics)
+
+    check.version = VersionInfo(14, 0, 0)
+    check._initialize_statement_metrics()
+
+    assert isinstance(check.statement_metrics, PostgresStatementMetricsV2)
+    assert check._async_job_registry['query-metrics'] is check.statement_metrics
+    assert list(check._async_job_registry) == ['query-metrics', 'query-samples', 'database-metadata']
+
+
+def test_initialize_statement_metrics_noop_without_dbm(pg_instance):
+    """Without DBM there is no query metrics job to build."""
+    pg_instance['dbm'] = False
+
+    check = PostgreSql('postgres', {}, [pg_instance])
+    check.version = VersionInfo(14, 0, 0)
+    check._initialize_statement_metrics()
+
+    assert check.statement_metrics is None
+    assert check._async_job_registry == {}
+
+
 def test_collect_column_statistics_updates_timestamp_on_failure(pg_instance):
     pg_instance['dbm'] = True
     pg_instance['collect_column_statistics'] = {'enabled': True, 'collection_interval': 60}