fix(taskworker) Retain task to host mapping on unavailable errors (#92426)

When a broker is unavailable we should retain the task : host mapping so
that we can redeliver the SetStatus request in the future.

Author: markstory

src/sentry/taskworker/client.py

@@ -196,23 +196,29 @@ def update_task(
             status=status,
             fetch_next_task=fetch_next_task,
         )
+
+        host = self._task_id_to_host.get(task_id, None)
+        if host is None:
+            metrics.incr("taskworker.client.task_id_not_in_client")
+            return None
         try:
             with metrics.timer("taskworker.update_task.rpc"):
-                if task_id not in self._task_id_to_host:
-                    metrics.incr("taskworker.client.task_id_not_in_client")
-                    return None
-                host = self._task_id_to_host.pop(task_id)
                 response = self._host_to_stubs[host].SetTaskStatus(request)
+                del self._task_id_to_host[task_id]
         except grpc.RpcError as err:
             metrics.incr(
                 "taskworker.client.rpc_error",
                 tags={"method": "SetTaskStatus", "status": err.code().name},
             )
             if err.code() == grpc.StatusCode.NOT_FOUND:
+                del self._task_id_to_host[task_id]
+
                 # The current broker is empty, switch.
                 self._num_tasks_before_rebalance = 0
+
                 return None
             raise
+
         if response.HasField("task"):
             self._task_id_to_host[response.task.id] = host
             return response.task

tests/sentry/taskworker/test_client.py

@@ -306,6 +306,25 @@ def test_update_task_not_found():
         assert result is None
 
 
+@django_db_all
+def test_update_task_unavailable_retain_task_to_host():
+    channel = MockChannel()
+    channel.add_response(
+        "/sentry_protos.taskbroker.v1.ConsumerService/SetTaskStatus",
+        MockGrpcError(grpc.StatusCode.UNAVAILABLE, "broker down"),
+    )
+    with patch("sentry.taskworker.client.grpc.insecure_channel") as mock_channel:
+        mock_channel.return_value = channel
+        client = TaskworkerClient("localhost:50051", 1)
+        client._task_id_to_host = {"abc123": "localhost-0:50051"}
+        with pytest.raises(MockGrpcError) as err:
+            client.update_task(
+                "abc123", TASK_ACTIVATION_STATUS_RETRY, FetchNextTask(namespace=None)
+            )
+        assert "broker down" in str(err.value)
+        assert client._task_id_to_host == {"abc123": "localhost-0:50051"}
+
+
 @django_db_all
 def test_client_loadbalance():
     channel_0 = MockChannel()