feat(taskworker): Add compression in taskworker producer and worker (#95153)

The taskbroker system does not handle large payloads well. We first
experienced the effects of this with ingest-profiles when rolling out to
smaller environments. We also observed slow throughput and sqlite issues
for processing pools handling sentryapp and seer tasks (2-8MB payloads).
This PR is responsible for adding the ability to enable task parameter
compression in taskworkers.

Flow:
1. User defines a `CompressionType` attribute on the
`@instrumented_task` decorator which chooses the compression algorithm
(only supports ZSTD and PLAINTEXT). Defaults to PLAINTEXT.
2. In the taskworker producer layer, parameters gets compression and
serialized. A task header is added indicating the compression type
3. Parameters stay compressed in kafka and taskbroker storage.
4. Using the tasks header, the worker determines whether the task needs
to be decompressed. If so, base64 decode and decompress the message.

Rollout:
To rollout this change, we update a task's decorator with
`CompressionType`, then update the sentry option's compression rollout
rate. Depending on the rate, tasks will be sampled for compression. By
using the task header, this enables same tasks to be incrementally
rolled out.

Testing:
- Local testing 
- Unit tests
- Will start with 1% of a single task in S4S. Observe compression
metrics (duration and size) to ensure similar performance.

Author: enochtangg

src/sentry/options/defaults.py

@@ -3452,6 +3452,14 @@
     flags=FLAG_AUTOMATOR_MODIFIABLE,
 )
 
+# Taskbroker compression flag
+register(
+    "taskworker.enable_compression.rollout",
+    default=0.0,
+    type=Float,
+    flags=FLAG_AUTOMATOR_MODIFIABLE,
+)
+
 # Orgs for which compression should be disabled in the chunk upload endpoint.
 # This is intended to circumvent sporadic 503 errors reported by some customers.
 register("chunk-upload.no-compression", default=[], flags=FLAG_AUTOMATOR_MODIFIABLE)

src/sentry/taskworker/constants.py

@@ -1,3 +1,5 @@
+from enum import Enum
+
 DEFAULT_PROCESSING_DEADLINE = 10
 """
 The fallback/default processing_deadline that tasks
@@ -32,3 +34,12 @@
 The number of tasks a worker child process will process
 before being restarted.
 """
+
+
+class CompressionType(Enum):
+    """
+    The type of compression used for task parameters.
+    """
+
+    ZSTD = "zstd"
+    PLAINTEXT = "plaintext"

src/sentry/taskworker/registry.py

@@ -15,7 +15,7 @@
 from sentry_sdk.consts import OP, SPANDATA
 
 from sentry.conf.types.kafka_definition import Topic
-from sentry.taskworker.constants import DEFAULT_PROCESSING_DEADLINE
+from sentry.taskworker.constants import DEFAULT_PROCESSING_DEADLINE, CompressionType
 from sentry.taskworker.retry import Retry
 from sentry.taskworker.router import TaskRouter
 from sentry.taskworker.task import P, R, Task
@@ -83,6 +83,7 @@ def register(
         processing_deadline_duration: int | datetime.timedelta | None = None,
         at_most_once: bool = False,
         wait_for_delivery: bool = False,
+        compression_type: CompressionType = CompressionType.PLAINTEXT,
     ) -> Callable[[Callable[P, R]], Task[P, R]]:
         """
         Register a task.
@@ -108,6 +109,8 @@ def register(
         wait_for_delivery: bool
             If true, the task will wait for the delivery report to be received
             before returning.
+        compression_type: CompressionType
+            The compression type to use to compress the task parameters.
         """
 
         def wrapped(func: Callable[P, R]) -> Task[P, R]:
@@ -125,6 +128,7 @@ def wrapped(func: Callable[P, R]) -> Task[P, R]:
                 ),
                 at_most_once=at_most_once,
                 wait_for_delivery=wait_for_delivery,
+                compression_type=compression_type,
             )
             # TODO(taskworker) tasks should be registered into the registry
             # so that we can ensure task names are globally unique

src/sentry/taskworker/task.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
+import base64
 import datetime
+import random
+import time
 from collections.abc import Callable, Collection, Mapping, MutableMapping
 from functools import update_wrapper
 from typing import TYPE_CHECKING, Any, Generic, ParamSpec, TypeVar
 from uuid import uuid4
 
 import orjson
 import sentry_sdk
+import zstandard as zstd
 from django.conf import settings
 from django.utils import timezone
 from google.protobuf.timestamp_pb2 import Timestamp
@@ -17,8 +21,10 @@
     TaskActivation,
 )
 
-from sentry.taskworker.constants import DEFAULT_PROCESSING_DEADLINE
+from sentry import options
+from sentry.taskworker.constants import DEFAULT_PROCESSING_DEADLINE, CompressionType
 from sentry.taskworker.retry import Retry
+from sentry.utils import metrics
 
 if TYPE_CHECKING:
     from sentry.taskworker.registry import TaskNamespace
@@ -39,6 +45,7 @@ def __init__(
         processing_deadline_duration: int | datetime.timedelta | None = None,
         at_most_once: bool = False,
         wait_for_delivery: bool = False,
+        compression_type: CompressionType = CompressionType.PLAINTEXT,
     ):
         self.name = name
         self._func = func
@@ -58,6 +65,7 @@ def __init__(
         self._retry = retry
         self.at_most_once = at_most_once
         self.wait_for_delivery = wait_for_delivery
+        self.compression_type = compression_type
         update_wrapper(self, func)
 
     @property
@@ -169,12 +177,41 @@ def create_activation(
                     f"The `{key}` header value is of type {type(value)}"
                 )
 
+        parameters_json = orjson.dumps({"args": args, "kwargs": kwargs})
+        if self.compression_type == CompressionType.ZSTD:
+            # TODO(taskworker): Nesting this conditional avoids django_db fixtures in tests.
+            # Once we have rolled out compression safely, we can remove this conditional.
+            compression_rollout_rate = options.get("taskworker.enable_compression.rollout")
+            if compression_rollout_rate and compression_rollout_rate > random.random():
+                # Worker uses this header to determine if the parameters are decompressed
+                headers["compression-type"] = CompressionType.ZSTD.value
+                start_time = time.perf_counter()
+                parameters_data = zstd.compress(parameters_json)
+                # Compressed data is binary and needs base64 encoding for transport
+                parameters_str = base64.b64encode(parameters_data).decode("utf8")
+                end_time = time.perf_counter()
+
+                metrics.distribution(
+                    "taskworker.producer.compressed_parameters_size",
+                    len(parameters_str),
+                    tags={"namespace": self._namespace.name, "taskname": self.name},
+                )
+                metrics.distribution(
+                    "taskworker.producer.compression_time",
+                    end_time - start_time,
+                    tags={"namespace": self._namespace.name, "taskname": self.name},
+                )
+            else:
+                parameters_str = parameters_json.decode("utf8")
+        else:
+            parameters_str = parameters_json.decode("utf8")
+
         return TaskActivation(
             id=uuid4().hex,
             namespace=self._namespace.name,
             taskname=self.name,
             headers=headers,
-            parameters=orjson.dumps({"args": args, "kwargs": kwargs}).decode("utf8"),
+            parameters=parameters_str,
             retry_state=self._create_retry_state(),
             received_at=received_at,
             processing_deadline_duration=processing_deadline,

src/sentry/taskworker/tasks/examples.py

@@ -4,6 +4,7 @@
 from time import sleep
 from typing import Any
 
+from sentry.taskworker.constants import CompressionType
 from sentry.taskworker.namespaces import exampletasks
 from sentry.taskworker.retry import LastAction, NoRetriesRemainingError, Retry, RetryError
 from sentry.taskworker.retry import retry_task as retry_task_helper
@@ -87,3 +88,9 @@ def at_most_once_task() -> None:
 def timed_task(sleep_seconds: float | str, *args: list[Any], **kwargs: dict[str, Any]) -> None:
     sleep(float(sleep_seconds))
     logger.debug("timed_task complete")
+
+
+@exampletasks.register(name="examples.simple_task", compression_type=CompressionType.ZSTD)
+def simple_task_compressed(*args: list[Any], **kwargs: dict[str, Any]) -> None:
+    sleep(0.1)
+    logger.debug("simple_task_compressed complete")

src/sentry/taskworker/workerchild.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import base64
 import contextlib
 import logging
 import queue
@@ -13,6 +14,7 @@
 # XXX: Don't import any modules that will import django here, do those within child_process
 import orjson
 import sentry_sdk
+import zstandard as zstd
 from sentry_protos.taskbroker.v1.taskbroker_pb2 import (
     TASK_ACTIVATION_STATUS_COMPLETE,
     TASK_ACTIVATION_STATUS_FAILURE,
@@ -25,6 +27,7 @@
 
 from sentry.taskworker.client.inflight_task_activation import InflightTaskActivation
 from sentry.taskworker.client.processing_result import ProcessingResult
+from sentry.taskworker.constants import CompressionType
 
 logger = logging.getLogger("sentry.taskworker.worker")
 
@@ -77,6 +80,19 @@ def get_at_most_once_key(namespace: str, taskname: str, task_id: str) -> str:
     return f"tw:amo:{namespace}:{taskname}:{task_id}"
 
 
+def load_parameters(data: str, headers: dict[str, str]) -> dict[str, Any]:
+    compression_type = headers.get("compression-type", None)
+    if not compression_type or compression_type == CompressionType.PLAINTEXT.value:
+        return orjson.loads(data)
+    elif compression_type == CompressionType.ZSTD.value:
+        return orjson.loads(zstd.decompress(base64.b64decode(data)))
+    else:
+        logger.error(
+            "Unsupported compression type: %s. Continuing with plaintext.", compression_type
+        )
+        return orjson.loads(data)
+
+
 def status_name(status: TaskActivationStatus.ValueType) -> str:
     """Convert a TaskActivationStatus to a human readable name"""
     if status == TASK_ACTIVATION_STATUS_COMPLETE:
@@ -316,10 +332,11 @@ def handle_alarm(signum: int, frame: FrameType | None) -> None:
 
     def _execute_activation(task_func: Task[Any, Any], activation: TaskActivation) -> None:
         """Invoke a task function with the activation parameters."""
-        parameters = orjson.loads(activation.parameters)
+        headers = {k: v for k, v in activation.headers.items()}
+        parameters = load_parameters(activation.parameters, headers)
+
         args = parameters.get("args", [])
         kwargs = parameters.get("kwargs", {})
-        headers = {k: v for k, v in activation.headers.items()}
 
         transaction = sentry_sdk.continue_trace(
             environ_or_headers=headers,

tests/sentry/taskworker/test_registry.py

@@ -1,18 +1,23 @@
+import base64
 from concurrent.futures import Future
 from unittest.mock import Mock
 
+import orjson
 import pytest
+import zstandard as zstd
 from django.test.utils import override_settings
 from sentry_protos.taskbroker.v1.taskbroker_pb2 import (
     ON_ATTEMPTS_EXCEEDED_DEADLETTER,
     ON_ATTEMPTS_EXCEEDED_DISCARD,
 )
 
 from sentry.conf.types.kafka_definition import Topic
+from sentry.taskworker.constants import CompressionType
 from sentry.taskworker.registry import TaskNamespace, TaskRegistry
 from sentry.taskworker.retry import LastAction, Retry
 from sentry.taskworker.router import DefaultRouter
 from sentry.taskworker.task import Task
+from sentry.testutils.helpers.options import override_options
 
 
 def test_namespace_register_task() -> None:
@@ -133,6 +138,34 @@ def simple_task() -> None:
     assert proto_message == activation.SerializeToString()
 
 
+@pytest.mark.django_db
+def test_namespace_send_task_with_compression() -> None:
+    namespace = TaskNamespace(
+        name="tests",
+        router=DefaultRouter(),
+        retry=None,
+    )
+
+    @namespace.register(name="test.compression_task", compression_type=CompressionType.ZSTD)
+    def simple_task_with_compression(param: str) -> None:
+        raise NotImplementedError
+
+    with override_options({"taskworker.enable_compression.rollout": 1.0}):
+        activation = simple_task_with_compression.create_activation(
+            args=["test_arg"], kwargs={"test_key": "test_value"}
+        )
+
+    assert activation.headers.get("compression-type") == CompressionType.ZSTD.value
+
+    expected_params = {"args": ["test_arg"], "kwargs": {"test_key": "test_value"}}
+
+    decoded_data = base64.b64decode(activation.parameters.encode("utf-8"))
+    decompressed_data = zstd.decompress(decoded_data)
+    actual_params = orjson.loads(decompressed_data)
+
+    assert actual_params == expected_params
+
+
 @pytest.mark.django_db
 def test_namespace_send_task_with_retry() -> None:
     namespace = TaskNamespace(

tests/sentry/taskworker/test_worker.py

@@ -1,10 +1,13 @@
+import base64
 import queue
 import time
 from multiprocessing import Event
 from unittest import mock
 
 import grpc
+import orjson
 import pytest
+import zstandard as zstd
 from sentry_protos.taskbroker.v1.taskbroker_pb2 import (
     ON_ATTEMPTS_EXCEEDED_DISCARD,
     TASK_ACTIVATION_STATUS_COMPLETE,
@@ -18,6 +21,7 @@
 
 from sentry.taskworker.client.inflight_task_activation import InflightTaskActivation
 from sentry.taskworker.client.processing_result import ProcessingResult
+from sentry.taskworker.constants import CompressionType
 from sentry.taskworker.retry import NoRetriesRemainingError
 from sentry.taskworker.state import current_task
 from sentry.taskworker.worker import TaskWorker
@@ -120,6 +124,30 @@
     ),
 )
 
+COMPRESSED_TASK = InflightTaskActivation(
+    host="localhost:50051",
+    receive_timestamp=0,
+    activation=TaskActivation(
+        id="compressed_task_123",
+        taskname="examples.simple_task",
+        namespace="examples",
+        parameters=base64.b64encode(
+            zstd.compress(
+                orjson.dumps(
+                    {
+                        "args": ["test_arg1", "test_arg2"],
+                        "kwargs": {"test_key": "test_value", "number": 42},
+                    }
+                )
+            )
+        ).decode("utf8"),
+        headers={
+            "compression-type": CompressionType.ZSTD.value,
+        },
+        processing_deadline_duration=2,
+    ),
+)
+
 
 @pytest.mark.django_db
 class TestTaskWorker(TestCase):
@@ -672,3 +700,28 @@ def test_child_process_terminate_task(mock_capture: mock.Mock) -> None:
     assert result.status == TASK_ACTIVATION_STATUS_FAILURE
     assert mock_capture.call_count == 1
     assert type(mock_capture.call_args.args[0]) is ProcessingDeadlineExceeded
+
+
+@pytest.mark.django_db
+@mock.patch("sentry.taskworker.workerchild.capture_checkin")
+def test_child_process_decompression(mock_capture_checkin) -> None:
+
+    todo: queue.Queue[InflightTaskActivation] = queue.Queue()
+    processed: queue.Queue[ProcessingResult] = queue.Queue()
+    shutdown = Event()
+
+    todo.put(COMPRESSED_TASK)
+    child_process(
+        todo,
+        processed,
+        shutdown,
+        max_task_count=1,
+        processing_pool_name="test",
+        process_type="fork",
+    )
+
+    assert todo.empty()
+    result = processed.get()
+    assert result.task_id == COMPRESSED_TASK.activation.id
+    assert result.status == TASK_ACTIVATION_STATUS_COMPLETE
+    assert mock_capture_checkin.call_count == 0