Skip to content

[ENH](fn-consumer): Add concurrency gauge #7670

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tanujnay112 merged 1 commit into from
Sep 2, 2026
+63 −3
Merged

[ENH](fn-consumer): Add concurrency gauge #7670

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 63 additions & 3 deletions rust/worker/src/fn_consumer/fn_consumer_manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ use std::collections::HashMap;
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use thiserror::Error;
use tokio::sync::mpsc;
Expand Down Expand Up @@ -77,6 +79,51 @@ impl ChromaError for DispatchError {
type FnDispatchOutput = Result<FnDispatchOutcome, DispatchError>;
type FnDispatchFuture = Pin<Box<dyn Future<Output = FnDispatchOutput> + Send>>;

#[derive(Clone, Debug)]
struct FnConsumerMetrics {
current_compactions: Arc<AtomicU64>,
}

impl Default for FnConsumerMetrics {
fn default() -> Self {
let current_compactions = Arc::new(AtomicU64::new(0));
let observed_count = current_compactions.clone();
opentelemetry::global::meter("chroma_fn_consumer")
.u64_observable_gauge("fn_consumer_current_compactions")
.with_description("Number of compaction jobs currently running in fn-consumer")
.with_callback(move |observer| {
observer.observe(observed_count.load(Ordering::Relaxed), &[]);
})
.build();
Comment on lines +91 to +97

@graphite-app graphite-app Bot Sep 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The observable gauge instrument returned by .build() is being dropped immediately, which will likely cause the callback to stop being invoked and the metric to not be reported.

OpenTelemetry observable instruments typically need to be kept alive for their callbacks to continue working. The return value should be stored in the FnConsumerMetrics struct:

struct FnConsumerMetrics {
    current_compactions: Arc<AtomicU64>,
    _gauge: ObservableGauge<u64>, // Keep instrument alive
}

impl Default for FnConsumerMetrics {
    fn default() -> Self {
        let current_compactions = Arc::new(AtomicU64::new(0));
        let observed_count = current_compactions.clone();
        let gauge = opentelemetry::global::meter("chroma_fn_consumer")
            .u64_observable_gauge("fn_consumer_current_compactions")
            .with_description("Number of compaction jobs currently running in fn-consumer")
            .with_callback(move |observer| {
                observer.observe(observed_count.load(Ordering::Relaxed), &[]);
            })
            .build();
        Self {
            current_compactions,
            _gauge: gauge,
        }
    }
}

Without storing the instrument, the metric collection will not work in production.

Spotted by Graphite

Fix in Graphite

Is this helpful? React 👍 or 👎 to let us know.

@tanujnay112 tanujnay112 Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I checked the OpenTelemetry Rust 0.27 implementation used by this repo. During .build(), the SDK passes the callback to Pipelines::register_callback, which stores it as an Arc in PipelineInner.callbacks. The returned ObservableGauge is only a PhantomData marker and has no Drop behavior, so dropping that value does not unregister the callback or stop collection. The pipeline continues to own the callback and the captured atomic value, so this metric remains live.

Self {
current_compactions,
}
}
}

impl FnConsumerMetrics {
fn track_compaction(&self) -> ActiveCompactionGuard {
self.current_compactions.fetch_add(1, Ordering::Relaxed);
ActiveCompactionGuard {
metrics: self.clone(),
}
}
}

struct ActiveCompactionGuard {
metrics: FnConsumerMetrics,
}

impl Drop for ActiveCompactionGuard {
fn drop(&mut self) {
let previous = self
.metrics
.current_compactions
.fetch_sub(1, Ordering::Relaxed);
debug_assert!(previous > 0, "active compaction count underflowed");
}
}

struct FnDispatchTask {
fn_id: AttachedFunctionUuid,
future: FnDispatchFuture,
Expand Down Expand Up @@ -182,8 +229,10 @@ impl FnConsumerManager {
// in-progress slot until that completion is drained. Therefore, pending
// completions are bounded by max_concurrent_workers and need no backpressure.
let (completion_tx, completion_rx) = mpsc::unbounded_channel::<FnDispatchCompletion>();
let metrics = FnConsumerMetrics::default();

@rescrv rescrv Sep 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe that dropping this may lose the metrics. Otherwise aI don't get the reason for the clone, either. It reads funny.

@tanujnay112 tanujnay112 Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The extra clone is indeed unnecessary and reads oddly: the original metrics value is dropped immediately, so it could be moved directly into the spawned task. That is a harmless readability issue rather than a correctness issue. I also checked the gauge lifetime against OpenTelemetry Rust 0.27: .build() registers the callback into the SDK pipeline, which owns it independently; the returned ObservableGauge is a marker with no Drop behavior, so dropping it does not lose the metric.

let awaiter_metrics = metrics.clone();
let dispatch_awaiter = tokio::spawn(async move {
fn_dispatch_awaiter_loop(dispatch_awaiter_rx, completion_tx).await;
fn_dispatch_awaiter_loop(dispatch_awaiter_rx, completion_tx, awaiter_metrics).await;
});
Self {
context,
Expand Down Expand Up @@ -502,6 +551,7 @@ fn panic_message(panic_payload: &(dyn Any + Send)) -> String {
async fn fn_dispatch_awaiter_loop(
mut task_rx: mpsc::Receiver<FnDispatchTask>,
completion_tx: mpsc::UnboundedSender<FnDispatchCompletion>,
metrics: FnConsumerMetrics,
) {
let mut futures = FuturesUnordered::new();
loop {
Expand All @@ -513,7 +563,9 @@ async fn fn_dispatch_awaiter_loop(
}
}
Some(task) = task_rx.recv() => {
let metrics = metrics.clone();
futures.push(async move {
let _active_compaction = metrics.track_compaction();
let FnDispatchTask {
fn_id,
future,
Expand Down Expand Up @@ -602,7 +654,11 @@ mod tests {
async fn dispatch_awaiter_completes_later_tasks_while_one_is_running() {
let (task_tx, task_rx) = mpsc::channel(2);
let (completion_tx, mut completion_rx) = mpsc::unbounded_channel();
let awaiter = tokio::spawn(fn_dispatch_awaiter_loop(task_rx, completion_tx));
let awaiter = tokio::spawn(fn_dispatch_awaiter_loop(
task_rx,
completion_tx,
FnConsumerMetrics::default(),
));
let slow_fn_id = AttachedFunctionUuid::new();
let fast_fn_id = AttachedFunctionUuid::new();
let (slow_started_tx, slow_started_rx) = oneshot::channel();
Expand Down Expand Up @@ -660,7 +716,11 @@ mod tests {
async fn dispatch_awaiter_completes_panicked_tasks() {
let (task_tx, task_rx) = mpsc::channel(1);
let (completion_tx, mut completion_rx) = mpsc::unbounded_channel();
let awaiter = tokio::spawn(fn_dispatch_awaiter_loop(task_rx, completion_tx));
let awaiter = tokio::spawn(fn_dispatch_awaiter_loop(
task_rx,
completion_tx,
FnConsumerMetrics::default(),
));
let fn_id = AttachedFunctionUuid::new();

task_tx
Expand Down
Loading