Add concurrency tracking (#55)

actualwitch · Brett Beutell · web-flow · commit 88164e7389a6 · 2023-06-23T17:56:58.000+02:00
* Add concurrency tracking

* Fix concurrency tests for otel and prom trackers

* Comment out opentelemetry concurrency test since gauges are not supported

* Update docs

* Add a test for passing objectives to autometrics then using async decorator

* Update version to 0.6

---------

Co-authored-by: Brett Beutell &lt;brett@fiberplane.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Exemplars support (#51)
+- Optional concurrency tracking support (#55)
 
 ### Changed
 
@@ -28,6 +29,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed decorator async function handling (#55)
+
 ### Security
 
 - Update requests, starlette, fastapi dependencies used by the examples
diff --git a/README.md b/README.md
@@ -45,6 +45,8 @@ def sayHello:
 
 ```
 
+- You can also track the number of concurrent calls to a function by using the `track_concurrency` argument: `@autometrics(track_concurrency=True)`. Note: currently only supported by the `prometheus` tracker.
+
 - To access the PromQL queries for your decorated functions, run `help(yourfunction)` or `print(yourfunction.__doc__)`.
 
 - To show tooltips over decorated functions in VSCode, with links to Prometheus queries, try installing [the VSCode extension](https://marketplace.visualstudio.com/items?itemName=Fiberplane.autometrics).
diff --git a/examples/django_example/django_example/urls.py b/examples/django_example/django_example/urls.py
@@ -16,12 +16,14 @@
 """
 from django.urls import path
 
+from .views.concurrency import ConcurrencyView
 from .views.latency import RandomLatencyView
 from .views.metrics import metrics
 from .views.simple import simple_handler
 from .views.error import ErrorOrOkView
 
 urlpatterns = [
+    path("concurrency/", ConcurrencyView.as_view()),
     path("latency/", RandomLatencyView.as_view()),
     path("error/", ErrorOrOkView.as_view()),
     path("simple/", simple_handler),
diff --git a/examples/django_example/django_example/views/concurrency.py b/examples/django_example/django_example/views/concurrency.py
@@ -0,0 +1,15 @@
+import time
+from autometrics import autometrics
+from django.http import HttpResponse
+from django.views import View
+
+
+class ConcurrencyView(View):
+    """Here you can see how concurrency tracking works in autometrics.
+    Just add the `track_concurrency=True` argument, and autometrics
+    will track the number of concurrent requests to this endpoint."""
+
+    @autometrics(track_concurrency=True)
+    def get(self, request):
+        time.sleep(0.25)
+        return HttpResponse("Many clients wait for a reply from this endpoint!")
diff --git a/examples/django_example/django_example/views/latency.py b/examples/django_example/django_example/views/latency.py
@@ -10,8 +10,8 @@ class RandomLatencyView(View):
 
     @autometrics
     def get(self, request):
-        duration = random.randint(1, 500)
+        duration = random.randint(1, 10)
 
-        time.sleep(duration / 1000)
+        time.sleep(duration / 10)
 
         return HttpResponse("i was waiting for {}ms!".format(duration))
diff --git a/examples/django_example/locustfile.py b/examples/django_example/locustfile.py
@@ -0,0 +1,22 @@
+import time
+from locust import HttpUser, task, between
+
+
+class DjangoUser(HttpUser):
+    wait_time = between(1, 2.5)
+
+    @task(10)
+    def visit_concurrency_handler(self):
+        self.client.get("/concurrency/")
+
+    @task
+    def visit_error_handler(self):
+        self.client.get("/error/")
+
+    @task
+    def visit_simple_handler(self):
+        self.client.get("/simple/")
+
+    @task
+    def visit_latency_handler(self):
+        self.client.get("/latency/")
diff --git a/examples/django_example/run_example.sh b/examples/django_example/run_example.sh
@@ -0,0 +1,15 @@
+#!/bin/sh
+
+export AUTOMETRICS_COMMIT=67a1b3a
+export AUTOMETRICS_VERSION=0.1.0
+export AUTOMETRICS_BRANCH=main
+export AUTOMETRICS_TRACKER=prometheus
+
+# run the server itself
+poetry run python manage.py runserver 8080 &
+# run the locust load test and pipe stdout to dev/null
+poetry run locust --host=http://localhost:8080 --users=100 --headless &
+
+# kill all child processes on exit
+trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT
+wait
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "autometrics"
-version = "0.5"
+version = "0.6"
 description = "Easily add metrics to your system – and actually understand them using automatically customized Prometheus queries"
 authors = ["Fiberplane <info@fiberplane.com>"]
 license = "MIT OR Apache-2.0"
@@ -72,6 +72,7 @@ urllib3 = "1.26.15"
 uvicorn = "0.21.1"
 webencodings = "0.5.1"
 zipp = "3.15.0"
+locust = "^2.15.1"
 
 [build-system]
 requires = ["poetry-core"]
diff --git a/src/autometrics/constants.py b/src/autometrics/constants.py
@@ -2,14 +2,18 @@
 
 COUNTER_NAME = "function.calls.count"
 HISTOGRAM_NAME = "function.calls.duration"
+CONCURRENCY_NAME = "function.calls.concurrent"
 # NOTE - The Rust implementation does not use `build.info`, instead opts for just `build_info`
 BUILD_INFO_NAME = "build_info"
 
+
 COUNTER_NAME_PROMETHEUS = COUNTER_NAME.replace(".", "_")
 HISTOGRAM_NAME_PROMETHEUS = HISTOGRAM_NAME.replace(".", "_")
+CONCURRENCY_NAME_PROMETHEUS = CONCURRENCY_NAME.replace(".", "_")
 
 COUNTER_DESCRIPTION = "Autometrics counter for tracking function calls"
 HISTOGRAM_DESCRIPTION = "Autometrics histogram for tracking function call duration"
+CONCURRENCY_DESCRIPTION = "Autometrics gauge for tracking function call concurrency"
 BUILD_INFO_DESCRIPTION = (
     "Autometrics info metric for tracking software version and build details"
 )
diff --git a/src/autometrics/decorator.py b/src/autometrics/decorator.py
@@ -28,24 +28,33 @@ def autometrics(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
 
 # Decorator with arguments
 @overload
-def autometrics(*, objective: Optional[Objective] = None) -> Callable:
+def autometrics(
+    *, objective: Optional[Objective] = None, track_concurrency: Optional[bool] = False
+) -> Callable:
     ...
 
 
 def autometrics(
     func: Optional[Callable] = None,
     *,
     objective: Optional[Objective] = None,
+    track_concurrency: Optional[bool] = False,
 ):
     """Decorator for tracking function calls and duration. Supports synchronous and async functions."""
 
+    def track_start(function: str, module: str):
+        get_tracker().start(
+            function=function, module=module, track_concurrency=track_concurrency
+        )
+
     def track_result_ok(start_time: float, function: str, module: str, caller: str):
         get_tracker().finish(
             start_time,
             function=function,
             module=module,
             caller=caller,
             objective=objective,
+            track_concurrency=track_concurrency,
             result=Result.OK,
         )
 
@@ -61,6 +70,7 @@ def track_result_error(
             module=module,
             caller=caller,
             objective=objective,
+            track_concurrency=track_concurrency,
             result=Result.ERROR,
         )
 
@@ -77,6 +87,8 @@ def sync_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
             caller = get_caller_function()
 
             try:
+                if track_concurrency:
+                    track_start(module=module_name, function=func_name)
                 result = func(*args, **kwds)
                 track_result_ok(
                     start_time, function=func_name, module=module_name, caller=caller
@@ -110,6 +122,8 @@ async def async_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
             caller = get_caller_function()
 
             try:
+                if track_concurrency:
+                    track_start(module=module_name, function=func_name)
                 result = await func(*args, **kwds)
                 track_result_ok(
                     start_time, function=func_name, module=module_name, caller=caller
@@ -130,8 +144,14 @@ async def async_wrapper(*args: P.args, **kwds: P.kwargs) -> T:
         async_wrapper.__doc__ = append_docs_to_docstring(func, func_name, module_name)
         return async_wrapper
 
+    def pick_decorator(func: Callable) -> Callable:
+        """Pick the correct decorator based on the function type."""
+        if inspect.iscoroutinefunction(func):
+            return async_decorator(func)
+        return sync_decorator(func)
+
     if func is None:
-        return sync_decorator
+        return pick_decorator
     elif inspect.iscoroutinefunction(func):
         return async_decorator(func)
     else:
diff --git a/src/autometrics/test_decorator.py b/src/autometrics/test_decorator.py
@@ -148,6 +148,50 @@ def test_objectives(self):
         duration_sum = f"""function_calls_duration_sum{{function="{function_name}",module="test_decorator",objective_latency_threshold="{latency[0].value}",objective_name="{objective_name}",objective_percentile="{latency[1].value}"}}"""
         assert duration_sum in data
 
+    @pytest.mark.asyncio
+    async def test_objectives_async(self):
+        """This is a test that covers objectives for async functions."""
+
+        # set up the function + objective variables
+        caller = get_caller_function(depth=1)
+        assert caller is not None
+        assert caller != ""
+        objective_name = "test_objective"
+        success_rate = ObjectivePercentile.P90
+        latency = (ObjectiveLatency.Ms100, ObjectivePercentile.P99)
+        objective = Objective(
+            name=objective_name, success_rate=success_rate, latency=latency
+        )
+        function_name = basic_async_function.__name__
+        wrapped_function = autometrics(objective=objective)(basic_async_function)
+
+        sleep_duration = 0.25
+
+        # Test that the function is *still* async after we wrapped it
+        assert asyncio.iscoroutinefunction(wrapped_function) == True
+
+        await wrapped_function(sleep_duration)
+
+        # get the metrics
+        blob = generate_latest()
+        assert blob is not None
+        data = blob.decode("utf-8")
+
+        total_count = f"""function_calls_count_total{{caller="{caller}",function="{function_name}",module="test_decorator",objective_name="{objective_name}",objective_percentile="{success_rate.value}",result="ok"}} 1.0"""
+        assert total_count in data
+
+        # Check the latency buckets
+        for objective in ObjectiveLatency:
+            count = 0 if float(objective.value) <= sleep_duration else 1
+            query = f"""function_calls_duration_bucket{{function="{function_name}",le="{objective.value}",module="test_decorator",objective_latency_threshold="{latency[0].value}",objective_name="{objective_name}",objective_percentile="{latency[1].value}"}} {count}"""
+            assert query in data
+
+        duration_count = f"""function_calls_duration_count{{function="{function_name}",module="test_decorator",objective_latency_threshold="{latency[0].value}",objective_name="{objective_name}",objective_percentile="{latency[1].value}"}}"""
+        assert duration_count in data
+
+        duration_sum = f"""function_calls_duration_sum{{function="{function_name}",module="test_decorator",objective_latency_threshold="{latency[0].value}",objective_name="{objective_name}",objective_percentile="{latency[1].value}"}}"""
+        assert duration_sum in data
+
     def test_exception(self):
         """This is a test that covers exceptions."""
         caller = get_caller_function(depth=1)
diff --git a/src/autometrics/tracker/opentelemetry.py b/src/autometrics/tracker/opentelemetry.py
@@ -1,4 +1,3 @@
-import os
 import time
 from typing import Optional
 
@@ -17,6 +16,8 @@
 from .tracker import Result
 from ..objectives import Objective, ObjectiveLatency
 from ..constants import (
+    CONCURRENCY_NAME,
+    CONCURRENCY_DESCRIPTION,
     COUNTER_DESCRIPTION,
     COUNTER_NAME,
     HISTOGRAM_DESCRIPTION,
@@ -41,7 +42,8 @@ class OpenTelemetryTracker:
 
     __counter_instance: Counter
     __histogram_instance: Histogram
-    __up_down_counter_instance: UpDownCounter
+    __up_down_counter_build_info_instance: UpDownCounter
+    __up_down_counter_concurrency_instance: UpDownCounter
 
     def __init__(self):
         exporter = PrometheusMetricReader("")
@@ -63,10 +65,14 @@ def __init__(self):
             name=HISTOGRAM_NAME,
             description=HISTOGRAM_DESCRIPTION,
         )
-        self.__up_down_counter_instance = meter.create_up_down_counter(
+        self.__up_down_counter_build_info_instance = meter.create_up_down_counter(
             name=BUILD_INFO_NAME,
             description=BUILD_INFO_DESCRIPTION,
         )
+        self.__up_down_counter_concurrency_instance = meter.create_up_down_counter(
+            name=CONCURRENCY_NAME,
+            description=CONCURRENCY_DESCRIPTION,
+        )
         self._has_set_build_info = False
 
     def __count(
@@ -129,7 +135,7 @@ def __histogram(
     def set_build_info(self, commit: str, version: str, branch: str):
         if not self._has_set_build_info:
             self._has_set_build_info = True
-            self.__up_down_counter_instance.add(
+            self.__up_down_counter_build_info_instance.add(
                 1.0,
                 attributes={
                     "commit": commit,
@@ -138,6 +144,19 @@ def set_build_info(self, commit: str, version: str, branch: str):
                 },
             )
 
+    def start(
+        self, function: str, module: str, track_concurrency: Optional[bool] = False
+    ):
+        """Start tracking metrics for a function call."""
+        if track_concurrency:
+            self.__up_down_counter_concurrency_instance.add(
+                1.0,
+                attributes={
+                    "function": function,
+                    "module": module,
+                },
+            )
+
     def finish(
         self,
         start_time: float,
@@ -146,12 +165,22 @@ def finish(
         caller: str,
         result: Result = Result.OK,
         objective: Optional[Objective] = None,
+        track_concurrency: Optional[bool] = False,
     ):
         """Finish tracking metrics for a function call."""
+
         exemplar = None
         # Currently, exemplars are only supported by prometheus-client
         # https://github.com/autometrics-dev/autometrics-py/issues/41
         # if os.getenv("AUTOMETRICS_EXEMPLARS") == "true":
         #     exemplar = get_exemplar()
         self.__count(function, module, caller, objective, exemplar, result)
         self.__histogram(function, module, start_time, objective, exemplar)
+        if track_concurrency:
+            self.__up_down_counter_concurrency_instance.add(
+                -1.0,
+                attributes={
+                    "function": function,
+                    "module": module,
+                },
+            )
diff --git a/src/autometrics/tracker/prometheus.py b/src/autometrics/tracker/prometheus.py
diff --git a/src/autometrics/tracker/test_concurrency.py b/src/autometrics/tracker/test_concurrency.py
diff --git a/src/autometrics/tracker/tracker.py b/src/autometrics/tracker/tracker.py
