From cfc33053d93d19e1de247a9d38766d02939985c6 Mon Sep 17 00:00:00 2001
From: Lubos Mjachky <lmjachky@redhat.com>
Date: Thu, 3 Oct 2024 14:41:44 +0200
Subject: [PATCH] Collect artifact size metrics after caching the content

The caching machinery stores all headers from the response coming from
content-handler. With this commit, we are introducing an proprietary header
that contains the size of the served artifact or artifact available after the
redirect, X-PULP-ARTIFACT-SIZE.

closes #5817
---
 CHANGES/5817.bugfix         |  1 +
 pulpcore/cache/cache.py     |  6 ++++++
 pulpcore/cache/metrics.py   | 23 +++++++++++++++++++++++
 pulpcore/content/handler.py | 31 +++----------------------------
 4 files changed, 33 insertions(+), 28 deletions(-)
 create mode 100644 CHANGES/5817.bugfix
 create mode 100644 pulpcore/cache/metrics.py

diff --git a/CHANGES/5817.bugfix b/CHANGES/5817.bugfix
new file mode 100644
index 0000000000..d313546db7
--- /dev/null
+++ b/CHANGES/5817.bugfix
@@ -0,0 +1 @@
+Started collecting artifact size metrics even after caching the served content.
diff --git a/pulpcore/cache/cache.py b/pulpcore/cache/cache.py
index 0c4aa85eec..819a5d67e6 100644
--- a/pulpcore/cache/cache.py
+++ b/pulpcore/cache/cache.py
@@ -21,6 +21,8 @@
 )
 from pulpcore.responses import ArtifactResponse
 
+from .metrics import artifacts_size_counter
+
 DEFAULT_EXPIRES_TTL = settings.CACHE_SETTINGS["EXPIRES_TTL"]
 
 
@@ -352,6 +354,10 @@ async def cached_function(*args, **kwargs):
                 response = await self.make_entry(
                     key, bk, func, args, kwargs, self.default_expires_ttl
                 )
+
+            if size := response.headers.get("X-PULP-ARTIFACT-SIZE"):
+                artifacts_size_counter.add(size)
+
             return response
 
         return cached_function
diff --git a/pulpcore/cache/metrics.py b/pulpcore/cache/metrics.py
new file mode 100644
index 0000000000..caea4f0688
--- /dev/null
+++ b/pulpcore/cache/metrics.py
@@ -0,0 +1,23 @@
+from opentelemetry import metrics
+
+from pulpcore.app.util import MetricsEmitter, get_domain, get_worker_name
+
+
+class ArtifactsSizeCounter(MetricsEmitter):
+    def __init__(self):
+        self.meter = metrics.get_meter("artifacts.size.meter")
+        self.counter = self.meter.create_counter(
+            "artifacts.size.counter",
+            unit="Bytes",
+            description="Counts the size of served artifacts",
+        )
+
+    def add(self, amount):
+        attributes = {
+            "domain_name": get_domain().name,
+            "worker_process": get_worker_name(),
+        }
+        self.counter.add(int(amount), attributes)
+
+
+artifacts_size_counter = ArtifactsSizeCounter.build()
diff --git a/pulpcore/content/handler.py b/pulpcore/content/handler.py
index 7be860b5df..11b128f677 100644
--- a/pulpcore/content/handler.py
+++ b/pulpcore/content/handler.py
@@ -21,8 +21,6 @@
 
 import django
 
-from opentelemetry import metrics
-
 from pulpcore.constants import STORAGE_RESPONSE_MAP
 from pulpcore.responses import ArtifactResponse
 
@@ -52,9 +50,7 @@
 )
 from pulpcore.app import mime_types  # noqa: E402: module level not at top of file
 from pulpcore.app.util import (  # noqa: E402: module level not at top of file
-    MetricsEmitter,
     get_domain,
-    get_worker_name,
     cache_key,
 )
 
@@ -161,20 +157,6 @@ class Handler:
 
     distribution_model = None
 
-    class ArtifactsSizeCounter(MetricsEmitter):
-        def __init__(self):
-            self.meter = metrics.get_meter("artifacts.size.meter")
-            self.counter = self.meter.create_counter(
-                "artifacts.size.counter",
-                unit="Bytes",
-                description="Counts the size of served artifacts",
-            )
-
-        def add(self, amount, attributes):
-            self.counter.add(amount, attributes)
-
-    artifacts_size_counter = ArtifactsSizeCounter.build()
-
     @staticmethod
     def _reset_db_connection():
         """
@@ -1010,7 +992,7 @@ def _build_url(**kwargs):
             size = artifact_file.size or "*"
             raise HTTPRequestRangeNotSatisfiable(headers={"Content-Range": f"bytes */{size}"})
 
-        self._report_served_artifact_size(content_length)
+        headers["X-PULP-ARTIFACT-SIZE"] = str(content_length)
 
         if domain.storage_class == "pulpcore.app.models.storage.FileSystem":
             path = storage.path(artifact_name)
@@ -1144,9 +1126,9 @@ async def finalize():
         download_result = await downloader.run()
 
         if content_length := response.headers.get("Content-Length"):
-            self._report_served_artifact_size(int(content_length))
+            response.headers["X-PULP-ARTIFACT-SIZE"] = content_length
         else:
-            self._report_served_artifact_size(size)
+            response.headers["X-PULP-ARTIFACT-SIZE"] = str(size)
 
         if save_artifact and remote.policy != Remote.STREAMED:
             await asyncio.shield(
@@ -1157,10 +1139,3 @@ async def finalize():
         if response.status == 404:
             raise HTTPNotFound()
         return response
-
-    def _report_served_artifact_size(self, size):
-        attributes = {
-            "domain_name": get_domain().name,
-            "worker_name": get_worker_name(),
-        }
-        self.artifacts_size_counter.add(size, attributes)