feat(profiles): accept compressed data

Author: john-z-yang

src/sentry/profiles/consumers/process/factory.py

@@ -19,7 +19,7 @@ def process_message(message: Message[KafkaPayload]) -> None:
 
     if sampled or options.get("profiling.profile_metrics.unsampled_profiles.enabled"):
         b64encoded = b64encode(message.payload.value).decode("utf-8")
-        process_profile_task.delay(payload=b64encoded, sampled=sampled)
+        process_profile_task.delay(payload=b64encoded, sampled=sampled, compressed_profile=False)
 
         if random.random() < options.get("taskworker.try_compress.profile_metrics"):
             import time

src/sentry/profiles/task.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import io
+import logging
+import zlib
 from base64 import b64decode, b64encode
 from copy import deepcopy
 from datetime import datetime, timezone
@@ -86,8 +88,23 @@ def _get_profiles_producer_from_topic(topic: Topic) -> KafkaProducer:
     max_futures=settings.SENTRY_PROFILE_CHUNKS_FUTURES_MAX_LIMIT,
 )
 
+logger = logging.getLogger(__name__)
 
-def decode_payload(encoded: str) -> dict[str, Any]:
+
+def decode_payload(encoded: str, compressed_profile: bool) -> dict[str, Any]:
+    if compressed_profile:
+        try:
+            res = msgpack.unpackb(
+                zlib.decompress(b64decode(encoded.encode("utf-8"))), use_list=False
+            )
+            metrics.incr("profiling.profile_metrics.decompress", tags={"status": "ok"})
+            return res
+        except Exception as e:
+            logger.exception("Failed to decompress compressed profile", extra={"error": e})
+            metrics.incr("profiling.profile_metrics.decompress", tags={"status": "err"})
+            raise
+
+    # not compressed
     return msgpack.unpackb(b64decode(encoded.encode("utf-8")), use_list=False)
 
 
@@ -117,18 +134,16 @@ def encode_payload(message: dict[str, Any]) -> str:
 )
 def process_profile_task(
     profile: Profile | None = None,
-    payload: str | bytes | None = None,
+    payload: str | None = None,
     sampled: bool = True,
+    compressed_profile: bool = False,
     **kwargs: Any,
 ) -> None:
     if not sampled and not options.get("profiling.profile_metrics.unsampled_profiles.enabled"):
         return
 
     if payload:
-        if isinstance(payload, str):  # It's been b64encoded for taskworker
-            message_dict = decode_payload(payload)
-        else:
-            message_dict = msgpack.unpackb(payload, use_list=False)
+        message_dict = decode_payload(payload, compressed_profile)
 
         profile = json.loads(message_dict["payload"], use_rapid_json=True)
 

tests/sentry/profiles/consumers/test_process.py

@@ -54,7 +54,9 @@ def test_basic_profile_to_celery(self, process_profile_task):
         processing_strategy.terminate()
 
         process_profile_task.assert_called_with(
-            payload=b64encode(payload).decode("utf-8"), sampled=True
+            payload=b64encode(payload).decode("utf-8"),
+            sampled=True,
+            compressed_profile=False,
         )
 
 

tests/sentry/profiles/test_task.py

@@ -6,7 +6,6 @@
 from typing import Any
 from unittest.mock import patch
 
-import msgpack
 import pytest
 from django.core.files.uploadedfile import SimpleUploadedFile
 from django.urls import reverse
@@ -1033,19 +1032,11 @@ def test_unknown_sdk(
 @patch("sentry.profiles.task._push_profile_to_vroom")
 @patch("sentry.profiles.task._symbolicate_profile")
 @patch("sentry.models.projectsdk.get_sdk_index")
-@pytest.mark.parametrize(
-    ["should_encode"],
-    [
-        (True,),
-        (False,),
-    ],
-)
 @django_db_all
 def test_track_latest_sdk_with_payload(
     get_sdk_index: Any,
     _symbolicate_profile: Any,
     _push_profile_to_vroom: Any,
-    should_encode: bool,
     organization: Organization,
     project: Project,
     request: Any,
@@ -1065,11 +1056,7 @@ def test_track_latest_sdk_with_payload(
         "received": "2024-01-02T03:04:05",
         "payload": json.dumps(profile),
     }
-    payload: str | bytes
-    if should_encode:
-        payload = encode_payload(kafka_payload)
-    else:
-        payload = msgpack.packb(kafka_payload)
+    payload = encode_payload(kafka_payload)
 
     with Feature("organizations:profiling-sdks"):
         process_profile_task(payload=payload)