diff --git a/postgres/changelog.d/24852.fixed b/postgres/changelog.d/24852.fixed new file mode 100644 index 0000000000000..257952970e618 --- /dev/null +++ b/postgres/changelog.d/24852.fixed @@ -0,0 +1 @@ +Only tear the check down once when it is cancelled, and keep logging usable through the end of teardown. diff --git a/postgres/datadog_checks/postgres/postgres.py b/postgres/datadog_checks/postgres/postgres.py index 50528750971a8..a9f478a4c6f92 100644 --- a/postgres/datadog_checks/postgres/postgres.py +++ b/postgres/datadog_checks/postgres/postgres.py @@ -199,6 +199,7 @@ def __init__(self, name, init_config, instances): self._cancel_lock = threading.Lock() self._is_running = False self._cancelled = False + self._finalized = False def database_monitoring_column_statistics(self, raw_event: str): self.event_platform_event(raw_event, "dbm-column-statistics") @@ -537,19 +538,25 @@ def _register_async_jobs(self): 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.""" + """Tear down check state. Runs at most once, and never while check() is executing.""" + with self._cancel_lock: + if self._finalized: + return + self._finalized = True self.log.debug("Finalizing check: closing connections and clearing state") self.shutdown_async_jobs() self._clean_state() self.check_initializations.clear() # TODO: move diagnosis cleanup into AgentCheck.cancel() in the base class self._diagnosis = None - self.log.check = None self._query_manager = None self.health = None self._close_db() self._close_db_pool() 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 _clean_state(self): self.log.debug("Cleaning state") diff --git a/postgres/tests/test_unit.py b/postgres/tests/test_unit.py index 2046284065d8b..7a00721af7cdb 100644 --- a/postgres/tests/test_unit.py +++ b/postgres/tests/test_unit.py @@ -519,6 +519,21 @@ def test_run_after_cancel_returns_immediately(pg_instance): assert result == '' +def test_finalize_runs_once_across_repeated_cancels(pg_instance): + """Verify that teardown is idempotent.""" + check = PostgreSql('postgres', {}, [pg_instance]) + conn = mock.MagicMock() + check._db = conn + + with mock.patch.object(check.db_pool, 'close_all', wraps=check.db_pool.close_all) as close_all: + check.cancel() + check.cancel() + check._finalize() + + conn.close.assert_called_once() + close_all.assert_called_once() + + @pytest.mark.parametrize( 'dbm, data_observability_enabled, expected_jobs', [