Update Taskbroker / Taskworker Client to Support Push Mode

Author: george-sentry

pyproject.toml

@@ -101,7 +101,7 @@ dependencies = [
   "statsd>=3.3.0",
   "structlog>=22.1.0",
   "symbolic>=12.14.1",
-  "taskbroker-client>=0.1.7",
+  "taskbroker-client>=0.1.8",
   "tiktoken>=0.8.0",
   "tokenizers>=0.22.0",
   "tldextract>=5.1.2",

src/sentry/runner/commands/run.py

@@ -135,11 +135,20 @@ def taskworker_scheduler(redis_cluster: str, **options: Any) -> None:
 
 
 @run.command()
+@click.option(
+    "--push-mode", help="Whether to run in PUSH or PULL mode.", default=False, is_flag=True
+)
 @click.option(
     "--rpc-host",
-    help="The hostname and port for the taskworker-rpc. When using num-brokers the hostname will be appended with `-{i}` to connect to individual brokers.",
+    help="The hostname and port for the taskbroker gRPC server. When using num-brokers the hostname will be appended with `-{i}` to connect to individual brokers.",
     default="127.0.0.1:50051",
 )
+@click.option(
+    "--worker-rpc-port",
+    help="Port for the taskworker gRPC server to listen on when it is running in push mode.",
+    default=50052,
+    type=int,
+)
 @click.option(
     "--num-brokers", help="Number of brokers available to connect to", default=None, type=int
 )
@@ -198,6 +207,8 @@ def taskworker(**options: Any) -> None:
 
 
 def run_taskworker(
+    push_mode: bool,
+    worker_rpc_port: int,
     rpc_host: str,
     num_brokers: int | None,
     rpc_host_list: str | None,
@@ -233,7 +244,8 @@ def run_taskworker(
             processing_pool_name=processing_pool_name,
             health_check_file_path=health_check_file_path,
             health_check_sec_per_touch=health_check_sec_per_touch,
-            **options,
+            grpc_port=worker_rpc_port,
+            push_mode=push_mode**options,
         )
         exitcode = worker.start()
         raise SystemExit(exitcode)

src/sentry/taskworker/adapters.py

@@ -59,6 +59,18 @@ def incr(
             sample_rate = settings.SENTRY_METRICS_SAMPLE_RATE
         sentry_metrics.incr(name, amount=int(value), tags=tags, sample_rate=sample_rate)
 
+    def gauge(
+        self,
+        key: str,
+        value: float,
+        instance: str | None = None,
+        tags: Tags | None = None,
+        sample_rate: float = 1,
+        unit: str | None = None,
+        stacklevel: int = 0,
+    ) -> None:
+        return sentry_metrics.gauge(key, value, tags=tags)
+
     def distribution(
         self,
         name: str,