implement basic consumer push model

src/sentry/runner/commands/run.py

@@ -259,24 +259,20 @@ def worker(ignore_unknown_queues: bool, **options: Any) -> None:
 @log_options()
 @configuration
 def taskworker(**options: Any) -> None:
-    from sentry.taskworker.worker import Worker
+    from sentry.taskworker.worker import serve
 
     with managed_bgtasks(role="taskworker"):
-        worker = Worker(
-            namespace=options.get("namespace"),
-        )
-        worker.start()
-        raise SystemExit(worker.exitcode)
+        serve(**options)
 
 
 @run.command()
 @log_options()
 @configuration
-def kafka_task_grpc_server(**options: Any) -> None:
-    from sentry.taskworker.grpc_server import serve
+def kafka_task_grpc(**options: Any) -> None:
+    from sentry.taskworker.consumer_grpc import start
 
     with managed_bgtasks(role="taskworker"):
-        serve()
+        start()
 
 
 @run.command()

src/sentry/taskworker/consumer_grpc.py

@@ -0,0 +1,57 @@
+import logging
+import time
+
+import grpc
+from sentry_protos.sentry.v1alpha.taskworker_pb2 import (
+    TASK_ACTIVATION_STATUS_PENDING,
+    DispatchRequest,
+)
+from sentry_protos.sentry.v1alpha.taskworker_pb2_grpc import WorkerServiceStub
+
+from sentry.taskworker.pending_task_store import PendingTaskStore
+
+logger = logging.getLogger("sentry.taskworker.grpc_server")
+
+
+class ConsumerGrpc:
+    def __init__(self) -> None:
+        self.pending_task_store = PendingTaskStore()
+        self.host = "localhost"
+        self.server_port = 50051
+        self.channel = grpc.insecure_channel(f"{self.host}:{self.server_port}")
+        self.stub = WorkerServiceStub(self.channel)
+
+    def start(self):
+        while True:
+            self.dispatch_task()
+
+    def dispatch_task(self):
+        in_flight_activation = self.pending_task_store.get_pending_task()
+        if not in_flight_activation:
+            logger.info("No tasks")
+            time.sleep(1)
+            return
+        try:
+            dispatch_task_response = self.stub.Dispatch(
+                DispatchRequest(task_activation=in_flight_activation.activation)
+            )
+            self.pending_task_store.set_task_status(
+                task_id=in_flight_activation.activation.id,
+                task_status=dispatch_task_response.status,
+            )
+        except grpc.RpcError as rpc_error:
+            logger.exception(
+                "Connection lost with worker, code: %s, details: %s",
+                rpc_error.code(),
+                rpc_error.details(),
+            )
+            self.pending_task_store.set_task_status(
+                task_id=in_flight_activation.activation.id,
+                task_status=TASK_ACTIVATION_STATUS_PENDING,
+            )
+            time.sleep(1)
+
+
+def start():
+    consumer_grpc = ConsumerGrpc()
+    consumer_grpc.start()

src/sentry/taskworker/worker.py

@@ -2,26 +2,35 @@
 
 import logging
 import time
+from concurrent import futures
 
+import grpc
 import orjson
 from django.conf import settings
 from sentry_protos.sentry.v1alpha.taskworker_pb2 import (
     TASK_ACTIVATION_STATUS_COMPLETE,
     TASK_ACTIVATION_STATUS_FAILURE,
     TASK_ACTIVATION_STATUS_RETRY,
+    DispatchRequest,
+    DispatchResponse,
 )
+from sentry_protos.sentry.v1alpha.taskworker_pb2_grpc import (
+    WorkerServiceServicer as BaseWorkerServiceServicer,
+)
+from sentry_protos.sentry.v1alpha.taskworker_pb2_grpc import add_WorkerServiceServicer_to_server
 
 from sentry.taskworker.config import TaskNamespace, taskregistry
 
 logger = logging.getLogger("sentry.taskworker")
 
 
-class Worker:
+class WorkerServicer(BaseWorkerServiceServicer):
     __namespace: TaskNamespace | None = None
 
-    def __init__(self, **options):
+    def __init__(self, **options) -> None:
+        super().__init__()
         self.options = options
-        self.exitcode = None
+        self.do_imports()
 
     @property
     def namespace(self) -> TaskNamespace:
@@ -36,25 +45,8 @@ def do_imports(self) -> None:
         for module in settings.TASKWORKER_IMPORTS:
             __import__(module)
 
-    def start(self) -> None:
-        self.do_imports()
-        try:
-            while True:
-                self.process_tasks(self.namespace)
-        except KeyboardInterrupt:
-            self.exitcode = 1
-        except Exception:
-            logger.exception("Worker process crashed")
-
-    def process_tasks(self, namespace: TaskNamespace) -> None:
-        from sentry.taskworker.service.client import task_client
-
-        activation = task_client.get_task(topic=namespace.topic)
-        if not activation:
-            logger.info("No tasks")
-            time.sleep(1)
-            return
-
+    def Dispatch(self, request: DispatchRequest, _) -> DispatchResponse:
+        activation = request.task_activation
         try:
             task_meta = self.namespace.get(activation.taskname)
         except KeyError:
@@ -78,17 +70,12 @@ def process_tasks(self, namespace: TaskNamespace) -> None:
         task_latency = execution_time - task_added_time
         logger.info("task.complete", extra={"latency": task_latency})
 
-        if next_state == TASK_ACTIVATION_STATUS_COMPLETE:
-            logger.info(
-                "taskworker.task.complete", extra={"task": activation.taskname, "id": activation.id}
-            )
-            task_client.complete_task(task_id=activation.id)
-        else:
-            logger.info(
-                "taskworker.task.change_status",
-                extra={"task": activation.taskname, "state": next_state},
-            )
-            task_client.set_task_status(
-                task_id=activation.id,
-                task_status=next_state,
-            )
+        return DispatchResponse(status=next_state)
+
+
+def serve(**options):
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
+    add_WorkerServiceServicer_to_server(WorkerServicer(**options), server)
+    server.add_insecure_port("[::]:50051")
+    server.start()
+    server.wait_for_termination()