diff --git a/src/common/core/constants.py b/src/common/core/constants.py index d4d74ea5..b55a6ca0 100644 --- a/src/common/core/constants.py +++ b/src/common/core/constants.py @@ -1 +1 @@ -DEFAULT_PROMETHEUS_MULTIPROC_DIR = "/tmp/flagsmith-prometheus" +DEFAULT_PROMETHEUS_MULTIPROC_DIR_NAME = "flagsmith-prometheus" diff --git a/src/common/core/main.py b/src/common/core/main.py index 47cd10ea..62774fd3 100644 --- a/src/common/core/main.py +++ b/src/common/core/main.py @@ -1,7 +1,6 @@ import contextlib import logging import os -import shutil import sys import typing @@ -10,7 +9,7 @@ ) from common.core.cli import healthcheck -from common.core.constants import DEFAULT_PROMETHEUS_MULTIPROC_DIR +from common.prometheus.multiprocessing import prepare_prom_multiproc_dir logger = logging.getLogger(__name__) @@ -35,6 +34,9 @@ def ensure_cli_env() -> typing.Generator[None, None, None]: # TODO @khvn26 Move logging setup to here + # Prometheus multiproc support + prepare_prom_multiproc_dir() + # Currently we don't install Flagsmith modules as a package, so we need to add # $CWD to the Python path to be able to import them sys.path.append(os.getcwd()) @@ -43,18 +45,6 @@ def ensure_cli_env() -> typing.Generator[None, None, None]: # without resorting to it being set outside of the application os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings.dev") - # Set up Prometheus' multiprocess mode - prometheus_multiproc_dir_name = os.environ.setdefault( - "PROMETHEUS_MULTIPROC_DIR", - DEFAULT_PROMETHEUS_MULTIPROC_DIR, - ) - shutil.rmtree(prometheus_multiproc_dir_name, ignore_errors=True) - os.makedirs(prometheus_multiproc_dir_name, exist_ok=True) - logger.info( - "Re-created %s for Prometheus multi-process mode", - prometheus_multiproc_dir_name, - ) - if "docgen" in sys.argv: os.environ["DOCGEN_MODE"] = "true" diff --git a/src/common/prometheus/__init__.py b/src/common/prometheus/__init__.py index c0550324..fc52dd36 100644 --- a/src/common/prometheus/__init__.py +++ b/src/common/prometheus/__init__.py @@ -1,3 +1,16 @@ -from common.prometheus.utils import Histogram +from typing import Any -__all__ = ("Histogram",) +_utils = ("Histogram",) + + +def __getattr__(name: str) -> Any: + """ + Since utils imports django.conf.settings, we lazy load any objects that + we want to import to prevent django.core.exceptions.ImproperlyConfigured + due to settings not being configured. + """ + if name in _utils: + from common.prometheus import utils + + return getattr(utils, name) + raise AttributeError(name) diff --git a/src/common/prometheus/multiprocessing.py b/src/common/prometheus/multiprocessing.py new file mode 100644 index 00000000..5fd1c6ad --- /dev/null +++ b/src/common/prometheus/multiprocessing.py @@ -0,0 +1,36 @@ +import os +import pathlib +import shutil +from tempfile import gettempdir + +from common.core.constants import DEFAULT_PROMETHEUS_MULTIPROC_DIR_NAME + + +def prepare_prom_multiproc_dir() -> None: + prom_dir = pathlib.Path( + os.environ.setdefault( + "PROMETHEUS_MULTIPROC_DIR", + os.path.join(gettempdir(), DEFAULT_PROMETHEUS_MULTIPROC_DIR_NAME), + ) + ) + + if prom_dir.exists(): + for p in prom_dir.rglob("*"): + try: + # Ensure that the cleanup doesn't silently fail on + # files and subdirs created by other users. + p.chmod(0o777) + except Exception: # pragma: no cover + pass + + shutil.rmtree(prom_dir, ignore_errors=True) + + prom_dir.mkdir(parents=True, exist_ok=True) + + # While `mkdir` sets mode=0o777 by default, this can be affected by umask resulting in + # lesser permissions for other users. This step ensures the directory is writable for + # all users. + try: + prom_dir.chmod(0o777) + except PermissionError: + pass diff --git a/src/task_processor/processor.py b/src/task_processor/processor.py index daab3aee..159b20ca 100644 --- a/src/task_processor/processor.py +++ b/src/task_processor/processor.py @@ -194,7 +194,7 @@ def _run_task( "result": result.lower(), } - timer.labels(**labels) # type: ignore[no-untyped-call] + timer.labels(**labels) ctx.close() metrics.flagsmith_task_processor_finished_tasks_total.labels(**labels).inc()