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

[enochtangg]

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.

[cursor]

Bug: Compression Rollout Bug Causes UnboundLocalError

An UnboundLocalError occurs for the parameters_str variable when compression_type is ZSTD but the compression rollout condition is not met. The code's else block only assigns parameters_str when compression_type is not ZSTD, leaving parameters_str undefined for tasks configured with ZSTD compression that do not pass the rollout check.

src/sentry/taskworker/task.py#L180-L207

sentry/src/sentry/taskworker/task.py

Lines 180 to 207 in fd30ab8

	parameters_json = orjson.dumps({"args": args, "kwargs": kwargs})
	if self.compression_type == CompressionType.ZSTD:
	option_flag = f"taskworker.{self._namespace.name}.compression.rollout"
	compression_rollout_rate = options.get(option_flag)
	# TODO(taskworker): This option is for added safety and requires a rollout
	# percentage to be set to actually use compression.
	# We can remove this later.
	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")

Fix in Cursor • Fix in Web

---

Bug: Task Name Conflict Overwrites Existing Task

The new simple_task_compressed function is registered with the name "examples.simple_task", which conflicts with an existing simple_task function. This overwrites the original task in the registry, rendering it unreachable. The new task requires a unique name.

src/sentry/taskworker/tasks/examples.py#L92-L96

sentry/src/sentry/taskworker/tasks/examples.py

Lines 92 to 96 in fd30ab8

	@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")

Fix in Cursor • Fix in Web

---

Was this report helpful? Give feedback by reacting with 👍 or 👎

[evanh]

This looks good to me. I'm interested to see how this will impact things.

[markstory]

Do we need to enable compression per namespace? Once we know that the general compression flow works, we could opt-tasks into compression via deploys.

[markstory]

Could we add this to the docstring so it shows up in LSP tooltips?

[markstory]

We could have a single option for enabling compression. We might also want the option check to be done with the self.compression_type comparison so that both the task has to have enabled compression and compression needs to be enabled.

[markstory]

It would be good to have a task covering payload compression when an activation is created as well. There are some existing tests in tests/sentry/taskworker/test_task.py and test_registry.py covering activation creation.