diff options
Diffstat (limited to 'buildscripts/metrics/tooling_metrics_utils.py')
| -rw-r--r-- | buildscripts/metrics/tooling_metrics_utils.py | 63 |
1 files changed, 31 insertions, 32 deletions
diff --git a/buildscripts/metrics/tooling_metrics_utils.py b/buildscripts/metrics/tooling_metrics_utils.py index 29e4e475bc9..f668e457909 100644 --- a/buildscripts/metrics/tooling_metrics_utils.py +++ b/buildscripts/metrics/tooling_metrics_utils.py @@ -1,19 +1,17 @@ +import atexit import logging import os -from typing import Optional -from git import Repo +from typing import Any, Callable, Dict import pymongo -from buildscripts.metrics.metrics_datatypes import ToolingMetrics - -logger = logging.getLogger('tooling_metrics_utils') +logger = logging.getLogger('tooling_metrics') INTERNAL_TOOLING_METRICS_HOSTNAME = "mongodb+srv://dev-metrics-pl-0.kewhj.mongodb.net" INTERNAL_TOOLING_METRICS_USERNAME = "internal_tooling_user" INTERNAL_TOOLING_METRICS_PASSWORD = "internal_tooling_user" -def _get_internal_tooling_metrics_client(): +def _get_internal_tooling_metrics_client() -> pymongo.MongoClient: """Retrieve client for internal MongoDB tooling metrics cluster.""" return pymongo.MongoClient( host=INTERNAL_TOOLING_METRICS_HOSTNAME, @@ -27,25 +25,16 @@ def _get_internal_tooling_metrics_client(): ) -EXPECTED_TOOLCHAIN_LOCATION = "/opt/mongodbtoolchain" - - -def _toolchain_exists() -> bool: - """Check if the internal MongoDB toolchain exists.""" - return os.path.exists(EXPECTED_TOOLCHAIN_LOCATION) - - -def _git_user_exists() -> Optional[str]: - """Check if a git user email exists.""" - try: - return Repo('.').config_reader().get_value("user", "email", None) - except Exception: # pylint: disable=broad-except - return None +MONGOD_INTENRAL_DISTRO_FILEPATH = '/etc/mongodb-distro-name' def _is_virtual_workstation() -> bool: """Detect whether this is a MongoDB internal virtual workstation.""" - return _toolchain_exists() and _git_user_exists() + try: + with open(MONGOD_INTENRAL_DISTRO_FILEPATH, 'r') as file: + return file.read().strip() == 'ubuntu1804-workstation' + except Exception as _: # pylint: disable=broad-except + return False TOOLING_METRICS_OPT_OUT = "TOOLING_METRICS_OPT_OUT" @@ -56,22 +45,32 @@ def _has_metrics_opt_out() -> bool: return os.environ.get(TOOLING_METRICS_OPT_OUT, None) == '1' -def should_collect_metrics() -> bool: +def _should_collect_metrics() -> bool: """Determine whether to collect tooling metrics.""" return _is_virtual_workstation() and not _has_metrics_opt_out() -def _save_metrics(metrics: ToolingMetrics) -> None: - """Save tooling metrics data.""" - client = _get_internal_tooling_metrics_client() - client.metrics.tooling_metrics.insert_one(metrics.dict()) - - -def save_tooling_metrics(tooling_metrics: ToolingMetrics) -> None: - """Persist tooling metrics data to MongoDB Internal Atlas Cluster.""" +# DO NOT USE DIRECTLY -- This is only to be used when metrics collection is registered atexit +def _save_metrics( + generate_metrics_function: Callable, + generate_metrics_args: Dict[str, Any], +) -> None: + """Save metrics to the atlas cluster.""" try: - _save_metrics(tooling_metrics) + client = _get_internal_tooling_metrics_client() + metrics = generate_metrics_function(**generate_metrics_args) + client.metrics.tooling_metrics.insert_one(metrics.dict()) except Exception as exc: # pylint: disable=broad-except logger.warning( - "\n%s\n\nUnexpected: Tooling metrics collection is not available -- this is a non-issue.\nIf this message persists, feel free to reach out to #server-development-platform", + "%s\n\nInternal Metrics Collection Failed -- this is a non-issue.\nIf this message persists, feel free to reach out to #server-dev-platform", exc) + + +# This is the only util that should be used externally +def register_metrics_collection_atexit( + generate_metrics_function: Callable, + generate_metrics_args: Dict[str, Any], +) -> None: + """Register metrics collection on atexit.""" + if _should_collect_metrics(): + atexit.register(_save_metrics, generate_metrics_function, generate_metrics_args) |