From d02a95ce40d067ef8a45b9d9b6670cb119f38c1d Mon Sep 17 00:00:00 2001 From: Tanuj Nayak Date: Tue, 28 Jul 2026 18:07:24 -0700 Subject: [PATCH 1/3] [BUG](fn-consumer): Skip stale work fetch --- .../src/execution/orchestration/compact.rs | 11 +- .../orchestration/function_execution.rs | 122 +++++++++++++++++- 2 files changed, 129 insertions(+), 4 deletions(-) diff --git a/rust/worker/src/execution/orchestration/compact.rs b/rust/worker/src/execution/orchestration/compact.rs index 33f7dc4d712..207456d66b3 100644 --- a/rust/worker/src/execution/orchestration/compact.rs +++ b/rust/worker/src/execution/orchestration/compact.rs @@ -16,7 +16,7 @@ use chroma_segment::{ spann_provider::SpannProvider, types::{ChromaSegmentWriter, VectorSegmentWriter}, }; -use chroma_sysdb::SysDb; +use chroma_sysdb::{sysdb::GetAttachedFunctionError, SysDb}; use chroma_system::{ wrap, ComponentHandle, Dispatcher, Orchestrator, OrchestratorContext, PanicError, System, TaskError, @@ -39,6 +39,7 @@ use super::register_orchestrator::{CollectionRegisterInfo, RegisterOrchestrator} use crate::execution::{ operators::{ + finish_async_work::FinishAsyncWorkError, get_attached_function::{GetAttachedFunctionInput, GetAttachedFunctionOperator}, materialize_logs::MaterializeLogOutput, }, @@ -293,6 +294,10 @@ pub enum CompactionError { CompactionContextError(#[from] CompactionContextError), #[error("Error fetching logs: {0}")] DataFetchError(#[from] LogFetchOrchestratorError), + #[error("Error resolving attached function state: {0}")] + AttachedFunctionState(#[from] GetAttachedFunctionError), + #[error("Error finishing async attached function work: {0}")] + FinishAsyncWork(#[from] FinishAsyncWorkError), #[error("Error registering collection: {0}")] RegisterError(#[from] RegisterOrchestratorError), #[error("Panic during compaction: {0}")] @@ -322,6 +327,8 @@ impl ChromaError for CompactionError { CompactionError::AttachedFunction(e) => e.code(), CompactionError::CompactionContextError(e) => e.code(), CompactionError::DataFetchError(e) => e.code(), + CompactionError::AttachedFunctionState(e) => e.code(), + CompactionError::FinishAsyncWork(e) => e.code(), CompactionError::RegisterError(e) => e.code(), CompactionError::PanicError(e) => e.code(), CompactionError::InvariantViolation(_) => ErrorCodes::Internal, @@ -335,6 +342,8 @@ impl ChromaError for CompactionError { Self::AttachedFunction(e) => e.should_trace_error(), Self::CompactionContextError(e) => e.should_trace_error(), Self::DataFetchError(e) => e.should_trace_error(), + Self::AttachedFunctionState(e) => e.should_trace_error(), + Self::FinishAsyncWork(e) => e.should_trace_error(), Self::PanicError(e) => e.should_trace_error(), Self::RegisterError(e) => e.should_trace_error(), Self::InvariantViolation(_) => true, diff --git a/rust/worker/src/execution/orchestration/function_execution.rs b/rust/worker/src/execution/orchestration/function_execution.rs index 5ba35afe3d9..ab51996ed79 100644 --- a/rust/worker/src/execution/orchestration/function_execution.rs +++ b/rust/worker/src/execution/orchestration/function_execution.rs @@ -1,11 +1,13 @@ use chroma_error::source_chain_contains; use chroma_log::grpc_log::GrpcPullLogsError; -use chroma_system::System; +use chroma_system::{Operator, System}; use chroma_types::{AttachedFunction, AttachedFunctionUuid, CollectionUuid, DatabaseName}; use uuid::Uuid; use crate::execution::operators::{ - fetch_log::FetchLogError, materialize_logs::MaterializeLogOutput, + fetch_log::FetchLogError, + finish_async_work::{FinishAsyncWorkInput, FinishAsyncWorkItem, FinishAsyncWorkOperator}, + materialize_logs::MaterializeLogOutput, }; use super::{ @@ -50,6 +52,15 @@ fn has_reached_queue_frontier(completion_offset: i64, queue_compaction_offset: i completion_offset >= queue_compaction_offset } +fn can_finish_stale_work( + completion_offset: i64, + collection_head: i64, + queue_compaction_offset: i64, +) -> bool { + completion_offset == collection_head + && has_reached_queue_frontier(completion_offset, queue_compaction_offset) +} + impl FunctionExecutionContext { pub fn new(compaction_context: &CompactionContext) -> Self { Self { @@ -177,6 +188,68 @@ impl FunctionExecutionContext { } } + /// Looks up the current completion frontier without materializing any logs. + /// + /// Work queue records can outlive the invocation they represent. Check this + /// frontier before fetching logs so a stale queue item does not cause an + /// expensive backfill to be loaded only to be discarded later. + async fn get_attached_function_completion_offset( + compaction_context: CompactionContext, + collection_id: CollectionUuid, + attached_function_id: AttachedFunctionUuid, + ) -> Result { + let mut sysdb = compaction_context.sysdb.clone(); + let attached_function = sysdb + .get_attached_functions(None, Some(collection_id), vec![attached_function_id], true) + .await? + .into_iter() + .find(|attached_function| attached_function.id == attached_function_id) + .ok_or(CompactionError::InvariantViolation( + "Missing resolved attached function state for fn-consumer input collection", + ))?; + + Ok(attached_function.completion_offset as i64) + } + + async fn get_collection_head( + compaction_context: CompactionContext, + collection_id: CollectionUuid, + ) -> Result { + let mut sysdb = compaction_context.sysdb.clone(); + let collection_info = sysdb + .get_collection_with_segments(None, collection_id) + .await + .map_err(|_| { + CompactionError::InvariantViolation("Failed to resolve fn-consumer collection head") + })?; + + Ok(collection_info.collection.log_position) + } + + async fn finish_stale_work( + compaction_context: CompactionContext, + attached_function_id: AttachedFunctionUuid, + collection_id: CollectionUuid, + new_completion_offset: i64, + ) -> Result<(), CompactionError> { + let work_queue_client = compaction_context.work_queue_client.clone().ok_or( + CompactionError::InvariantViolation("Work queue client not available for fn-consumer"), + )?; + + FinishAsyncWorkOperator::new() + .run(&FinishAsyncWorkInput::new( + attached_function_id, + vec![FinishAsyncWorkItem { + input_collection_id: collection_id, + completion_offset: new_completion_offset, + }], + work_queue_client, + )) + .await?; + + Ok(()) + } + async fn resolve_shared_input_database_name( compaction_context: CompactionContext, fn_inputs: &[FunctionExecutionInput], @@ -223,6 +296,42 @@ impl FunctionExecutionContext { Self::resolve_shared_input_database_name(base_context.clone(), &fn_inputs).await?; let mut input_collection_data = Vec::with_capacity(fn_inputs.len()); for input in fn_inputs { + let completion_offset = Self::get_attached_function_completion_offset( + base_context.clone(), + input.collection_id, + attached_function_id, + ) + .await?; + + if has_reached_queue_frontier(completion_offset, input.queue_compaction_offset) { + let collection_head = + Self::get_collection_head(base_context.clone(), input.collection_id).await?; + + if can_finish_stale_work( + completion_offset, + collection_head, + input.queue_compaction_offset, + ) { + Self::finish_stale_work( + base_context.clone(), + attached_function_id, + input.collection_id, + collection_head, + ) + .await?; + continue; + } + + tracing::info!( + collection_id = %input.collection_id, + completion_offset, + collection_head, + queue_compaction_offset = input.queue_compaction_offset, + "Skipping stale fn-consumer work item before fetching logs because completion does not equal the collection head" + ); + continue; + } + let collection_data = Box::pin(Self::fetch_function_input_collection_data( base_context.clone(), input.collection_id, @@ -288,7 +397,7 @@ impl FunctionExecutionContext { #[cfg(test)] mod tests { - use super::{has_reached_queue_frontier, FunctionExecutionContext}; + use super::{can_finish_stale_work, has_reached_queue_frontier, FunctionExecutionContext}; use crate::execution::{ operators::fetch_log::FetchLogError, orchestration::{ @@ -314,6 +423,13 @@ mod tests { assert!(has_reached_queue_frontier(0, 0)); } + #[test] + fn acknowledges_stale_work_only_at_the_collection_head() { + assert!(can_finish_stale_work(20_300, 20_300, 20_300)); + assert!(!can_finish_stale_work(20_299, 20_300, 20_300)); + assert!(!can_finish_stale_work(20_300, 20_301, 20_300)); + } + #[test] fn generic_not_found_does_not_trigger_backfill() { let err = CompactionError::DataFetchError(LogFetchOrchestratorError::FetchLog( From d913edeaf945882ae50c48e28c28d8fe49fbe98c Mon Sep 17 00:00:00 2001 From: Tanuj Nayak Date: Wed, 29 Jul 2026 13:34:17 -0700 Subject: [PATCH 2/3] [BUG](fn-consumer): Fix SysDB stale-work lookup --- .../worker/src/execution/orchestration/function_execution.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/rust/worker/src/execution/orchestration/function_execution.rs b/rust/worker/src/execution/orchestration/function_execution.rs index b168a12aa08..15a8d93cf68 100644 --- a/rust/worker/src/execution/orchestration/function_execution.rs +++ b/rust/worker/src/execution/orchestration/function_execution.rs @@ -200,7 +200,10 @@ impl FunctionExecutionContext { ) -> Result { let mut sysdb = compaction_context.sysdb.clone(); let attached_function = sysdb - .get_attached_functions(None, Some(collection_id), vec![attached_function_id], true) + // Do not pass a single ID here: the SysDB client populates both the + // deprecated `id` field and the newer `ids` field for that request. + // The coordinator rejects requests containing both fields. + .get_attached_functions(None, Some(collection_id), vec![], true) .await? .into_iter() .find(|attached_function| attached_function.id == attached_function_id) From 8aae8ab2c853d488b8cc38b1221293f487863d64 Mon Sep 17 00:00:00 2001 From: Tanuj Nayak Date: Wed, 29 Jul 2026 13:36:52 -0700 Subject: [PATCH 3/3] [BUG](fn-consumer): Let WQS finish stale work --- .../orchestration/function_execution.rs | 65 +++---------------- 1 file changed, 8 insertions(+), 57 deletions(-) diff --git a/rust/worker/src/execution/orchestration/function_execution.rs b/rust/worker/src/execution/orchestration/function_execution.rs index 15a8d93cf68..2033f97309a 100644 --- a/rust/worker/src/execution/orchestration/function_execution.rs +++ b/rust/worker/src/execution/orchestration/function_execution.rs @@ -55,15 +55,6 @@ fn has_reached_queue_frontier(completion_offset: i64, queue_compaction_offset: i completion_offset >= queue_compaction_offset } -fn can_finish_stale_work( - completion_offset: i64, - collection_head: i64, - queue_compaction_offset: i64, -) -> bool { - completion_offset == collection_head - && has_reached_queue_frontier(completion_offset, queue_compaction_offset) -} - impl FunctionExecutionContext { pub fn new(compaction_context: &CompactionContext) -> Self { Self { @@ -214,22 +205,7 @@ impl FunctionExecutionContext { Ok(attached_function.completion_offset as i64) } - async fn get_collection_head( - compaction_context: CompactionContext, - collection_id: CollectionUuid, - ) -> Result { - let mut sysdb = compaction_context.sysdb.clone(); - let collection_info = sysdb - .get_collection_with_segments(None, collection_id) - .await - .map_err(|_| { - CompactionError::InvariantViolation("Failed to resolve fn-consumer collection head") - })?; - - Ok(collection_info.collection.log_position) - } - - async fn finish_stale_work( + async fn finish_completed_work( compaction_context: CompactionContext, attached_function_id: AttachedFunctionUuid, collection_id: CollectionUuid, @@ -395,31 +371,13 @@ impl FunctionExecutionContext { .await?; if has_reached_queue_frontier(completion_offset, input.queue_compaction_offset) { - let collection_head = - Self::get_collection_head(base_context.clone(), input.collection_id).await?; - - if can_finish_stale_work( - completion_offset, - collection_head, - input.queue_compaction_offset, - ) { - Self::finish_stale_work( - base_context.clone(), - attached_function_id, - input.collection_id, - collection_head, - ) - .await?; - continue; - } - - tracing::info!( - collection_id = %input.collection_id, + Self::finish_completed_work( + base_context.clone(), + attached_function_id, + input.collection_id, completion_offset, - collection_head, - queue_compaction_offset = input.queue_compaction_offset, - "Skipping stale fn-consumer work item before fetching logs because completion does not equal the collection head" - ); + ) + .await?; continue; } @@ -488,7 +446,7 @@ impl FunctionExecutionContext { #[cfg(test)] mod tests { - use super::{can_finish_stale_work, has_reached_queue_frontier, FunctionExecutionContext}; + use super::{has_reached_queue_frontier, FunctionExecutionContext}; use crate::execution::{ operators::fetch_log::FetchLogError, orchestration::{ @@ -527,13 +485,6 @@ mod tests { assert!(has_reached_queue_frontier(0, 0)); } - #[test] - fn acknowledges_stale_work_only_at_the_collection_head() { - assert!(can_finish_stale_work(20_300, 20_300, 20_300)); - assert!(!can_finish_stale_work(20_299, 20_300, 20_300)); - assert!(!can_finish_stale_work(20_300, 20_301, 20_300)); - } - #[test] fn generic_not_found_does_not_trigger_backfill() { let pull_error: Box = Box::new(GrpcPullLogsError::FailedToPullLogs(