feat(taskbroker): Dual-write new parameters_bytes (#602)

* feat(taskbroker): Dual-write parameters_bytes (msgpack) alongside parameters (JSON)

Introduces the new parameters_bytes field on TaskActivation so tasks can
carry raw bytes via msgpack, and adds the worker-side reader that prefers
parameters_bytes with a fallback to the legacy JSON parameters field.

The producer dual-writes both fields by default so this rolls out in a
single commit regardless of worker/broker upgrade order. The
TASKBROKER_CLIENT_PARAMETERS_FORMAT env var (both|json|msgpack) narrows
this once everything is on the new reader.

First step toward STREAM-882 (taskbroker passthrough mode for arbitrary
Kafka topics): https://linear.app/getsentry/issue/STREAM-882

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* ci: Satisfy clippy, black, and mypy on dual-write branch

- upkeep.rs: move #[allow(deprecated)] from individual rstest-expanded
  tests (where it gets swallowed by the macro expansion) to the test
  module as a whole.
- task.py: black formatting.
- pyproject.toml: tell mypy to ignore missing msgpack stubs.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* ref(upkeep tests): Scope allow(deprecated) to the specific lines

Replaces the module-wide allow with three narrow block-scoped allows
exactly where the tests touch the legacy parameters field, so we keep
deprecation warnings active for everything else in the test module.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* switch to sum of both

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: untitaker

Cargo.lock

@@ -42,7 +42,7 @@ sentry = { version = "0.41.0", default-features = false, features = [
     "tracing",
     "logs"
 ] }
-sentry_protos = "0.8.5"
+sentry_protos = "0.8.13"
 serde = "1.0.214"
 serde_yaml = "0.9.34"
 sha2 = "0.10.8"

Cargo.toml

@@ -6,10 +6,11 @@ readme = "README.md"
 dependencies = [
     "sentry-arroyo>=2.38.7",
     "sentry-sdk[http2]>=2.43.0",
-    "sentry-protos>=0.8.5",
+    "sentry-protos>=0.8.13",
     "confluent_kafka>=2.3.0",
     "cronsim>=2.6",
     "grpcio>=1.67.0",
+    "msgpack>=1.0.0",
     "orjson>=3.10.10",
     "protobuf>=5.28.3",
     "redis>=3.4.1",

clients/python/pyproject.toml

@@ -66,3 +66,16 @@ class CompressionType(Enum):
 
     ZSTD = "zstd"
     PLAINTEXT = "plaintext"
+
+
+class ParametersFormat(Enum):
+    """
+    How the producer populates the legacy `parameters` (JSON) and new
+    `parameters_bytes` (msgpack) fields on TaskActivation.
+
+    Set via env var `TASKBROKER_CLIENT_PARAMETERS_FORMAT`. Default BOTH.
+    """
+
+    BOTH = "both"
+    JSON = "json"
+    MSGPACK = "msgpack"

clients/python/src/taskbroker_client/constants.py

@@ -2,12 +2,14 @@
 
 import base64
 import datetime
+import os
 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 msgpack
 import orjson
 import sentry_sdk
 import zstandard as zstd
@@ -22,9 +24,22 @@
     DEFAULT_PROCESSING_DEADLINE,
     MAX_PARAMETER_BYTES_BEFORE_COMPRESSION,
     CompressionType,
+    ParametersFormat,
 )
 from taskbroker_client.retry import Retry
 
+
+def _get_parameters_format() -> ParametersFormat:
+    raw = os.environ.get("TASKBROKER_CLIENT_PARAMETERS_FORMAT", ParametersFormat.BOTH.value)
+    try:
+        return ParametersFormat(raw.lower())
+    except ValueError:
+        raise ValueError(
+            f"Invalid TASKBROKER_CLIENT_PARAMETERS_FORMAT={raw!r}. "
+            f"Expected one of: {', '.join(f.value for f in ParametersFormat)}"
+        )
+
+
 if TYPE_CHECKING:
     from taskbroker_client.registry import TaskNamespace
 
@@ -197,37 +212,58 @@ def create_activation(
                     f"The `{key}` header value is of type {type(value)}"
                 )
 
-        parameters_json = orjson.dumps({"args": args, "kwargs": kwargs})
-        if (
-            len(parameters_json) > MAX_PARAMETER_BYTES_BEFORE_COMPRESSION
-            or self.compression_type == CompressionType.ZSTD
-        ):
-            # Worker uses this header to determine if the parameters are decompressed
+        parameters_format = _get_parameters_format()
+        data = {"args": args, "kwargs": kwargs}
+
+        msgpack_bytes = (
+            msgpack.packb(data, use_bin_type=True)
+            if parameters_format in (ParametersFormat.BOTH, ParametersFormat.MSGPACK)
+            else b""
+        )
+        # JSON can't encode some values msgpack can (e.g. raw bytes). In
+        # JSON-only mode we surface the TypeError; in BOTH mode we silently
+        # skip the legacy field so msgpack-aware workers can still run.
+        json_bytes: bytes | None = None
+        if parameters_format in (ParametersFormat.BOTH, ParametersFormat.JSON):
+            try:
+                json_bytes = orjson.dumps(data)
+            except TypeError:
+                if parameters_format == ParametersFormat.JSON:
+                    raise
+
+        should_compress = (
+            self.compression_type == CompressionType.ZSTD
+            or (len(msgpack_bytes) + len(json_bytes or b""))
+            > MAX_PARAMETER_BYTES_BEFORE_COMPRESSION
+        )
+
+        if should_compress:
             headers["compression-type"] = CompressionType.ZSTD.value
             start_time = time.perf_counter()
-            parameters_str = base64.b64encode(zstd.compress(parameters_json)).decode("utf8")
-            end_time = time.perf_counter()
+            parameters_bytes_val = zstd.compress(msgpack_bytes) if msgpack_bytes else b""
+            parameters_str = (
+                base64.b64encode(zstd.compress(json_bytes)).decode("utf8") if json_bytes else ""
+            )
+            elapsed = time.perf_counter() - start_time
 
+            metric_tags = {
+                "namespace": self._namespace.name,
+                "taskname": self.name,
+                "topic": self._namespace.topic,
+            }
             self.namespace.metrics.distribution(
                 "taskworker.producer.compressed_parameters_size",
-                len(parameters_str),
-                tags={
-                    "namespace": self._namespace.name,
-                    "taskname": self.name,
-                    "topic": self._namespace.topic,
-                },
+                len(parameters_bytes_val) or len(parameters_str),
+                tags=metric_tags,
             )
             self.namespace.metrics.distribution(
                 "taskworker.producer.compression_time",
-                end_time - start_time,
-                tags={
-                    "namespace": self._namespace.name,
-                    "taskname": self.name,
-                    "topic": self._namespace.topic,
-                },
+                elapsed,
+                tags=metric_tags,
             )
         else:
-            parameters_str = parameters_json.decode("utf8")
+            parameters_bytes_val = msgpack_bytes
+            parameters_str = json_bytes.decode("utf8") if json_bytes else ""
 
         return TaskActivation(
             id=uuid4().hex,
@@ -236,6 +272,7 @@ def create_activation(
             taskname=self.name,
             headers=headers,
             parameters=parameters_str,
+            parameters_bytes=parameters_bytes_val,
             retry_state=self._create_retry_state(),
             received_at=received_at,
             processing_deadline_duration=processing_deadline,

clients/python/src/taskbroker_client/task.py

@@ -12,6 +12,7 @@
 from typing import Any
 
 # XXX: Don't import any modules that will import django here, do those within child_process
+import msgpack
 import orjson
 import sentry_sdk
 import zstandard as zstd
@@ -59,17 +60,28 @@ def timeout_alarm(
         signal.signal(signal.SIGALRM, original)
 
 
-def load_parameters(data: str, headers: dict[str, str]) -> dict[str, Any]:
+def load_parameters(activation: TaskActivation) -> dict[str, Any]:
+    headers = dict(activation.headers)
     compression_type = headers.get("compression-type", None)
+
+    # Prefer new msgpack field
+    if activation.parameters_bytes:
+        data = activation.parameters_bytes
+        if compression_type == CompressionType.ZSTD.value:
+            data = zstd.decompress(data)
+        return msgpack.unpackb(data, raw=False)
+
+    # Legacy JSON fallback
+    data_str = activation.parameters
     if not compression_type or compression_type == CompressionType.PLAINTEXT.value:
-        return orjson.loads(data)
+        return orjson.loads(data_str)
     elif compression_type == CompressionType.ZSTD.value:
-        return orjson.loads(zstd.decompress(base64.b64decode(data)))
+        return orjson.loads(zstd.decompress(base64.b64decode(data_str)))
     else:
         logger.error(
             "Unsupported compression type: %s. Continuing with plaintext.", compression_type
         )
-        return orjson.loads(data)
+        return orjson.loads(data_str)
 
 
 def status_name(status: TaskActivationStatus.ValueType) -> str:
@@ -314,12 +326,12 @@ def _execute_activation(
         context_hooks: Sequence[ContextHook] = (),
     ) -> None:
         """Invoke a task function with the activation parameters."""
-        headers = {k: v for k, v in activation.headers.items()}
-        parameters = load_parameters(activation.parameters, headers)
+        parameters = load_parameters(activation)
 
         args = parameters.get("args", [])
         kwargs = parameters.get("kwargs", {})
 
+        headers = dict(activation.headers)
         transaction = sentry_sdk.continue_trace(
             environ_or_headers=headers,
             op="queue.task.taskworker",

clients/python/src/taskbroker_client/worker/workerchild.py

@@ -2,6 +2,7 @@
 from concurrent.futures import Future
 from unittest.mock import Mock
 
+import msgpack
 import orjson
 import pytest
 import zstandard as zstd
@@ -174,12 +175,14 @@ def simple_task_with_compression(param: str) -> None:
 
     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)
+    decompressed_data = zstd.decompress(activation.parameters_bytes)
+    actual_params = msgpack.unpackb(decompressed_data, raw=False)
 
     assert actual_params == expected_params
 
+    legacy_decompressed = zstd.decompress(base64.b64decode(activation.parameters.encode("utf-8")))
+    assert orjson.loads(legacy_decompressed) == expected_params
+
 
 def test_namespace_send_task_with_auto_compression() -> None:
     namespace = TaskNamespace(
@@ -204,12 +207,14 @@ def simple_task_with_compression(param: str) -> None:
 
     expected_params = {"args": big_args, "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)
+    decompressed_data = zstd.decompress(activation.parameters_bytes)
+    actual_params = msgpack.unpackb(decompressed_data, raw=False)
 
     assert actual_params == expected_params
 
+    legacy_decompressed = zstd.decompress(base64.b64decode(activation.parameters.encode("utf-8")))
+    assert orjson.loads(legacy_decompressed) == expected_params
+
 
 def test_namespace_send_task_with_retry() -> None:
     namespace = TaskNamespace(

clients/python/tests/test_registry.py

@@ -4,6 +4,7 @@
 from typing import Any
 from unittest.mock import patch
 
+import msgpack
 import orjson
 import pytest
 import sentry_sdk
@@ -78,9 +79,9 @@ def test_func(*args: Any, **kwargs: Any) -> None:
 
     activation = call_params.args[0]
     assert activation.expires == 10
-    assert activation.parameters == orjson.dumps(
-        {"args": ["arg2"], "kwargs": {"org_id": 2}}
-    ).decode("utf-8")
+    expected_params = {"args": ["arg2"], "kwargs": {"org_id": 2}}
+    assert msgpack.unpackb(activation.parameters_bytes, raw=False) == expected_params
+    assert orjson.loads(activation.parameters) == expected_params
 
 
 def test_apply_async_countdown(task_namespace: TaskNamespace) -> None:
@@ -99,9 +100,9 @@ def test_func(*args: Any, **kwargs: Any) -> None:
 
     activation = call_params.args[0]
     assert activation.delay == 600
-    assert activation.parameters == orjson.dumps(
-        {"args": ["arg2"], "kwargs": {"org_id": 2}}
-    ).decode("utf-8")
+    expected_params = {"args": ["arg2"], "kwargs": {"org_id": 2}}
+    assert msgpack.unpackb(activation.parameters_bytes, raw=False) == expected_params
+    assert orjson.loads(activation.parameters) == expected_params
 
 
 def test_delay_immediate_mode(task_namespace: TaskNamespace) -> None:
@@ -268,11 +269,14 @@ def with_parameters(one: str, two: int, org_id: int) -> None:
         raise NotImplementedError
 
     activation = with_parameters.create_activation(["one", 22], {"org_id": 99})
-    params = orjson.loads(activation.parameters)
+    params = msgpack.unpackb(activation.parameters_bytes, raw=False)
     assert params["args"]
     assert params["args"] == ["one", 22]
     assert params["kwargs"] == {"org_id": 99}
 
+    json_params = orjson.loads(activation.parameters)
+    assert json_params == params
+
 
 def test_create_activation_tracing(task_namespace: TaskNamespace) -> None:
     @task_namespace.register(name="test.parameters")