feat(taskworker): Add Push Mode to Taskworker by james-mcnulty · Pull Request #576 · getsentry/taskbroker

james-mcnulty · 2026-03-19T00:20:00Z

Linear

Description

Currently, taskworkers pull tasks from taskbrokers via RPC. This approach works, but has some drawbacks. Therefore, we want taskbrokers to push tasks to taskworkers instead. Read this page on Notion for more information.

This PR allows users to run the taskworker in push mode via the --push-mode and --grpc-port CLI options.

Option	Type	Default	Description
`--push-mode`	`bool`	`False`	Whether to run in push or pull mode.
`--grpc-port`	`int`	`50052`	The port to use for the taskworker gRPC server.

Details

Dependencies

Upgrade sentry-protos from 0.4.11 to 0.8.5 (to use the new worker service schema)

Additions

Define WorkerServicer class in worker.py
Add push_mode and grpc_port fields to TaskWorker and TaskbrokerClient classes
Add push_task method to TaskWorker class
Add gauge method to metrics backend abstract class
Add unit tests for push mode

Modifications

Start gRPC server if push mode, enter pull loop if pull mode
Add _get_stub method to TaskbrokerClient to connect to brokers on the fly by using the activation callback URL

linear-code · 2026-03-19T00:20:04Z

STREAM-822 Add Push Mode to Taskworker

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

clients/python/src/examples/cli.py

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

clients/python/src/examples/cli.py

config.yaml

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

evanh · 2026-03-19T13:50:45Z

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

            channel = grpc.intercept_channel(channel, RequestSignatureInterceptor(secrets))
        return ConsumerServiceStub(channel)

+    def _get_stub(self, host: str) -> ConsumerServiceStub:


Once we switch the taskbroker to be a deployment instead of a statefulset all of this logic will go away, the workers don't need to keep track of individual hosts.

Will the workers use the callback_url on the task requests instead?

I actually don't think we will even need that, we can just have one service that gets configured that all the workers talk too.

One good thing about this is that you can still use push mode with SQLite if you'd like. If we drop callback_url completely, you will have no choice but to use a shared store. It's more flexibility for minimal cost, no?

In our case, callback_url will point to the single broker service that all workers will talk to. But if someone else wants to continue using SQLite, it would point to the specific taskbroker it came from.

But if someone else wants to continue using SQLite, it would point to the specific taskbroker it came from.

If we're not going to be running sqlite long term, then we should remove it (after providing a way to transition to postgres/in memory)

Yes that's true, we want to leave it as we do the migration. Once everything is migrated we can point self-hosted to the existing postgres DB.

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

evanh · 2026-03-19T13:53:20Z

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

+            if timeout is None:
+                self._child_tasks.put(inflight)
+            else:
+                self._child_tasks.put(inflight, timeout=timeout)


I think you can just do this line, there's no need for a conditional.

self.child_tasks.put(inflight, timeout=None) is equivalent to self._child_tasks.put(inflight).

markstory · 2026-03-19T16:37:52Z

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

            channel = grpc.intercept_channel(channel, RequestSignatureInterceptor(secrets))
        return ConsumerServiceStub(channel)

+    def _get_stub(self, host: str) -> ConsumerServiceStub:


Will the workers use the callback_url on the task requests instead?

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

markstory · 2026-03-19T16:45:28Z

clients/python/tests/worker/test_worker.py

+        with mock.patch("taskbroker_client.worker.worker.grpc.server") as mock_grpc_server:
+            mock_grpc_server.return_value = mock_server
+            with mock.patch("taskbroker_client.worker.worker.ThreadPoolExecutor") as mock_tpe:
+                with mock.patch(
+                    "taskbroker_client.worker.worker.taskbroker_pb2_grpc.add_WorkerServiceServicer_to_server"
+                ) as mock_add_servicer:
+                    exitcode = taskworker.start()


This is a lot of mocks. How much of the actual implementation is being exercised here now? Should we consider having an integration test that tests the worker processing messages end to end?

Yes definitely. Once the taskbroker and taskworker push changes are both done, I want to add integration tests for push mode.

clients/python/tests/worker/test_worker.py

sentry · 2026-03-19T20:51:17Z

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

+                    server.stop(grace=5)
+
+                self.shutdown()
+        else:


Bug: In push mode, self.shutdown() is not called if the gRPC server terminates without a KeyboardInterrupt, causing child processes and threads to leak.
_{Severity: MEDIUM}

Suggested Fix

Wrap the server.wait_for_termination() call in a try...finally block. The self.shutdown() method should be moved into the finally block to guarantee its execution, ensuring that resources are always cleaned up regardless of how the server terminates.

Prompt for AI Agent

Review the code at the location below. A potential bug has been identified by an AI agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's not valid. Location: clients/python/src/taskbroker_client/worker/worker.py#L193-L196 Potential issue: In push mode, the gRPC server's `start` method has a potential resource leak. If `server.wait_for_termination()` returns for any reason other than a `KeyboardInterrupt` (e.g., if `server.stop()` is called from another thread), the `self.shutdown()` method is never executed. This prevents proper cleanup, leaving child worker processes and background threads running after the server has stopped. The current implementation only calls `self.shutdown()` within the `except KeyboardInterrupt` block, failing to account for other graceful shutdown scenarios.

cursor

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

cursor · 2026-03-19T20:55:36Z

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

+        with self._host_to_stubs_lock:
+            if host not in self._host_to_stubs:
+                self._host_to_stubs[host] = self._connect_to_host(host)
+            return self._host_to_stubs[host]


Inconsistent lock usage on _host_to_stubs dict

Low Severity

The new _host_to_stubs_lock protects _host_to_stubs in _get_stub, but _get_cur_stub reads and writes the same dict (lines 260–264) without acquiring the lock. In pull mode, _get_cur_stub runs on the main thread while _get_stub runs on the result thread, creating a race on the shared dict with inconsistent locking.

Additional Locations (1)

clients/python/src/taskbroker_client/worker/client.py#L259-L261

This is a valid bug.

evanh

Still a couple bugs from the AI, but otherwise looks good!

evanh · 2026-03-20T15:58:30Z

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

+        with self._host_to_stubs_lock:
+            if host not in self._host_to_stubs:
+                self._host_to_stubs[host] = self._connect_to_host(host)
+            return self._host_to_stubs[host]


This is a valid bug.

Add Push Mode to Taskworker + Unit Tests

70ee047

james-mcnulty requested a review from a team as a code owner March 19, 2026 00:20

sentry bot reviewed Mar 19, 2026

View reviewed changes

clients/python/src/taskbroker_client/worker/worker.py Show resolved Hide resolved

clients/python/src/examples/cli.py Outdated Show resolved Hide resolved

cursor bot reviewed Mar 19, 2026

View reviewed changes

clients/python/src/taskbroker_client/worker/worker.py Show resolved Hide resolved

clients/python/src/taskbroker_client/worker/worker.py Show resolved Hide resolved

clients/python/src/examples/cli.py Show resolved Hide resolved

Fix Worker CLI Options

357099b

sentry bot reviewed Mar 19, 2026

View reviewed changes

clients/python/src/examples/cli.py Outdated Show resolved Hide resolved

cursor bot reviewed Mar 19, 2026

View reviewed changes

clients/python/src/examples/cli.py Outdated Show resolved Hide resolved

config.yaml Outdated Show resolved Hide resolved

Fix gRPC Port CLI Option Duplicate

eef475a

cursor bot reviewed Mar 19, 2026

View reviewed changes

clients/python/src/taskbroker_client/worker/worker.py Show resolved Hide resolved

Stop Server Only if Exists

1988ba9

sentry bot reviewed Mar 19, 2026

View reviewed changes

clients/python/src/taskbroker_client/worker/worker.py Show resolved Hide resolved

Make Sure Server Variable Exists

1df88aa

sentry bot reviewed Mar 19, 2026

View reviewed changes

clients/python/src/taskbroker_client/worker/worker.py Show resolved Hide resolved

evanh reviewed Mar 19, 2026

View reviewed changes

markstory reviewed Mar 19, 2026

View reviewed changes

Pass Timeout to Queue Put Directly

43ec727

cursor bot reviewed Mar 19, 2026

View reviewed changes

clients/python/tests/worker/test_worker.py Outdated Show resolved Hide resolved

Fix Python Client Tests

b61581c

sentry bot reviewed Mar 19, 2026

View reviewed changes

cursor bot reviewed Mar 19, 2026

View reviewed changes

evanh reviewed Mar 20, 2026

View reviewed changes

Uh oh!

Conversation

james-mcnulty commented Mar 19, 2026

Linear

Description

Details

Dependencies

Additions

Modifications

Uh oh!

linear-code bot commented Mar 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sentry bot Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

cursor bot left a comment

Choose a reason for hiding this comment

Uh oh!

cursor bot Mar 19, 2026

Choose a reason for hiding this comment

Inconsistent lock usage on _host_to_stubs dict

Uh oh!

Choose a reason for hiding this comment

Uh oh!

evanh left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Inconsistent lock usage on `_host_to_stubs` dict