Add Application to GetTaskRequest and FetchNextTask Initializations

Author: george-sentry

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

@@ -250,7 +250,9 @@ def _get_cur_stub(self) -> tuple[str, ConsumerServiceStub]:
         self._num_tasks_before_rebalance -= 1
         return self._cur_host, self._host_to_stubs[self._cur_host]
 
-    def get_task(self, namespace: str | None = None) -> InflightTaskActivation | None:
+    def get_task(
+        self, namespace: str | None = None, application: str | None = None
+    ) -> InflightTaskActivation | None:
         """
         Fetch a pending task.
 
@@ -259,7 +261,7 @@ def get_task(self, namespace: str | None = None) -> InflightTaskActivation | Non
         """
         self._emit_health_check()
 
-        request = GetTaskRequest(namespace=namespace)
+        request = GetTaskRequest(namespace=namespace, application=application)
         try:
             host, stub = self._get_cur_stub()
             with self._metrics.timer("taskworker.get_task.rpc", tags={"host": host}):

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

@@ -51,6 +51,7 @@ def __init__(
         app_module: str,
         broker_hosts: list[str],
         max_child_task_count: int | None = None,
+        application: str | None = None,
         namespace: str | None = None,
         concurrency: int = 1,
         child_tasks_queue_maxsize: int = DEFAULT_WORKER_QUEUE_SIZE,
@@ -65,6 +66,7 @@ def __init__(
         self.options = kwargs
         self._app_module = app_module
         self._max_child_task_count = max_child_task_count
+        self._application = application
         self._namespace = namespace
         self._concurrency = concurrency
         app = import_app(app_module)
@@ -269,7 +271,7 @@ def _send_result(self, result: ProcessingResult, fetch: bool = True) -> bool:
         if fetch:
             fetch_next = None
             if not self._child_tasks.full():
-                fetch_next = FetchNextTask(namespace=self._namespace)
+                fetch_next = FetchNextTask(namespace=self._namespace, application=self._application)
 
             next = self._send_update_task(result, fetch_next)
             if next:
@@ -377,7 +379,7 @@ def fetch_task(self) -> InflightTaskActivation | None:
         # Use the shutdown_event as a sleep mechanism
         self._shutdown_event.wait(self._gettask_backoff_seconds)
         try:
-            activation = self.client.get_task(self._namespace)
+            activation = self.client.get_task(self._namespace, self._application)
         except grpc.RpcError as e:
             logger.info(
                 "taskworker.fetch_task.failed",

src/grpc/server.rs

@@ -65,6 +65,8 @@ impl ConsumerService for TaskbrokerServer {
         &self,
         request: Request<SetTaskStatusRequest>,
     ) -> Result<Response<SetTaskStatusResponse>, Status> {
+        println!("Set task status...");
+
         let start_time = Instant::now();
         let id = request.get_ref().id.clone();
 

src/kafka/deserialize_activation.rs

@@ -34,6 +34,8 @@ pub fn new(
         let namespace = activation.namespace.clone();
         let taskname = activation.taskname.clone();
 
+        println!("Task activation ingested - application: {}", application);
+
         metrics::histogram!(
             "consumer.message.payload_size_bytes",
             "namespace" => namespace.clone(),