Use statsd pipelines (#3504)

Avoid sending a udp packet per bit of information, instead use a
pipeline per request.  Add a MaybeStatsd utility class so we can skip
sending the stats when not in a flask app context.

Author: jcristau

src/auslib/db.py

@@ -16,7 +16,6 @@
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.sql.expression import null
 from sqlalchemy.sql.functions import max as sql_max
-from statsd.defaults.env import statsd
 
 from auslib.blobs.base import createBlob, merge_dicts
 from auslib.errors import PermissionDeniedError, ReadOnlyError, SignoffRequiredError
@@ -32,6 +31,7 @@
     matchSimpleExpression,
     matchVersion,
 )
+from auslib.util.statsd import statsd
 from auslib.util.timestamp import getMillisecondTimestamp
 from auslib.util.versions import get_version_class
 

src/auslib/util/auth.py

@@ -2,7 +2,8 @@
 import requests
 from auth0.authentication import Users as auth0_Users
 from repoze.lru import lru_cache
-from statsd.defaults.env import statsd
+
+from auslib.util.statsd import statsd
 
 
 class AuthError(Exception):
@@ -50,9 +51,9 @@ def get_jwks(auth_domain):
 
 
 @lru_cache(2048)
-@statsd.timer("auth0_userinfo")
 def get_additional_userinfo(auth_domain, access_token):
-    return auth0_Users(auth_domain).userinfo(access_token)
+    with statsd.timer("auth0_userinfo"):
+        return auth0_Users(auth_domain).userinfo(access_token)
 
 
 def verified_userinfo(request, auth_domain, auth_audience):

src/auslib/util/autograph.py

@@ -3,9 +3,9 @@
 
 import requests
 from requests_hawk import HawkAuth
-from statsd.defaults.env import statsd
 
 from auslib.util.retry import retry_sync
+from auslib.util.statsd import statsd
 
 SIGNATURE_PREFIX = "Content-Signature:\x00"
 

src/auslib/util/cache.py

@@ -1,7 +1,8 @@
 from copy import deepcopy
 
 from repoze.lru import ExpiringLRUCache
-from statsd.defaults.env import statsd
+
+from auslib.util.statsd import statsd
 
 uncached_sentinel = object()
 

src/auslib/util/statsd.py

@@ -0,0 +1,20 @@
+from contextlib import contextmanager
+
+from flask import g
+
+
+class MaybeStatsd:
+    @contextmanager
+    def timer(self, *args, **kwargs):
+        if g:
+            with g.statsd.timer(*args, **kwargs):
+                yield
+        else:
+            yield
+
+    def incr(self, *args, **kwargs):
+        if g:
+            return g.statsd.incr(*args, **kwargs)
+
+
+statsd = MaybeStatsd()

src/auslib/web/admin/base.py

@@ -55,7 +55,8 @@ def create_app():
     create_dockerflow_endpoints(flask_app)
 
     @flask_app.before_request
-    def setup_timer():
+    def setup_statsd():
+        g.statsd = statsd.pipeline()
         g.request_timer = None
         if should_time_request():
             # do some massaging to get the metric name right
@@ -70,7 +71,7 @@ def setup_timer():
                 .removeprefix("auslib_web_common_")
             )
             metric = f"endpoint_{metric}"
-            g.request_timer = statsd.timer(metric)
+            g.request_timer = g.statsd.timer(metric)
             g.request_timer.start()
 
     @flask_app.before_request
@@ -203,6 +204,7 @@ def add_security_headers(response):
     def send_stats(response):
         if hasattr(g, "request_timer") and g.request_timer:
             g.request_timer.stop()
+        g.statsd.send()
 
         return response
 

src/auslib/web/public/base.py

@@ -4,7 +4,7 @@
 from os import path
 
 import connexion
-from flask import Response, make_response, request, send_from_directory
+from flask import Response, g, make_response, request, send_from_directory
 from sentry_sdk import capture_exception
 from specsynthase.specbuilder import SpecBuilder
 from statsd.defaults.env import statsd
@@ -108,6 +108,15 @@ def robots():
     def contributejson():
         return send_from_directory(flask_app.static_folder, "contribute.json")
 
+    @flask_app.before_request
+    def create_statsd_pipeline():
+        g.statsd = statsd.pipeline()
+
+    @flask_app.after_request
+    def send_statsd_pipeline(response):
+        g.statsd.send()
+        return response
+
     @flask_app.before_request
     def set_cache_control():
         # By default, we want a cache that can be shared across requests from
@@ -141,7 +150,7 @@ def log_request(response):
         prefix = request.path.split("/")[1]
         if prefix not in ("update", "json", "api"):
             prefix = "unknown"
-        statsd.incr(f"response.{prefix}.{response.status_code}")
+        g.statsd.incr(f"response.{prefix}.{response.status_code}")
 
         return response
 

src/auslib/web/public/client.py

@@ -3,8 +3,7 @@
 import sys
 
 from flask import current_app as app
-from flask import make_response, request
-from statsd.defaults.env import statsd
+from flask import g, make_response, request
 
 from auslib.AUS import FORCE_FALLBACK_MAPPING, FORCE_MAIN_MAPPING
 from auslib.blobs.base import XMLBlob, createBlob
@@ -138,7 +137,7 @@ def extract_query_version(request_url):
 
 @with_transaction
 def get_update_blob(transaction, **url):
-    with statsd.timer("client.parse_query"):
+    with g.statsd.timer("client.parse_query"):
         url["queryVersion"] = extract_query_version(request.url)
         # Underlying code depends on osVersion being set. Since this route only
         # exists to support ancient queries, and all newer versions have osVersion
@@ -152,7 +151,7 @@ def get_update_blob(transaction, **url):
         query = getQueryFromURL(url)
         LOG.debug("Got query: %s", query)
 
-    with statsd.timer("client.evaluate_rules"):
+    with g.statsd.timer("client.evaluate_rules"):
         release, update_type, eval_metadata = AUS.evaluateRules(query, transaction=transaction)
 
     response_blobs = []
@@ -162,12 +161,12 @@ def get_update_blob(transaction, **url):
         response_products = release.getResponseProducts()
         response_blob_names = release.getResponseBlobs()
         if response_products:
-            with statsd.timer("client.process_response_products"):
+            with g.statsd.timer("client.process_response_products"):
                 # if we have a SuperBlob of gmp, we process the response products and
                 # concatenate their inner XMLs
                 response_blobs.extend(evaluate_response_products(response_products, query, transaction))
         elif response_blob_names:
-            with statsd.timer("client.process_response_blobs"):
+            with g.statsd.timer("client.process_response_blobs"):
                 # if we have a SuperBlob of systemaddons, we process the response products and
                 # concatenate their inner XMLs
                 response_blobs.extend(evaluate_response_blobs(response_blob_names, update_type, query, transaction))
@@ -179,7 +178,7 @@ def get_update_blob(transaction, **url):
                 squash_response = True
                 LOG.debug("Busted nightly detected, will squash xml response")
 
-    with statsd.timer("client.make_response"):
+    with g.statsd.timer("client.make_response"):
         return construct_response(release, query, update_type, response_blobs, squash_response, eval_metadata)
 
 

src/auslib/web/public/json.py

@@ -2,7 +2,7 @@
 
 from flask import Response
 from flask import current_app as app
-from statsd.defaults.env import statsd
+from flask import g
 
 from auslib.AUS import FORCE_FALLBACK_MAPPING, FORCE_MAIN_MAPPING
 from auslib.web.public.helpers import AUS, get_aus_metadata_headers, get_content_signature_headers, with_transaction
@@ -12,13 +12,13 @@
 def get_update(transaction, **parameters):
     force = parameters.get("force")
     parameters["force"] = {FORCE_MAIN_MAPPING.query_value: FORCE_MAIN_MAPPING, FORCE_FALLBACK_MAPPING.query_value: FORCE_FALLBACK_MAPPING}.get(force)
-    with statsd.timer("json.evaluate_rules"):
+    with g.statsd.timer("json.evaluate_rules"):
         release, _, eval_metadata = AUS.evaluateRules(parameters, transaction=transaction)
 
     if not release:
         return Response(status=404)
 
-    with statsd.timer("json.make_response"):
+    with g.statsd.timer("json.make_response"):
         headers = get_aus_metadata_headers(eval_metadata)
         response = json.dumps(release.getResponse(parameters, app.config["ALLOWLISTED_DOMAINS"]))
         headers.update(get_content_signature_headers(response, ""))

tests/admin/views/test_releases_v2.py

@@ -2195,8 +2195,9 @@ def test_statsd(api, method, endpoint, metric):
     Timings are handled centrally, so it's not worth the trouble to test
     every single one. Requests don't need to succeed for these to pass;
     they just need to be routable."""
-    with mock.patch("auslib.web.admin.base.statsd.timer") as mocked_timer:
+    with mock.patch("auslib.web.admin.base.statsd.pipeline") as mocked_pipeline:
         api.open(endpoint, method=method)
+        mocked_timer = mocked_pipeline().timer
         assert mocked_timer.call_count == 1
         mocked_timer.assert_has_calls([mock.call(metric)])
 
@@ -2209,9 +2210,8 @@ def test_statsd_gcs_creation(api, firefox_62_0_build1):
         # 20 locales to upload for + top level
         # a creation uploads twice
         # 21 * 2 = 42
-        # plus 1 for the overall request
-        assert mocked_timer.call_count == 43
-        mocked_timer.assert_has_calls([mock.call("async_gcs_upload"), mock.call("endpoint_releases_v2_ensure")], any_order=True)
+        assert mocked_timer.call_count == 42
+        assert mocked_timer.call_args_list == [mock.call("async_gcs_upload")] * 42
 
 
 @pytest.mark.usefixtures("releases_db", "mock_verified_userinfo")
@@ -2224,9 +2224,9 @@ def test_statsd_gcs_update(api):
         assert old_data_versions["platforms"]["Darwin_x86_64-gcc3-u-i386-x86_64"]["locales"]["de"]
         ret = api.post("/v2/releases/Firefox-60.0b3-build1", json={"blob": blob, "old_data_versions": old_data_versions})
         assert ret.status_code == 200, ret.data
-        # one call for top level modification, one for the changed locale, one for overall request
-        assert mocked_timer.call_count == 3
-        mocked_timer.assert_has_calls([mock.call("async_gcs_upload"), mock.call("endpoint_releases_v2_update")], any_order=True)
+        # one call for top level modification, one for the changed locale
+        assert mocked_timer.call_count == 2
+        assert mocked_timer.call_args_list == [mock.call("async_gcs_upload")] * 2
 
 
 @pytest.mark.usefixtures("releases_db", "mock_verified_userinfo")
@@ -2235,6 +2235,6 @@ def test_statsd_gcs_delete(api):
         ret = api.delete("/v2/releases/Firefox-65.0-build1")
         assert ret.status_code == 200, ret.data
 
-        # one call for top level, one call for each locale, one for overall request
-        assert mocked_timer.call_count == 22
-        mocked_timer.assert_has_calls([mock.call("async_gcs_upload"), mock.call("endpoint_releases_v2_delete")], any_order=True)
+        # one call for top level, one call for each locale
+        assert mocked_timer.call_count == 21
+        assert mocked_timer.call_args_list == [mock.call("async_gcs_upload")] * 21