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 b751ec26879..2033f97309a 100644 --- a/rust/worker/src/execution/orchestration/function_execution.rs +++ b/rust/worker/src/execution/orchestration/function_execution.rs @@ -179,6 +179,56 @@ 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 + // 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) + .ok_or(CompactionError::InvariantViolation( + "Missing resolved attached function state for fn-consumer input collection", + ))?; + + Ok(attached_function.completion_offset as i64) + } + + async fn finish_completed_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(()) + } + fn is_purged_pull_error(pull_error: &(dyn ChromaError + 'static)) -> bool { let pull_error = pull_error as &(dyn Error + 'static); @@ -313,6 +363,24 @@ impl FunctionExecutionContext { ))?; let mut input_collection_data = Vec::with_capacity(live_inputs.len()); for input in live_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) { + Self::finish_completed_work( + base_context.clone(), + attached_function_id, + input.collection_id, + completion_offset, + ) + .await?; + continue; + } + let collection_data = Box::pin(Self::fetch_function_input_collection_data( base_context.clone(), input.collection_id,