[ENH](fn-consumer): List in-progress jobs

idl/chromadb/proto/fn_consumer.proto

@@ -0,0 +1,18 @@
+syntax = "proto3";
+
+package chroma;
+
+message ListFnConsumerInProgressJobsRequest {}
+
+message FnConsumerInProgressJobInfo {
+  string fn_id = 1;
+  int64 expires_at_epoch_secs = 2;
+}
+
+message ListFnConsumerInProgressJobsResponse {
+  repeated FnConsumerInProgressJobInfo jobs = 1;
+}
+
+service FnConsumer {
+  rpc ListInProgressJobs(ListFnConsumerInProgressJobsRequest) returns (ListFnConsumerInProgressJobsResponse) {}
+}

rust/memberlist/src/client_manager.rs

@@ -7,8 +7,9 @@ use chroma_error::ChromaError;
 use chroma_system::{Component, ComponentContext, Handler};
 use chroma_tracing::GrpcClientTraceService;
 use chroma_types::chroma_proto::{
-    compactor_client::CompactorClient, heap_tender_service_client::HeapTenderServiceClient,
-    log_service_client::LogServiceClient, query_executor_client::QueryExecutorClient,
+    compactor_client::CompactorClient, fn_consumer_client::FnConsumerClient,
+    heap_tender_service_client::HeapTenderServiceClient, log_service_client::LogServiceClient,
+    query_executor_client::QueryExecutorClient,
 };
 use parking_lot::RwLock;
 use std::{
@@ -561,6 +562,20 @@ impl ClientFactory
     }
 }
 
+impl ClientFactory
+    for FnConsumerClient<chroma_tracing::GrpcClientTraceService<tonic::transport::Channel>>
+{
+    fn new_from_channel(channel: GrpcClientTraceService<Channel>) -> Self {
+        FnConsumerClient::new(channel)
+    }
+    fn max_encoding_message_size(self, max_size: usize) -> Self {
+        self.max_encoding_message_size(max_size)
+    }
+    fn max_decoding_message_size(self, max_size: usize) -> Self {
+        self.max_decoding_message_size(max_size)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::super::memberlist_provider::Member;

rust/types/build.rs

@@ -9,6 +9,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         "idl/chromadb/proto/query_executor.proto",
         "idl/chromadb/proto/garbage_collector.proto",
         "idl/chromadb/proto/fault_injection.proto",
+        "idl/chromadb/proto/fn_consumer.proto",
         "idl/chromadb/proto/workqueue.proto",
     ];
 

rust/worker/src/fn_consumer/fn_consumer_manager.rs

@@ -16,7 +16,7 @@ use std::panic::AssertUnwindSafe;
 use std::pin::Pin;
 use std::time::{Duration, SystemTime};
 use thiserror::Error;
-use tokio::sync::mpsc;
+use tokio::sync::{mpsc, oneshot};
 use tracing::{instrument, span};
 
 use crate::compactor::config::CompactorConfig;
@@ -46,6 +46,35 @@ impl InProgressFn {
     }
 }
 
+#[derive(Debug, PartialEq, Eq)]
+pub struct InProgressFnEntry {
+    pub fn_id: AttachedFunctionUuid,
+    pub expires_at_epoch_secs: i64,
+}
+
+#[derive(Debug)]
+pub struct ListInProgressJobsMessage {
+    pub response_tx: oneshot::Sender<Vec<InProgressFnEntry>>,
+}
+
+fn snapshot_in_progress_jobs(
+    in_progress: &HashMap<AttachedFunctionUuid, InProgressFn>,
+) -> Vec<InProgressFnEntry> {
+    let mut entries: Vec<_> = in_progress
+        .iter()
+        .map(|(fn_id, job)| InProgressFnEntry {
+            fn_id: *fn_id,
+            expires_at_epoch_secs: job
+                .expires_at
+                .duration_since(std::time::UNIX_EPOCH)
+                .map(|duration| duration.as_secs() as i64)
+                .unwrap_or(0),
+        })
+        .collect();
+    entries.sort_unstable_by_key(|entry| entry.fn_id.to_string());
+    entries
+}
+
 #[derive(Error, Debug)]
 pub enum DispatchError {
     #[error("Dispatcher not initialized")]
@@ -92,6 +121,18 @@ struct FnDispatchCompletion {
     result: FnDispatchOutput,
 }
 
+fn drain_dispatch_completions(
+    receiver: &mut mpsc::UnboundedReceiver<FnDispatchCompletion>,
+    in_progress: &mut HashMap<AttachedFunctionUuid, InProgressFn>,
+) -> Vec<FnDispatchCompletion> {
+    let mut completions = Vec::new();
+    while let Ok(completion) = receiver.try_recv() {
+        in_progress.remove(&completion.fn_id);
+        completions.push(completion);
+    }
+    completions
+}
+
 #[derive(Clone)]
 pub struct FnConsumerContext {
     pub system: System,
@@ -307,9 +348,10 @@ impl FnConsumerManager {
     }
 
     fn process_completions(&mut self) {
-        while let Ok(completion) = self.dispatch_awaiter_completion_channel.try_recv() {
-            self.in_progress.remove(&completion.fn_id);
-
+        for completion in drain_dispatch_completions(
+            &mut self.dispatch_awaiter_completion_channel,
+            &mut self.in_progress,
+        ) {
             match completion.result {
                 Ok(FnDispatchOutcome::Completed) => {
                     tracing::debug!(
@@ -592,12 +634,91 @@ impl Handler<ScheduledPollMessage> for FnConsumerManager {
     }
 }
 
+#[async_trait]
+impl Handler<ListInProgressJobsMessage> for FnConsumerManager {
+    type Result = ();
+
+    async fn handle(&mut self, message: ListInProgressJobsMessage, _ctx: &ComponentContext<Self>) {
+        self.process_completions();
+        let entries = snapshot_in_progress_jobs(&self.in_progress);
+        if let Err(entries) = message.response_tx.send(entries) {
+            tracing::warn!(
+                job_count = entries.len(),
+                "Failed to send fn-consumer in-progress jobs response"
+            );
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use tokio::sync::oneshot;
     use tokio::time::{timeout, Duration};
 
+    #[test]
+    fn snapshots_in_progress_jobs() {
+        let first_fn_id = AttachedFunctionUuid::new();
+        let second_fn_id = AttachedFunctionUuid::new();
+        let mut in_progress = HashMap::new();
+        in_progress.insert(
+            first_fn_id,
+            InProgressFn {
+                expires_at: std::time::UNIX_EPOCH + Duration::from_secs(20),
+                expiry_logged: false,
+            },
+        );
+        in_progress.insert(
+            second_fn_id,
+            InProgressFn {
+                expires_at: std::time::UNIX_EPOCH + Duration::from_secs(10),
+                expiry_logged: false,
+            },
+        );
+
+        let entries = snapshot_in_progress_jobs(&in_progress);
+        assert_eq!(entries.len(), 2);
+        assert!(entries
+            .windows(2)
+            .all(|pair| pair[0].fn_id.to_string() < pair[1].fn_id.to_string()));
+        assert!(entries
+            .iter()
+            .any(|entry| { entry.fn_id == first_fn_id && entry.expires_at_epoch_secs == 20 }));
+        assert!(entries
+            .iter()
+            .any(|entry| { entry.fn_id == second_fn_id && entry.expires_at_epoch_secs == 10 }));
+    }
+
+    #[test]
+    fn snapshots_empty_in_progress_jobs() {
+        assert!(snapshot_in_progress_jobs(&HashMap::new()).is_empty());
+    }
+
+    #[test]
+    fn draining_completions_removes_finished_jobs() {
+        let fn_id = AttachedFunctionUuid::new();
+        let mut in_progress = HashMap::from([(
+            fn_id,
+            InProgressFn {
+                expires_at: std::time::UNIX_EPOCH + Duration::from_secs(20),
+                expiry_logged: false,
+            },
+        )]);
+        let (completion_tx, mut completion_rx) = mpsc::unbounded_channel();
+        completion_tx
+            .send(FnDispatchCompletion {
+                fn_id,
+                batch_size: 1,
+                result: Ok(FnDispatchOutcome::Completed),
+            })
+            .unwrap();
+
+        let completions = drain_dispatch_completions(&mut completion_rx, &mut in_progress);
+
+        assert_eq!(completions.len(), 1);
+        assert!(in_progress.is_empty());
+    }
+
     #[tokio::test]
     async fn dispatch_awaiter_completes_later_tasks_while_one_is_running() {
         let (task_tx, task_rx) = mpsc::channel(2);

rust/worker/src/fn_consumer/grpc.rs

@@ -0,0 +1,51 @@
+use async_trait::async_trait;
+use chroma_system::ComponentHandle;
+use chroma_types::chroma_proto::{
+    fn_consumer_server::{FnConsumer, FnConsumerServer},
+    FnConsumerInProgressJobInfo, ListFnConsumerInProgressJobsRequest,
+    ListFnConsumerInProgressJobsResponse,
+};
+use tonic::{Request, Response, Status};
+
+use super::fn_consumer_manager::{FnConsumerManager, ListInProgressJobsMessage};
+
+pub struct FnConsumerGrpcServer {
+    manager: ComponentHandle<FnConsumerManager>,
+}
+
+impl FnConsumerGrpcServer {
+    pub fn new(manager: ComponentHandle<FnConsumerManager>) -> Self {
+        Self { manager }
+    }
+
+    pub fn into_service(self) -> FnConsumerServer<Self> {
+        FnConsumerServer::new(self)
+    }
+}
+
+#[async_trait]
+impl FnConsumer for FnConsumerGrpcServer {
+    async fn list_in_progress_jobs(
+        &self,
+        _request: Request<ListFnConsumerInProgressJobsRequest>,
+    ) -> Result<Response<ListFnConsumerInProgressJobsResponse>, Status> {
+        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
+        self.manager
+            .receiver()
+            .send(ListInProgressJobsMessage { response_tx }, None)
+            .await
+            .map_err(|error| Status::internal(error.to_string()))?;
+
+        let jobs = response_rx
+            .await
+            .map_err(|error| Status::internal(format!("Failed to receive response: {error}")))?
+            .into_iter()
+            .map(|entry| FnConsumerInProgressJobInfo {
+                fn_id: entry.fn_id.to_string(),
+                expires_at_epoch_secs: entry.expires_at_epoch_secs,
+            })
+            .collect();
+
+        Ok(Response::new(ListFnConsumerInProgressJobsResponse { jobs }))
+    }
+}

rust/worker/src/fn_consumer/mod.rs

@@ -1,5 +1,6 @@
 pub mod config;
 pub mod fn_consumer_manager;
+mod grpc;
 pub mod server;
 
 pub use server::fn_consumer_service_entrypoint;

rust/worker/src/fn_consumer/server.rs

@@ -1,5 +1,6 @@
 use crate::config::RootConfig;
 use crate::fn_consumer::fn_consumer_manager::FnConsumerManager;
+use crate::fn_consumer::grpc::FnConsumerGrpcServer;
 use crate::work_queue::work_queue_client::WorkQueueClient;
 use chroma_blockstore::provider::BlockfileProvider;
 use chroma_config::registry::Registry;
@@ -172,10 +173,16 @@ pub async fn fn_consumer_service_entrypoint() {
         spann_provider,
     );
     manager.set_dispatcher(dispatcher_handle);
-    let _manager_handle = system.start_component(manager);
+    let manager_handle = system.start_component(manager);
 
     // Create health service for readiness probe
-    let (_health_reporter, health_service) = tonic_health::server::health_reporter();
+    let (health_reporter, health_service) = tonic_health::server::health_reporter();
+    health_reporter
+        .set_serving::<chroma_types::chroma_proto::fn_consumer_server::FnConsumerServer<
+            FnConsumerGrpcServer,
+        >>()
+        .await;
+    let fn_consumer_service = FnConsumerGrpcServer::new(manager_handle).into_service();
 
     let addr = format!("0.0.0.0:{}", service_config.my_port)
         .parse()
@@ -186,6 +193,7 @@ pub async fn fn_consumer_service_entrypoint() {
     // Start server (this blocks forever)
     Server::builder()
         .add_service(health_service)
+        .add_service(fn_consumer_service)
         .serve(addr)
         .await
         .expect("Failed to start fn-consumer service");