[DBMON-6881] Move mysql onto the DatabaseCheck job registry and cancellation lifecycle

mysql/changelog.d/24936.fixed

@@ -0,0 +1 @@
+Manage the DBM async jobs through the ``DatabaseCheck`` registry.

mysql/datadog_checks/mysql/activity.py

@@ -247,6 +247,7 @@ def _check_version(self):
     def _collect_activity(self):
         # type: () -> None
         # do not emit any dd.internal metrics for DBM specific check code
+        self._raise_if_cancelled()
         tags = [t for t in self._tags if not t.startswith('dd.internal')]
         with closing(self._get_db_connection().cursor(CommenterDictCursor)) as cursor:
             rows = self._get_activity(cursor)
@@ -289,6 +290,7 @@ def _get_activity_query(self):
     @tracked_method(agent_check_getter=agent_check_getter, track_result_length=True)
     def _get_activity(self, cursor):
         # type: (pymysql.cursor) -> List[Dict[str]]
+        self._raise_if_cancelled()
         query = self._get_activity_query()
         self._log.debug("Running activity query [%s]", query)
         cursor.execute(query)
@@ -432,6 +434,12 @@ def _json_event_encoding(o):
             return int(o.total_seconds())
         raise TypeError
 
+    def shutdown(self) -> None:
+        self._close_db_conn()
+        self._check = None
+        # A bound method of the check, so it pins the check independently of _check above.
+        self._connection_args_provider = None
+
     def _close_db_conn(self):
         # type: () -> None
         if self._db:

mysql/datadog_checks/mysql/databases_data.py

@@ -144,11 +144,11 @@ def __init__(self, mysql_metadata, check, config):
             config.schemas_config.get('max_execution_time', self.DEFAULT_MAX_EXECUTION_TIME), collection_interval
         )
 
-    def shut_down(self):
-        self._data_submitter.submit()
-
     def _cursor_run(self, cursor, query, params=None):
         """Run the query, log it, and emit a metric on database error."""
+        cancel_event = getattr(self._metadata, '_cancel_event', None)
+        if cancel_event is not None and cancel_event.is_set():
+            raise Exception("Job loop cancelled. Aborting query.")
         try:
             params_repr = "({} params)".format(len(params)) if isinstance(params, list) else params
             self._log.debug("Running query [{}] params={}".format(query, params_repr))

mysql/datadog_checks/mysql/metadata.py

@@ -110,6 +110,15 @@ def get_db_connection(self):
             self._db.ping()
         return self._db
 
+    def shutdown(self) -> None:
+        self._close_db_conn()
+        self._check = None
+        # A bound method of the check, so it pins the check independently of _check above.
+        self._connection_args_provider = None
+        # DatabasesData points back at both this job and the check, so dropping it is what lets
+        # refcounting reclaim the pair.
+        self._databases_data = None
+
     def _close_db_conn(self):
         if self._db:
             try:
@@ -123,6 +132,7 @@ def _cursor_run(self, cursor, query, params=None):
         """
         Run and log the query. If provided, obfuscated params are logged in place of the regular params.
         """
+        self._raise_if_cancelled()
         try:
             self._log.debug("Running query [{}] params={}".format(query, params))
             cursor.execute(query, params)
@@ -158,9 +168,6 @@ def run_job(self):
                                 These may be unavailable until the error is resolved. The error - {}""".format(e)
                 )
 
-    def shut_down(self):
-        self._databases_data.shut_down()
-
     @tracked_method(agent_check_getter=attrgetter('_check'))
     def report_mysql_metadata(self):
         settings = []

mysql/datadog_checks/mysql/mysql.py

@@ -149,15 +149,12 @@ def __init__(self, name, init_config, instances):
             and self.cloud_metadata['aws']['managed_authentication'].get('enabled', False)
         )
 
-        # Pass function reference and managed auth flag to async jobs
-        self._statement_metrics = MySQLStatementMetrics(
-            self, self._config, self._get_connection_args, self._uses_aws_managed_auth
-        )
-        self._statement_samples = MySQLStatementSamples(
-            self, self._config, self._get_connection_args, self._uses_aws_managed_auth
-        )
-        self._mysql_metadata = MySQLMetadata(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
-        self._query_activity = MySQLActivity(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        # Jobs stay None unless the configuration enables DBM.
+        self._statement_metrics = None
+        self._statement_samples = None
+        self._mysql_metadata = None
+        self._query_activity = None
+        self._register_async_jobs()
         self._index_metrics = MySqlIndexMetrics(self._config)
         # _database_instance_emitted: limit the collection and transmission of the database instance metadata
         self._database_instance_emitted = TTLCache(
@@ -171,6 +168,33 @@ def __init__(self, name, init_config, instances):
 
         self._submit_initialization_health_event()
 
+    def shutdown(self) -> None:
+        """Release the resources this check holds for its whole lifetime."""
+        self._query_manager = None
+        self._runtime_queries_cached = None
+        self.health = None
+
+    def _register_async_jobs(self):
+        """Build and register the async jobs enabled by this check's configuration."""
+        # Every job requires DBM, and each job's own enabled flag defaults to true, so DBM is
+        # checked here rather than left to the jobs.
+        if not self._config.dbm_enabled:
+            return
+
+        # Pass function reference and managed auth flag to async jobs
+        self._statement_metrics = self.register_async_job(
+            MySQLStatementMetrics(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        )
+        self._statement_samples = self.register_async_job(
+            MySQLStatementSamples(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        )
+        self._mysql_metadata = self.register_async_job(
+            MySQLMetadata(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        )
+        self._query_activity = self.register_async_job(
+            MySQLActivity(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        )
+
     def _submit_initialization_health_event(self):
         try:
             # Handle the config validation result after we've set tags so those tags are included in the health event
@@ -398,10 +422,7 @@ def check(self, _):
 
                 if self._config.dbm_enabled:
                     dbm_tags = list(set(self.service_check_tags) | set(tags))
-                    self._statement_metrics.run_job_loop(dbm_tags)
-                    self._statement_samples.run_job_loop(dbm_tags)
-                    self._query_activity.run_job_loop(dbm_tags)
-                    self._mysql_metadata.run_job_loop(dbm_tags)
+                    self.run_async_jobs(dbm_tags)
 
                 # keeping track of these:
                 self._put_qcache_stats()
@@ -416,12 +437,6 @@ def check(self, _):
             self._conn = None
             self._report_warnings()
 
-    def cancel(self):
-        self._statement_samples.cancel()
-        self._statement_metrics.cancel()
-        self._query_activity.cancel()
-        self._mysql_metadata.cancel()
-
     def _new_query_executor(self, queries):
         return QueryExecutor(
             self.execute_query_raw,

mysql/datadog_checks/mysql/statement_samples.py

@@ -290,6 +290,12 @@ def _read_version_info(self):
                 self._global_status_table = "performance_schema.global_status"
             self._version_processed = True
 
+    def shutdown(self) -> None:
+        self._close_db_conn()
+        self._check = None
+        # A bound method of the check, so it pins the check independently of _check above.
+        self._connection_args_provider = None
+
     def _close_db_conn(self):
         if self._db:
             try:
@@ -330,6 +336,7 @@ def _cursor_run(self, cursor, query, params=None, obfuscated_params=None, obfusc
         """
         Run and log the query. If provided, obfuscated params are logged in place of the regular params.
         """
+        self._raise_if_cancelled()
         try:
             logged_query = obfuscated_query if obfuscated_query else query
             self._log.debug("Running query [%s] %s", logged_query, obfuscated_params if obfuscated_params else params)

mysql/datadog_checks/mysql/statements.py

@@ -147,6 +147,12 @@ def __init__(self, check, config, connection_args_provider, uses_managed_auth=Fa
             ttl=self._config.statement_rows_cache_ttl,
         )
 
+    def shutdown(self) -> None:
+        self._close_db_conn()
+        self._check = None
+        # A bound method of the check, so it pins the check independently of _check above.
+        self._connection_args_provider = None
+
     def _close_db_conn(self):
         if self._db:
             try:
@@ -236,6 +242,7 @@ def _collect_per_statement_metrics(self, tags):
         return rows
 
     def _get_statement_count(self, tags):
+        self._raise_if_cancelled()
         with closing(self._get_db_connection().cursor(CommenterDictCursor)) as cursor:
             cursor.execute("SELECT count(*) AS count from performance_schema.events_statements_summary_by_digest")
 
@@ -343,6 +350,7 @@ def _query_summary_per_statement(self):
                 LIMIT 10000
                 """
 
+        self._raise_if_cancelled()
         with closing(self._get_db_connection().cursor(CommenterDictCursor)) as cursor:
             args = [self._last_seen] if only_query_recent_statements else None
             cursor.execute(sql_statement_summary, args)

mysql/datadog_checks/mysql/util.py

@@ -100,6 +100,7 @@ class ManagedAuthConnectionMixin:
         self._uses_managed_auth (bool)
         self._db_created_at (float, timestamp)
         self._db (connection or None)
+        self._cancel_event (event.Event, used to abort queries if the Agent has unscheduled this check)
 
     Subclasses must implement:
         _close_db_conn() - closes self._db
@@ -113,6 +114,11 @@ def _should_reconnect_for_managed_auth(self):
             return False
         return (time.time() - self._db_created_at) >= self.MANAGED_AUTH_RECONNECT_INTERVAL
 
+    def _raise_if_cancelled(self):
+        """Abort before a query if the Agent has unscheduled this check."""
+        if self._cancel_event.is_set():
+            raise Exception("Job loop cancelled. Aborting query.")
+
     def _get_db_connection(self):
         """Get or create database connection, reconnecting periodically for managed auth."""
         if self._should_reconnect_for_managed_auth():

mysql/pyproject.toml

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

mysql/tests/test_metadata.py

@@ -3,11 +3,13 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
 import re
+from unittest import mock
 
 import pytest
 from packaging.version import parse as parse_version
 
 from datadog_checks.mysql import MySql
+from datadog_checks.mysql.databases_data import DatabasesData
 
 from . import common
 from .common import MYSQL_FLAVOR, MYSQL_REPLICATION, MYSQL_VERSION_PARSED
@@ -62,6 +64,19 @@ def normalize_values(actual_payload):
                             )
 
 
+@pytest.mark.unit
+def test_schema_collection_aborts_query_when_cancelled(dbm_instance):
+    """Schema collection is not a DBMAsyncJob; it must honor the metadata job's cancel event."""
+    check = MySql(common.CHECK_NAME, {}, instances=[dbm_instance])
+    check._mysql_metadata._cancel_event.set()
+    databases_data = DatabasesData(check._mysql_metadata, check, check._config)
+    cursor = mock.MagicMock()
+
+    with pytest.raises(Exception, match='cancelled'):
+        databases_data._cursor_run(cursor, 'SELECT 1')
+    cursor.execute.assert_not_called()
+
+
 @pytest.mark.integration
 @pytest.mark.usefixtures('dd_environment')
 def test_collect_mysql_settings(aggregator, dbm_instance, dd_run_check):

mysql/tests/test_query_activity.py

@@ -596,53 +596,6 @@ def test_activity_collection_rate_limit(aggregator, dd_run_check, dbm_instance):
     assert check._query_activity.collection_interval == collection_interval
 
 
-@pytest.mark.integration
-@pytest.mark.usefixtures('dd_environment')
-@pytest.mark.parametrize("activity_enabled", [True, False])
-def test_async_job_enabled(dd_run_check, dbm_instance, activity_enabled):
-    dbm_instance['query_activity'] = {'enabled': activity_enabled, 'run_sync': False}
-    check = MySql(CHECK_NAME, {}, [dbm_instance])
-    dd_run_check(check)
-    check.cancel()
-    if activity_enabled:
-        assert check._query_activity._job_loop_future is not None
-        check._query_activity._job_loop_future.result()
-    else:
-        assert check._query_activity._job_loop_future is None
-
-
-@pytest.mark.integration
-@pytest.mark.usefixtures('dd_environment')
-def test_async_job_inactive_stop(aggregator, dd_run_check, dbm_instance):
-    dbm_instance['query_activity']['run_sync'] = False
-    check = MySql(CHECK_NAME, {}, [dbm_instance])
-    dd_run_check(check)
-    check._query_activity._job_loop_future.result()
-    aggregator.assert_metric(
-        "dd.mysql.async_job.inactive_stop",
-        tags=_expected_dbm_job_err_tags(dbm_instance, check),
-        hostname='',
-    )
-
-
-@pytest.mark.integration
-@pytest.mark.usefixtures('dd_environment')
-def test_async_job_cancel(aggregator, dd_run_check, dbm_instance):
-    dbm_instance['query_activity']['run_sync'] = False
-    check = MySql(CHECK_NAME, {}, [dbm_instance])
-    dd_run_check(check)
-    check.cancel()
-    # wait for it to stop and make sure it doesn't throw any exceptions
-    check._query_activity._job_loop_future.result()
-    assert not check._query_activity._job_loop_future.running(), "activity thread should be stopped"
-    # if the thread doesn't start until after the cancel signal is set then the db connection will never
-    # be created in the first place
-    aggregator.assert_metric(
-        "dd.mysql.async_job.cancel",
-        tags=_expected_dbm_job_err_tags(dbm_instance, check),
-    )
-
-
 @pytest.mark.integration
 @pytest.mark.usefixtures('dd_environment')
 def test_events_wait_current_disabled(dbm_instance, dd_run_check, root_conn, aggregator):
@@ -729,24 +682,6 @@ def test_events_wait_current_disabled_no_warning_azure_flexible_server(
     assert not dbm_activity, "should not have collected any activity"
 
 
-# the inactive job metrics are emitted from the main integrations
-# directly to metrics-intake, so they should also be properly tagged with a resource
-def _expected_dbm_job_err_tags(dbm_instance, check):
-    _tags = dbm_instance['tags'] + (
-        'database_hostname:stubbed.hostname',
-        'database_instance:stubbed.hostname',
-        'job:query-activity',
-        'port:{}'.format(PORT),
-        'dd.internal.resource:database_instance:stubbed.hostname',
-        'dbms_flavor:{}'.format(MYSQL_FLAVOR.lower()),
-    )
-    if MYSQL_FLAVOR.lower() in ('mysql', 'percona'):
-        _tags += ("server_uuid:{}".format(check.server_uuid),)
-        if MYSQL_REPLICATION == 'classic':
-            _tags += ('cluster_uuid:{}'.format(check.cluster_uuid), 'replication_role:primary')
-    return _tags
-
-
 @pytest.mark.integration
 @pytest.mark.usefixtures('dd_environment')
 def test_if_deadlock_metric_is_collected(aggregator, dd_run_check, dbm_instance):