feat(profiles): accept compressed data

Author: john-z-yang

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,29 @@ 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 | bytes) -> dict[str, Any]:
+    if isinstance(encoded, bytes):
+        return msgpack.unpackb(encoded, use_list=False)
+
+    # It's been b64encoded for taskworker
+
+    # compressed
+    if encoded.startswith("!"):
+        try:
+            res = msgpack.unpackb(
+                zlib.decompress(b64decode(encoded[1:].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)
 
 
@@ -125,10 +148,7 @@ def process_profile_task(
         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)
 
         profile = json.loads(message_dict["payload"], use_rapid_json=True)