fix(taskworker) Rebalance when a broker has no tasks (#92127)

In multi-broker deployments that are lower volume it is possible for a
worker to select a broker that is empty. When we get a notfound
response, we should switch brokers.

Author: markstory

src/sentry/taskworker/client.py

@@ -166,6 +166,8 @@ def get_task(self, namespace: str | None = None) -> TaskActivation | None:
                 "taskworker.client.rpc_error", tags={"method": "GetTask", "status": err.code().name}
             )
             if err.code() == grpc.StatusCode.NOT_FOUND:
+                # Because our current broker doesn't have any tasks, try rebalancing.
+                self._num_tasks_before_rebalance = 0
                 return None
             raise
         if response.HasField("task"):
@@ -207,6 +209,8 @@ def update_task(
                 tags={"method": "SetTaskStatus", "status": err.code().name},
             )
             if err.code() == grpc.StatusCode.NOT_FOUND:
+                # The current broker is empty, switch.
+                self._num_tasks_before_rebalance = 0
                 return None
             raise
         if response.HasField("task"):

tests/sentry/taskworker/test_client.py

@@ -429,3 +429,83 @@ def test_client_loadbalance():
 
             client.update_task(task_3.id, TASK_ACTIVATION_STATUS_COMPLETE, None)
             assert client._task_id_to_host == {}
+
+
+@django_db_all
+def test_client_loadbalance_on_notfound():
+    channel_0 = MockChannel()
+    channel_0.add_response(
+        "/sentry_protos.taskbroker.v1.ConsumerService/GetTask",
+        MockGrpcError(grpc.StatusCode.NOT_FOUND, "no pending task found"),
+    )
+
+    channel_1 = MockChannel()
+    channel_1.add_response(
+        "/sentry_protos.taskbroker.v1.ConsumerService/GetTask",
+        GetTaskResponse(
+            task=TaskActivation(
+                id="1",
+                namespace="testing",
+                taskname="do_thing",
+                parameters="",
+                headers={},
+                processing_deadline_duration=10,
+            )
+        ),
+    )
+    channel_1.add_response(
+        "/sentry_protos.taskbroker.v1.ConsumerService/SetTaskStatus",
+        MockGrpcError(grpc.StatusCode.NOT_FOUND, "no pending task found"),
+    )
+
+    channel_2 = MockChannel()
+    channel_2.add_response(
+        "/sentry_protos.taskbroker.v1.ConsumerService/GetTask",
+        GetTaskResponse(
+            task=TaskActivation(
+                id="2",
+                namespace="testing",
+                taskname="do_thing",
+                parameters="",
+                headers={},
+                processing_deadline_duration=10,
+            )
+        ),
+    )
+
+    with patch("sentry.taskworker.client.grpc.insecure_channel") as mock_channel:
+        mock_channel.side_effect = [channel_0, channel_1, channel_2]
+        with patch("sentry.taskworker.client.random.choice") as mock_randchoice:
+            mock_randchoice.side_effect = [
+                "localhost-0:50051",
+                "localhost-1:50051",
+                "localhost-2:50051",
+            ]
+            client = TaskworkerClient(
+                "localhost:50051", num_brokers=3, max_tasks_before_rebalance=30
+            )
+
+            # Fetch from the first channel, it should return notfound
+            task_0 = client.get_task()
+            assert task_0 is None
+
+            # Fetch again, this time from channel_1
+            task_1 = client.get_task()
+            assert task_1 and task_1.id == "1"
+
+            assert client._task_id_to_host == {
+                "1": "localhost-1:50051",
+            }
+
+            res = client.update_task(task_1.id, TASK_ACTIVATION_STATUS_COMPLETE, None)
+            assert res is None
+            assert client._task_id_to_host == {}
+
+            # Because SetStatus on channel_1 returned notfound the client
+            # should switch brokers.
+            task_2 = client.get_task()
+            assert task_2 and task_2.id == "2"
+
+            assert client._task_id_to_host == {
+                "2": "localhost-2:50051",
+            }