From add7ce49b03df576560b43c56a5c60897fa80484 Mon Sep 17 00:00:00 2001 From: Bill Njoroge Date: Sat, 3 Oct 2026 12:16:27 -0400 Subject: [PATCH 1/4] Make idle broker polls reader-only and notify-driven Idle long-polling runners were the dominant writer-pool load: every poll opened a writer transaction, re-stamped last_seen_at, and re-ran the full claim scan, and the root broker handler re-polled every 3 seconds while holding the 50s request open. - poll_session now probes on the reader pool first (session, oldest message, active request, pending cancellation, ready check) and only opens the writer transaction when there is a cancellation to deliver or a claim to attempt. Ownership is still revalidated inside the writer transaction, and a probe hit that loses a race just comes back Empty, as before. - next_message_broker_ref_root waits on message_notify until the window deadline instead of re-polling every 3s, matching next_message_broker_ref and the azdo path. Heartbeats stay once per 50s window via the handlers' up-front touch_session. - retire_node_requests no longer selects already-settled requests when retiring (settle is first-result-wins, so the second settle was pure waste: ~5 statements matching zero rows per completion). - Regression tests: an idle poll leaves last_seen_at untouched, and a poll right after submit still claims. --- crates/preloop-runner-server/src/broker.rs | 10 +- .../src/control/pg/dispatch.rs | 112 ++++++++++++++++-- .../src/control/pg/tests.rs | 86 ++++++++++++++ 3 files changed, 194 insertions(+), 14 deletions(-) diff --git a/crates/preloop-runner-server/src/broker.rs b/crates/preloop-runner-server/src/broker.rs index ab789809f..5ee6919b2 100644 --- a/crates/preloop-runner-server/src/broker.rs +++ b/crates/preloop-runner-server/src/broker.rs @@ -777,10 +777,12 @@ pub async fn next_message_broker_ref_root( if wait == 0 || std::time::Instant::now() >= deadline { return Ok(Json(serde_json::Value::Null).into_response()); } - // Wake promptly on cancel/enqueue rather than fixed 250ms sleep. - let remaining = deadline.saturating_duration_since(std::time::Instant::now()); - let slice = remaining.min(Duration::from_secs(3)); - let _ = tokio::time::timeout(slice, shared.state.message_notify.notified()).await; + // One wait per window, like `next_message_broker_ref` and the azdo + // path: the poll above is a cheap reader probe now, and + // `message_notify` fires on every enqueue, cancellation and + // promotion. A lost wake only delays until the window ends. + let _ = + tokio::time::timeout_at(deadline.into(), shared.state.message_notify.notified()).await; } } diff --git a/crates/preloop-runner-server/src/control/pg/dispatch.rs b/crates/preloop-runner-server/src/control/pg/dispatch.rs index 8a93a4f9c..36ef310e3 100644 --- a/crates/preloop-runner-server/src/control/pg/dispatch.rs +++ b/crates/preloop-runner-server/src/control/pg/dispatch.rs @@ -619,9 +619,20 @@ pub(super) async fn retire_node_requests( node_id: &JobId, retirement: Retirement, ) -> Result, ControlError> { + // Settling is first-result-wins, so an already-settled request is a pure + // no-op: don't even select it. (`complete_job` settles the attempt before + // the sweep retires the node; without this the settle statements run a + // second time and match zero rows.) Purge keeps the unfiltered select — + // it replaces the node's rows wholesale. + let pending_only = matches!(retirement, Retirement::Settle(_)); let select = format!( - "{} WHERE q.run_id = $1::text::uuid AND q.job_id = $2 FOR UPDATE OF q", - lookups::REQUEST_SELECT + "{} WHERE q.run_id = $1::text::uuid AND q.job_id = $2{} FOR UPDATE OF q", + lookups::REQUEST_SELECT, + if pending_only { + " AND q.result IS NULL" + } else { + "" + }, ); let rows = tx .query(&select, &[&run_id.0.to_string(), &node_id.0]) @@ -4646,11 +4657,11 @@ impl PgBackend { /// The session's row (mapped id), or `None` when it is unknown/expired. pub(super) async fn session_ref( &self, - tx: &Transaction<'_>, + client: &impl GenericClient, session_id: &str, ) -> Result, ControlError> { let uuid = logic::session_uuid(session_id).to_string(); - Ok(tx + Ok(client .query_opt( "SELECT runner_id, protocol FROM runner_sessions WHERE session_id = $1::text::uuid", &[&uuid], @@ -4686,10 +4697,10 @@ impl PgBackend { /// The oldest unacknowledged message of a session (redelivery-first). pub(super) async fn oldest_session_message( &self, - tx: &Transaction<'_>, + client: &impl GenericClient, session_uuid: &str, ) -> Result, ControlError> { - Ok(tx + Ok(client .query_opt( "SELECT m.message_id, m.message_type, m.request_id, m.body::text, \ s.runner_id IS NULL \ @@ -4715,7 +4726,7 @@ impl PgBackend { /// The session's live request, when it holds one. pub(super) async fn session_active_request( &self, - tx: &Transaction<'_>, + client: &impl GenericClient, session_uuid: &str, ) -> Result, ControlError> { let select = format!( @@ -4723,7 +4734,11 @@ impl PgBackend { ORDER BY q.request_id DESC LIMIT 1", lookups::REQUEST_SELECT ); - match tx.query_opt(&select, &[&session_uuid]).await.map_err(db)? { + match client + .query_opt(&select, &[&session_uuid]) + .await + .map_err(db)? + { Some(row) => Ok(Some(lookups::request_from_row(&row)?)), None => Ok(None), } @@ -4775,10 +4790,10 @@ impl PgBackend { /// A pending (undelivered) cancellation for the attempt, if any. async fn pending_cancellation( &self, - tx: &Transaction<'_>, + client: &impl GenericClient, request_id: i64, ) -> Result, ControlError> { - Ok(tx + Ok(client .query_opt( "SELECT request_id FROM job_cancellations \ WHERE request_id = $1 AND delivered_at IS NULL LIMIT 1", @@ -5154,11 +5169,88 @@ impl PgBackend { .transpose() } + /// `poll_session` probe: the read-only pre-check, run on the reader pool. + /// + /// An idle poll — no session message, no active request, no pending + /// cancellation, no ready work — is the common case for long-polling + /// runners, and answering it here keeps thousands of them off the writer + /// pool entirely. Anything that needs a write (delivering a cancellation, + /// claiming a job) returns `None` and the caller runs the full writer + /// transaction, which re-checks everything under its locks: a probe hit + /// that loses a race just comes back `Empty`, exactly as before. + async fn probe_poll( + &self, + client: &impl GenericClient, + poll: &PollRequest, + ) -> Result, ControlError> { + // Ownership is revalidated inside the claim transaction too: the + // handler caches the runner across a long poll, and a liveness sweep + // can purge the session while it waits. + let Some(session) = self.session_ref(client, &poll.session_id).await? else { + return Err(ControlError::Forbidden( + "session has no runner owner".to_owned(), + )); + }; + if poll.verified_runner_id.is_some() && poll.verified_runner_id != session.runner_id { + return Err(ControlError::Forbidden( + "session belongs to another runner".to_owned(), + )); + } + if let Some(message) = self + .oldest_session_message(client, &session.session_uuid) + .await? + { + return Ok(Some(PollOutcome::Inflight(azdo::TaskAgentMessage { + message_id: message.message_id, + message_type: message.message_type.clone(), + body: message.runner_body(), + iv: None, + }))); + } + if let Some(request) = self + .session_active_request(client, &session.session_uuid) + .await? + { + if self + .pending_cancellation(client, request.request_id) + .await? + .is_some() + { + // The writer delivers the cancellation below. + return Ok(None); + } + let runner_id = session.runner_id.unwrap_or(0); + return Ok(Some(PollOutcome::ActiveRequest { request, runner_id })); + } + if poll.busy { + return Ok(Some(PollOutcome::Empty)); + } + let ready = client + .query_opt( + "SELECT 1 FROM jobs WHERE queue_state = 'ready' LIMIT 1", + &[], + ) + .await + .map_err(db)? + .is_some(); + if ready { + // The writer runs the claim below. + return Ok(None); + } + Ok(Some(PollOutcome::Empty)) + } + /// `poll_session`: redelivery, cancellation, active request, then a claim. pub(super) async fn poll_session( &self, poll: PollRequest, ) -> Result { + // Fast path first: idle polls never take a writer. + let reader = self.reader().await?; + if let Some(outcome) = self.probe_poll(&*reader, &poll).await? { + return Ok(outcome); + } + drop(reader); let mut client = self.writer().await?; let tx = client.transaction().await.map_err(db)?; // Ownership is revalidated inside the claim transaction: the handler diff --git a/crates/preloop-runner-server/src/control/pg/tests.rs b/crates/preloop-runner-server/src/control/pg/tests.rs index eed8bdcc9..909848a6b 100644 --- a/crates/preloop-runner-server/src/control/pg/tests.rs +++ b/crates/preloop-runner-server/src/control/pg/tests.rs @@ -605,6 +605,92 @@ async fn submit_and_claim( (claimed.request, runner.runner.id) } +/// The session's `last_seen_at` as text (microsecond precision), read through +/// the reader pool. +async fn session_last_seen_at(node: &PgBackend, session_id: &str) -> String { + let uuid = crate::control::logic::session_uuid(session_id).to_string(); + let reader = node.reader().await.expect("reader pool"); + reader + .query_one( + "SELECT last_seen_at::text FROM runner_sessions WHERE session_id = $1::text::uuid", + &[&uuid], + ) + .await + .expect("last_seen_at read") + .get(0) +} + +/// Idle polls must not take a writer: with no session message, no active +/// request, no pending cancellation and no ready work, `poll_session` +/// answers from the reader pool and leaves `last_seen_at` untouched. +/// (Previously every poll opened a writer transaction and re-stamped +/// liveness — about five statements per poll, per runner, every few +/// seconds — which is what made thousands of idle long-pollers the +/// dominant writer load.) +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn idle_poll_leaves_last_seen_at_untouched() { + let Some((_guard, url)) = fresh_database_opt().await else { + return skip_no_postgres(); + }; + let node = connect(&url).await; + let runner = node.register_runner(register_runner("r1")).await.unwrap(); + let session = node + .create_session(create_session(runner.runner.id)) + .await + .unwrap(); + let seen_at = session_last_seen_at(&node, &session.session_id).await; + // A re-stamp would land on a later microsecond; sleep so the assertion + // below can actually observe one. + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + let outcome = node + .poll_session(poll(&session.session_id, runner.runner.id)) + .await + .unwrap(); + assert!( + matches!(outcome, PollOutcome::Empty), + "expected an idle poll, got {outcome:?}" + ); + assert_eq!( + session_last_seen_at(&node, &session.session_id).await, + seen_at, + "an idle poll must not rewrite last_seen_at" + ); +} + +/// The probe must not swallow work: a poll right after a submit still +/// claims through the writer path. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn probe_poll_still_claims_after_submit() { + let Some((_guard, url)) = fresh_database_opt().await else { + return skip_no_postgres(); + }; + let node = connect(&url).await; + let runner = node.register_runner(register_runner("r1")).await.unwrap(); + let session = node + .create_session(create_session(runner.runner.id)) + .await + .unwrap(); + // Sanity: idle first. + let outcome = node + .poll_session(poll(&session.session_id, runner.runner.id)) + .await + .unwrap(); + assert!(matches!(outcome, PollOutcome::Empty)); + // Then work arrives: the probe defers to the writer and the claim lands. + let run_id = RunId::new(); + node.submit_run(submit_run(run_id, vec![submit_job(run_id, "build", 1)])) + .await + .unwrap(); + let outcome = node + .poll_session(poll(&session.session_id, runner.runner.id)) + .await + .unwrap(); + let PollOutcome::Claimed(claimed) = outcome else { + panic!("expected a claim after submit, got {outcome:?}"); + }; + assert_eq!(claimed.queued.job_id, JobId("build".to_owned())); +} + /// Far more concurrent acquires than pooled readers must all complete. An /// acquire that held one reader while checking out a second deadlocked the /// pool as soon as every reader was held by an acquire waiting for another — From bbeafd75e6fe72d37312030f7749b6ebb943cccf Mon Sep 17 00:00:00 2001 From: Bill Njoroge Date: Sat, 3 Oct 2026 14:25:03 -0400 Subject: [PATCH 2/4] Sample the ready-queue depth from the 5s sampler instead of counting per operation Every submit, claim, completion, cancel and promotion ran SELECT count(*) over the jobs table to refresh a gauge whose only reader is the runner pool supervisor. At target throughput that is ~140 full scans a second for a number nobody needs exact. - The 5s state sampler already computes the ready depth (status_inputs' grouped bucket count, published as jobs.ready). It now also stores it into the node-local queue_depth atomic the co-hosted pool scales off, and mirrors it into pool_status. - Removed queue_depth from SubmitOutcome, ClaimedJob, CompleteOutcome, CancelOutcome, PromoteOutcome, EnvironmentApprovalOutcome and the azdo claim variant, plus JobSettled.queue_len (write-only). - settle_job and the cancel paths use cheap EXISTS checks for the queue_nonempty wake signal instead of the count. - Deleted the now-dead pg queue_depth()/queue_depth_on() and the uncalled pg queue_gauges helper; lite's queue_gauges no longer counts. - The per-operation ready-front labels read stays: it is one indexed LIMIT 1 row, not a full scan. --- crates/preloop-runner-server/src/bootstrap.rs | 37 ++++++----- crates/preloop-runner-server/src/broker.rs | 10 --- .../src/control/backend.rs | 2 - .../src/control/lite/dispatch.rs | 21 +++---- .../src/control/lite/fork_gate.rs | 1 - .../src/control/lite/lifecycle.rs | 1 - .../src/control/lite/poll.rs | 2 - .../src/control/lite/submit.rs | 4 -- .../src/control/pg/dispatch.rs | 63 +++++++------------ .../src/control/pg/lifecycle.rs | 26 -------- .../src/control/tests.rs | 12 ++-- .../src/control/types.rs | 24 ++----- .../src/distributed_task.rs | 23 ++----- crates/preloop-runner-server/src/runs.rs | 22 +------ .../src/runtime_scheduling.rs | 7 +-- crates/preloop-runner-server/src/state.rs | 6 +- 16 files changed, 74 insertions(+), 187 deletions(-) diff --git a/crates/preloop-runner-server/src/bootstrap.rs b/crates/preloop-runner-server/src/bootstrap.rs index 529f222a9..f99101ba2 100644 --- a/crates/preloop-runner-server/src/bootstrap.rs +++ b/crates/preloop-runner-server/src/bootstrap.rs @@ -20,8 +20,9 @@ pub struct ServerConfig { pub record_flows: Option, /// TLS mode (default: no TLS). pub tls: TlsMode, - /// Shared counter published with the number of jobs still queued after - /// each claim. Supply one to let a co-hosted runner pool scale to demand. + /// Shared counter with the number of jobs still queued, refreshed by the + /// 5s state sampler. Supply one to let a co-hosted runner pool scale to + /// demand. pub queue_depth: Option>, /// Shared list, refreshed after each claim, of the `runs-on` labels of /// the job at the front of the dispatch queue. Supply one to let a @@ -595,11 +596,8 @@ pub async fn reap_once(shared: &Arc) { Ok(outcome) => { // Post-commit: refresh the node-local mirrors the runner // supervisor and the pool read, then wake them if the sweep - // changed what is schedulable. - shared - .state - .queue_depth - .store(outcome.queue_depth, std::sync::atomic::Ordering::Release); + // changed what is schedulable. The ready-queue depth itself now + // comes from the 5s sampler snapshot. if let Ok(mut guard) = shared.state.next_job_runs_on.write() { *guard = outcome.next_runs_on; } @@ -1272,6 +1270,15 @@ async fn run_state_sampler( // has a single exposition source (the SDK renderer). { let s = shared.state.status_snapshot.read(); + // The co-hosted runner pool scales off this atomic; it + // used to be refreshed by every submit/claim/complete + // (a full count(*) each). The sampler's grouped count + // is the same number at a fixed 5s cadence. + shared + .state + .queue_depth + .store(s.jobs.ready as usize, std::sync::atomic::Ordering::Release); + shared.state.pool_status.set_queue_depth(s.jobs.ready); shared.state.observability.metrics().pool.record( s.service.uptime_seconds, s.pool.desired as u64, @@ -1400,19 +1407,11 @@ pub async fn serve(config: ServerConfig) -> anyhow::Result<()> { } let queue = state.backend.queue_stats().await.unwrap_or_default(); if let Some(queue_depth) = config.queue_depth.clone() { + // A co-hosted pool shares this atomic to scale to demand. The state + // sampler refreshes it every 5s (immediate first tick), and the + // constructor already seeded it from the recovered store — no + // per-operation re-arm needed. state.queue_depth = queue_depth; - // The pool shares this same atomic and only forks a runner while it - // is non-zero; a freshly restarted server has no runners yet, so - // nothing will refresh it from a broker poll. Re-arm it with the - // ready-queue size recovered from the store, or every job queued - // before the restart sits forever with the pool asleep. - let queue_len = queue.ready; - state - .queue_depth - .store(queue_len, std::sync::atomic::Ordering::Release); - state - .pool_status - .set_queue_depth(state.queue_depth.load(std::sync::atomic::Ordering::Acquire) as u32); } if let Some(next_job_runs_on) = config.next_job_runs_on.clone() { state.next_job_runs_on = next_job_runs_on; diff --git a/crates/preloop-runner-server/src/broker.rs b/crates/preloop-runner-server/src/broker.rs index 5ee6919b2..0ae1064ff 100644 --- a/crates/preloop-runner-server/src/broker.rs +++ b/crates/preloop-runner-server/src/broker.rs @@ -350,13 +350,8 @@ pub async fn next_message_broker_ref( queued, request, runner_id, - queue_depth, next_runs_on, } = *claimed; - shared - .state - .queue_depth - .store(queue_depth, std::sync::atomic::Ordering::Release); *shared.state.next_job_runs_on.write().unwrap() = next_runs_on; record_claim_queue_wait(&shared, &Some(queued.clone())); let run_id = queued.run_id; @@ -746,13 +741,8 @@ pub async fn next_message_broker_ref_root( queued, request, runner_id, - queue_depth, next_runs_on, } = *claimed; - shared - .state - .queue_depth - .store(queue_depth, std::sync::atomic::Ordering::Release); *shared.state.next_job_runs_on.write().unwrap() = next_runs_on; record_claim_queue_wait(&shared, &Some(queued)); Some(broker_job_ref_root(&request, runner_id)) diff --git a/crates/preloop-runner-server/src/control/backend.rs b/crates/preloop-runner-server/src/control/backend.rs index 5afd257a1..adfbca130 100644 --- a/crates/preloop-runner-server/src/control/backend.rs +++ b/crates/preloop-runner-server/src/control/backend.rs @@ -43,8 +43,6 @@ pub(crate) struct PromoteOutcome { pub(crate) promoted: usize, /// Jobs failed closed by a denied or expired environment gate. pub(crate) failed: usize, - /// Global ready-queue depth after the pass. - pub(crate) queue_depth: usize, /// `runs-on` labels of the ready-queue front after the pass. pub(crate) next_runs_on: Vec, } diff --git a/crates/preloop-runner-server/src/control/lite/dispatch.rs b/crates/preloop-runner-server/src/control/lite/dispatch.rs index ac98c8771..ee7edb069 100644 --- a/crates/preloop-runner-server/src/control/lite/dispatch.rs +++ b/crates/preloop-runner-server/src/control/lite/dispatch.rs @@ -283,8 +283,10 @@ fn complete_node( /// Ready-queue gauges a completion/cancellation reports back: depth, /// front labels, and whether any work (or pending cancellation) remains. -fn queue_gauges(tx: &Transaction<'_>) -> Result<(usize, Vec, bool), ControlError> { - let depth = jobs::ready_count(tx)?; +/// Queue gauges for command outcomes: the ready-front labels plus whether +/// any work remains. The exact ready count is intentionally not computed +/// here — the runner supervisor reads it from the 5s sampler snapshot. +fn queue_gauges(tx: &Transaction<'_>) -> Result<(Vec, bool), ControlError> { let next = jobs::next_ready_labels(tx)?; let nonempty = jobs::queue_nonempty(tx)? || tx @@ -295,7 +297,7 @@ fn queue_gauges(tx: &Transaction<'_>) -> Result<(usize, Vec, bool), Cont .map_err(db)? .query_row([], |row| row.get::<_, bool>(0)) .map_err(db)?; - Ok((depth, next, nonempty)) + Ok((next, nonempty)) } impl LiteBackend { @@ -340,7 +342,7 @@ impl LiteBackend { } let mut scheduling = crate::runtime_scheduling::SchedulingOutcome::default(); promote::promote_run(tx, self, run_id, &mut scheduling)?; - let (queue_depth, _next_runs_on, queue_nonempty) = queue_gauges(tx)?; + let (_next_runs_on, queue_nonempty) = queue_gauges(tx)?; let record = jobs::run_record(tx, run_id)? .ok_or_else(|| ControlError::NotFound(format!("run {run_id}")))?; Ok(CompleteOutcome { @@ -351,7 +353,6 @@ impl LiteBackend { scheduling, live_log_key: applied.live_log_key, queue_nonempty, - queue_depth, replayed: applied.replayed, }) }) @@ -468,8 +469,7 @@ impl LiteBackend { } let mut scheduling = crate::runtime_scheduling::SchedulingOutcome::default(); promote::promote_run(tx, self, run_id, &mut scheduling)?; - let (queue_len, next_runs_on, _pending) = queue_gauges(tx)?; - let queue_nonempty = queue_len > 0; + let (next_runs_on, queue_nonempty) = queue_gauges(tx)?; Ok(SettleJobOutcome::Settled(Box::new(JobSettled { effective_status: applied.effective_status, cancelled_siblings: applied.cancelled_siblings, @@ -477,7 +477,6 @@ impl LiteBackend { queue_nonempty, newly_terminal_success: applied.newly_terminal_success, live_log_key: applied.live_log_key, - queue_len, next_runs_on, }))) }) @@ -517,7 +516,7 @@ impl LiteBackend { .map_err(db)?; rows.collect::, _>>().map_err(db)? }; - let (queue_depth, next_runs_on, queue_nonempty) = queue_gauges(tx)?; + let (next_runs_on, queue_nonempty) = queue_gauges(tx)?; let record = jobs::run_record(tx, run_id)?; Ok(CancelOutcome { cancellations, @@ -525,7 +524,6 @@ impl LiteBackend { queue_nonempty, record, cancelled_jobs, - queue_depth, next_runs_on, }) }) @@ -556,7 +554,7 @@ impl LiteBackend { .map_err(db)?; rows.collect::, _>>().map_err(db)? }; - let (queue_depth, next_runs_on, queue_nonempty) = queue_gauges(tx)?; + let (next_runs_on, queue_nonempty) = queue_gauges(tx)?; let record = jobs::run_record(tx, run_id)?; Ok(CancelOutcome { cancellations, @@ -564,7 +562,6 @@ impl LiteBackend { queue_nonempty, record, cancelled_jobs, - queue_depth, next_runs_on, }) }) diff --git a/crates/preloop-runner-server/src/control/lite/fork_gate.rs b/crates/preloop-runner-server/src/control/lite/fork_gate.rs index eee4c336b..97e1ad4d8 100644 --- a/crates/preloop-runner-server/src/control/lite/fork_gate.rs +++ b/crates/preloop-runner-server/src/control/lite/fork_gate.rs @@ -69,7 +69,6 @@ fn promote_ready_jobs_tx( outcome.promoted += scheduling.promoted; outcome.failed += scheduling.failed.len(); } - outcome.queue_depth = jobs::ready_count(tx)?; outcome.next_runs_on = jobs::next_ready_labels(tx)?; Ok(outcome) } diff --git a/crates/preloop-runner-server/src/control/lite/lifecycle.rs b/crates/preloop-runner-server/src/control/lite/lifecycle.rs index bca647208..5c7e91621 100644 --- a/crates/preloop-runner-server/src/control/lite/lifecycle.rs +++ b/crates/preloop-runner-server/src/control/lite/lifecycle.rs @@ -1442,7 +1442,6 @@ impl LiteBackend { let promote = self.promote_ready_jobs(Some(run_id), &rules).await?; Ok(EnvironmentApprovalOutcome { result, - queue_depth: promote.queue_depth, next_runs_on: promote.next_runs_on, promoted: promote.promoted, }) diff --git a/crates/preloop-runner-server/src/control/lite/poll.rs b/crates/preloop-runner-server/src/control/lite/poll.rs index a051e2303..681ff7788 100644 --- a/crates/preloop-runner-server/src/control/lite/poll.rs +++ b/crates/preloop-runner-server/src/control/lite/poll.rs @@ -583,7 +583,6 @@ impl LiteBackend { queued, request, runner_id, - queue_depth: jobs::ready_count(tx)?, next_runs_on: jobs::next_ready_labels(tx)?, }))) }) @@ -702,7 +701,6 @@ impl LiteBackend { message, run_id, job_id, - queue_depth: jobs::ready_count(tx)?, next_runs_on: jobs::next_ready_labels(tx)?, }) }) diff --git a/crates/preloop-runner-server/src/control/lite/submit.rs b/crates/preloop-runner-server/src/control/lite/submit.rs index 25de1e4c2..ac924a9e8 100644 --- a/crates/preloop-runner-server/src/control/lite/submit.rs +++ b/crates/preloop-runner-server/src/control/lite/submit.rs @@ -199,7 +199,6 @@ fn submit_run_tx( held: false, rejected: None, existing: Some(Box::new(existing)), - queue_depth: jobs::ready_count(tx)?, next_runs_on: jobs::next_ready_labels(tx)?, events: Vec::new(), }); @@ -221,7 +220,6 @@ fn submit_run_tx( held: false, rejected: Some(ExecutionStatus::Failure), existing: None, - queue_depth: jobs::ready_count(tx)?, next_runs_on: jobs::next_ready_labels(tx)?, events: Vec::new(), }); @@ -434,7 +432,6 @@ fn submit_run_tx( held: false, rejected: Some(ExecutionStatus::Cancelled), existing: None, - queue_depth: jobs::ready_count(tx)?, next_runs_on: jobs::next_ready_labels(tx)?, events: Vec::new(), }); @@ -830,7 +827,6 @@ fn submit_run_tx( held, rejected: None, existing: None, - queue_depth: jobs::ready_count(tx)?, next_runs_on: jobs::next_ready_labels(tx)?, events, }) diff --git a/crates/preloop-runner-server/src/control/pg/dispatch.rs b/crates/preloop-runner-server/src/control/pg/dispatch.rs index 36ef310e3..4fbf319d0 100644 --- a/crates/preloop-runner-server/src/control/pg/dispatch.rs +++ b/crates/preloop-runner-server/src/control/pg/dispatch.rs @@ -2436,7 +2436,7 @@ impl PgBackend { return Err(ControlError::NotFound(format!("run {run_id}"))); } let cancellations = cancel_run_tx(self, &tx, run_id, reason.as_deref()).await?; - let (cancelled_jobs, queue_depth, next_runs_on, pending_cancels) = + let (cancelled_jobs, queue_nonempty, next_runs_on, pending_cancels) = cancel_outcome_gauges(&tx, run_id).await?; let record = self .load_graph(&tx, run_id) @@ -2446,10 +2446,9 @@ impl PgBackend { Ok(CancelOutcome { cancellations, run_status: record.as_ref().map(|record| record.status), - queue_nonempty: queue_depth > 0 || pending_cancels, + queue_nonempty: queue_nonempty || pending_cancels, record, cancelled_jobs, - queue_depth, next_runs_on, }) } @@ -2466,7 +2465,7 @@ impl PgBackend { let tx = client.transaction().await.map_err(db)?; PgBackend::lock_run(&tx, run_id).await?; let cancellations = cancel_job_tx(self, &tx, run_id, job_id, None).await?; - let (cancelled_jobs, queue_depth, next_runs_on, pending_cancels) = + let (cancelled_jobs, queue_nonempty, next_runs_on, pending_cancels) = cancel_outcome_gauges(&tx, run_id).await?; let record = self .load_graph(&tx, run_id) @@ -2476,10 +2475,9 @@ impl PgBackend { Ok(CancelOutcome { cancellations, run_status: record.as_ref().map(|record| record.status), - queue_nonempty: queue_depth > 0 || pending_cancels, + queue_nonempty: queue_nonempty || pending_cancels, record, cancelled_jobs, - queue_depth, next_runs_on, }) } @@ -2491,7 +2489,7 @@ impl PgBackend { async fn cancel_outcome_gauges( tx: &Transaction<'_>, run_id: RunId, -) -> Result<(Vec, usize, Vec, bool), ControlError> { +) -> Result<(Vec, bool, Vec, bool), ControlError> { let cancelled_jobs = tx .query( "SELECT job_id FROM jobs WHERE run_id=$1::text::uuid \ @@ -2503,12 +2501,14 @@ async fn cancel_outcome_gauges( .iter() .map(|row| JobId(row.get::<_, String>(0))) .collect(); - let queue_depth: usize = tx - .query_one("SELECT count(*) FROM jobs WHERE queue_state='ready'", &[]) + let queue_nonempty: bool = tx + .query_one( + "SELECT EXISTS (SELECT 1 FROM jobs WHERE queue_state='ready')", + &[], + ) .await .map_err(db)? - .get::<_, i64>(0) - .max(0) as usize; + .get(0); let next_runs_on: Vec = tx .query_opt( "SELECT runs_on::text FROM jobs WHERE queue_state='ready' \ @@ -2528,7 +2528,12 @@ async fn cancel_outcome_gauges( .await .map_err(db)? .get(0); - Ok((cancelled_jobs, queue_depth, next_runs_on, pending_cancels)) + Ok(( + cancelled_jobs, + queue_nonempty, + next_runs_on, + pending_cancels, + )) } // ───────────────────────────────────────────────────────────────────────── @@ -3965,7 +3970,6 @@ impl PgBackend { held: false, rejected: Some(ExecutionStatus::Failure), existing: None, - queue_depth: self.queue_depth().await?, next_runs_on: Vec::new(), events: Vec::new(), }); @@ -4249,7 +4253,6 @@ impl PgBackend { held: false, rejected: Some(ExecutionStatus::Cancelled), existing: None, - queue_depth: self.queue_depth().await?, next_runs_on: self.ready_front_labels().await?, events: Vec::new(), }); @@ -4543,7 +4546,6 @@ impl PgBackend { concluded, held, rejected: None, - queue_depth: self.queue_depth().await?, next_runs_on: self.ready_front_labels().await?, existing: None, events, @@ -4584,21 +4586,6 @@ impl PgBackend { Ok(platforms) } - /// Ready-queue depth (the node-local gauge the runner supervisor reads). - pub(super) async fn queue_depth(&self) -> Result { - Self::queue_depth_on(&*self.reader().await?).await - } - - /// [`Self::queue_depth`] on a caller's connection or open transaction. - pub(super) async fn queue_depth_on(client: &impl GenericClient) -> Result { - let count = client - .query_one("SELECT count(*) FROM jobs WHERE queue_state = 'ready'", &[]) - .await - .map_err(db)? - .get::<_, i64>(0); - Ok(count.max(0) as usize) - } - /// `runs-on` labels of the ready-queue front, for `next_job_runs_on`. pub(super) async fn ready_front_labels(&self) -> Result, ControlError> { Self::ready_front_labels_on(&*self.reader().await?).await @@ -4633,7 +4620,6 @@ impl PgBackend { concluded: Vec::new(), held: existing.status == ExecutionStatus::Pending, rejected: None, - queue_depth: self.queue_depth().await?, next_runs_on: self.ready_front_labels().await?, existing: Some(Box::new(existing)), events: Vec::new(), @@ -5387,7 +5373,6 @@ impl PgBackend { queued, request, runner_id, - queue_depth: self.queue_depth().await?, next_runs_on: self.ready_front_labels().await?, }))) } @@ -5707,7 +5692,6 @@ impl PgBackend { scheduling, live_log_key: applied.live_log_key, queue_nonempty, - queue_depth: self.queue_depth().await?, replayed: applied.replayed, }) } @@ -5794,11 +5778,11 @@ impl PgBackend { } sweep.sweep().await?; let scheduling = std::mem::take(&mut sweep.outcome); - let (queue_len, next_runs_on) = ( - Self::queue_depth_on(&tx).await?, - Self::ready_front_labels_on(&tx).await?, - ); - let queue_nonempty = queue_len > 0; + // Cheap existence check: the handler only needs to know whether to + // wake pollers, not the exact depth (the supervisor reads that from + // the 5s sampler snapshot). + let queue_nonempty = Self::work_pending_on(&tx).await?; + let next_runs_on = Self::ready_front_labels_on(&tx).await?; sweep.flush().await?; tx.commit().await.map_err(db)?; Ok(SettleJobOutcome::Settled(Box::new(JobSettled { @@ -5808,7 +5792,6 @@ impl PgBackend { queue_nonempty, newly_terminal_success: applied.newly_terminal_success, live_log_key: applied.live_log_key, - queue_len, next_runs_on, }))) } @@ -6074,7 +6057,6 @@ impl PgBackend { message, run_id, job_id, - queue_depth: self.queue_depth().await?, next_runs_on: self.ready_front_labels().await?, }) } @@ -6562,7 +6544,6 @@ impl PgBackend { } sweep.flush().await?; } - outcome.queue_depth = Self::queue_depth_on(&tx).await?; outcome.next_runs_on = Self::ready_front_labels_on(&tx).await?; tx.commit().await.map_err(db)?; Ok(outcome) diff --git a/crates/preloop-runner-server/src/control/pg/lifecycle.rs b/crates/preloop-runner-server/src/control/pg/lifecycle.rs index 3fdc8e1e7..2b5bf3652 100644 --- a/crates/preloop-runner-server/src/control/pg/lifecycle.rs +++ b/crates/preloop-runner-server/src/control/pg/lifecycle.rs @@ -42,31 +42,6 @@ async fn token_request_for( .transpose() } -/// `queue_depth`/`next_runs_on` gauges for command outcomes. -async fn queue_gauges( - client: &tokio_postgres::Client, -) -> Result<(usize, Vec), ControlError> { - let rows = client - .query( - "SELECT runs_on::text FROM jobs WHERE queue_state='ready' \ - ORDER BY priority DESC, run_order, job_order LIMIT 64", - &[], - ) - .await - .map_err(db)?; - let depth = client - .query_one("SELECT count(*) FROM jobs WHERE queue_state='ready'", &[]) - .await - .map_err(db)? - .get::<_, i64>(0) as usize; - let front = rows - .first() - .map(|row| from_json::>(row.get::<_, String>(0).as_str())) - .transpose()? - .unwrap_or_default(); - Ok((depth, front)) -} - impl PgBackend { // ── Sessions ───────────────────────────────────────────────────── @@ -1750,7 +1725,6 @@ impl PgBackend { let promote = self.promote_ready_jobs(Some(run_id), &rules).await?; Ok(EnvironmentApprovalOutcome { result, - queue_depth: promote.queue_depth, next_runs_on: promote.next_runs_on, promoted: promote.promoted, }) diff --git a/crates/preloop-runner-server/src/control/tests.rs b/crates/preloop-runner-server/src/control/tests.rs index adabfcb72..69967b772 100644 --- a/crates/preloop-runner-server/src/control/tests.rs +++ b/crates/preloop-runner-server/src/control/tests.rs @@ -1813,7 +1813,6 @@ pub(crate) mod suite { .await .unwrap(); assert_eq!(outcome.promoted, 1, "the released hold admits the job"); - assert_eq!(outcome.queue_depth, 1); assert_eq!( backend.job_queue_state(run_id, &job_id).await.unwrap(), Some(("ready".to_owned(), "queued".to_owned())) @@ -2892,15 +2891,18 @@ pub(crate) mod suite { let run_a = RunId::new(); let mut submit_a = submit_run(run_a, vec![submit_job(run_a, "build", 1)]); submit_a.workflow_concurrency = Some(workflow_concurrency("g", true)); - let outcome_a = backend.submit_run(submit_a).await.unwrap(); - assert_eq!(outcome_a.queue_depth, 1); + backend.submit_run(submit_a).await.unwrap(); + let stats = backend.queue_stats().await.unwrap(); + assert_eq!((stats.ready, stats.claimed), (1, 0)); let run_b = RunId::new(); let mut submit_b = submit_run(run_b, vec![submit_job(run_b, "build", 2)]); submit_b.workflow_concurrency = Some(workflow_concurrency("g", true)); - let outcome_b = backend.submit_run(submit_b).await.unwrap(); + backend.submit_run(submit_b).await.unwrap(); + let stats = backend.queue_stats().await.unwrap(); assert_eq!( - outcome_b.queue_depth, 1, + (stats.ready, stats.claimed), + (1, 0), "cancelled job must not inflate depth" ); let record_a = backend.run_record(run_a).await.unwrap(); diff --git a/crates/preloop-runner-server/src/control/types.rs b/crates/preloop-runner-server/src/control/types.rs index b9fe53709..c7483ab2a 100644 --- a/crates/preloop-runner-server/src/control/types.rs +++ b/crates/preloop-runner-server/src/control/types.rs @@ -322,9 +322,6 @@ pub(crate) struct SubmitOutcome { /// A replayed webhook delivery found an existing run — the handler /// returns that acceptance instead of a new run. pub(crate) existing: Option>, - /// Ready-queue depth after the transition — the handler stores it in - /// the node-local `queue_depth` gauge that wakes the runner supervisor. - pub(crate) queue_depth: usize, /// `runs-on` labels of the next ready job, for `next_job_runs_on`. pub(crate) next_runs_on: Vec, /// Events the submit transaction already wrote to the outbox — the @@ -360,8 +357,6 @@ pub(crate) struct ClaimedJob { pub(crate) queued: QueuedJob, pub(crate) request: TaskAgentJobRequestRecord, pub(crate) runner_id: i64, - /// Queue depth after the claim (for the supervisor atomic). - pub(crate) queue_depth: usize, /// `runs-on` of the new queue front, for golden selection. pub(crate) next_runs_on: Vec, } @@ -412,8 +407,6 @@ pub(crate) struct CompleteOutcome { pub(crate) live_log_key: String, /// Whether the ready queue or cancellation queue is non-empty. pub(crate) queue_nonempty: bool, - /// Queue depth after the transition. - pub(crate) queue_depth: usize, /// Whether the completion was a replay of an already-terminal job. pub(crate) replayed: bool, } @@ -431,8 +424,6 @@ pub(crate) struct CancelOutcome { pub(crate) record: Option, /// Every job of the run whose status is now `Cancelled`, in job-id order. pub(crate) cancelled_jobs: Vec, - /// Global ready-queue depth after the transition. - pub(crate) queue_depth: usize, /// `runs-on` labels of the ready-queue front after the transition. pub(crate) next_runs_on: Vec, } @@ -1058,8 +1049,6 @@ pub(crate) enum EnvironmentApprovalResult { #[derive(Debug)] pub(crate) struct EnvironmentApprovalOutcome { pub(crate) result: EnvironmentApprovalResult, - /// Ready-queue depth after the command. - pub(crate) queue_depth: usize, /// `runs-on` labels of the ready queue's front job after the command. pub(crate) next_runs_on: Vec, /// Jobs released to ready by the post-approval promotion sweep. @@ -1170,9 +1159,7 @@ pub(crate) struct JobSettled { pub(crate) newly_terminal_success: bool, /// Live-log key of the attempt (agent job id, else logical job id). pub(crate) live_log_key: String, - /// Ready-queue depth after the transition and the `runs-on` of its - /// front job. - pub(crate) queue_len: usize, + /// `runs-on` of the ready-queue front job after the transition. pub(crate) next_runs_on: Vec, } @@ -1234,16 +1221,13 @@ pub(crate) enum AzdoPollOutcome { /// a `JobCancellation` session message (HTTP 200). Cancel(SessionMessage), /// A job was claimed: the new `PipelineAgentJobRequest` session message - /// (HTTP 202) and the job it assigns, plus the ready-queue snapshot the - /// committed claim produced so the handler can refresh the supervisor - /// gauges (`queue_depth`/`next_job_runs_on`) the same way the broker - /// claim does. + /// (HTTP 202) and the job it assigns, plus the ready-queue front labels + /// the committed claim produced so the handler can refresh + /// `next_job_runs_on` the same way the broker claim does. Claimed { message: SessionMessage, run_id: RunId, job_id: JobId, - /// Queue depth after the claim (for the supervisor atomic). - queue_depth: usize, /// `runs-on` of the new queue front, for `next_job_runs_on`. next_runs_on: Vec, }, diff --git a/crates/preloop-runner-server/src/distributed_task.rs b/crates/preloop-runner-server/src/distributed_task.rs index 519d7fb65..11cf9cfc1 100644 --- a/crates/preloop-runner-server/src/distributed_task.rs +++ b/crates/preloop-runner-server/src/distributed_task.rs @@ -54,16 +54,11 @@ pub async fn next_message( message, run_id, job_id, - queue_depth, next_runs_on, } => { - // Refresh the on-demand pool gauges from the committed - // claim, exactly like the broker claim path: the ready - // queue just shrank. - shared - .state - .queue_depth - .store(queue_depth, std::sync::atomic::Ordering::Release); + // Refresh the next-job labels from the committed claim, + // exactly like the broker claim path. The ready-queue depth + // itself now comes from the 5s sampler snapshot. *shared.state.next_job_runs_on.write().unwrap() = next_runs_on; github::report_check_run_in_progress(&shared, run_id, &job_id).await; shared @@ -772,7 +767,6 @@ pub(crate) async fn complete_job_settling( queue_nonempty, newly_terminal_success, live_log_key, - queue_len: tx_queue_len, next_runs_on: tx_next_labels, } = settled; @@ -783,14 +777,9 @@ pub(crate) async fn complete_job_settling( crate::live_logs::close_live_log(&mut inner, &live_log_key); inner.dap_ports.remove(&completion.run_id); } - // Refresh the on-demand pool wake atomic and the next-job labels from the - // committed scheduling state. - let (queue_len, next_labels) = (tx_queue_len, tx_next_labels); - shared - .state - .queue_depth - .store(queue_len, std::sync::atomic::Ordering::Release); - *shared.state.next_job_runs_on.write().unwrap() = next_labels; + // Refresh the next-job labels from the committed scheduling state. The + // ready-queue depth itself now comes from the 5s sampler snapshot. + *shared.state.next_job_runs_on.write().unwrap() = tx_next_labels; // Any reusable-caller or dynamic-matrix node the sweep above unblocked was // deferred rather than expanded under the lock. Build those subtrees now diff --git a/crates/preloop-runner-server/src/runs.rs b/crates/preloop-runner-server/src/runs.rs index 3bbda0103..9c733206a 100644 --- a/crates/preloop-runner-server/src/runs.rs +++ b/crates/preloop-runner-server/src/runs.rs @@ -1914,12 +1914,9 @@ async fn submit_run_inner_with_webhook_delivery_unreserved( let tx_ms = t_tx.elapsed().as_secs_f64() * 1000.0; let t_tail = std::time::Instant::now(); - // Post-commit: refresh the node-local gauges that wake the runner - // supervisor, then fan out the events the outcome carries. - shared - .state - .queue_depth - .store(outcome.queue_depth, std::sync::atomic::Ordering::Release); + // Post-commit: refresh the node-local labels the runner supervisor + // reads, then fan out the events the outcome carries. The ready-queue + // depth itself now comes from the 5s sampler snapshot. *shared.state.next_job_runs_on.write().unwrap() = outcome.next_runs_on; if let Some(existing) = outcome.existing { @@ -3415,7 +3412,6 @@ pub async fn cancel_run( cancellations: cancellation_count, record, cancelled_jobs, - queue_depth, next_runs_on, .. } = shared @@ -3429,10 +3425,6 @@ pub async fn cancel_run( "run not found".to_owned(), )) })?; - shared - .state - .queue_depth - .store(queue_depth, std::sync::atomic::Ordering::Release); *shared.state.next_job_runs_on.write().unwrap() = next_runs_on; if cancellation_count > 0 { shared.state.message_notify.notify_waiters(); @@ -3508,10 +3500,6 @@ pub async fn approve_job( }) .await .map_err(ApiError::from)?; - shared - .state - .queue_depth - .store(outcome.queue_depth, std::sync::atomic::Ordering::Release); *shared.state.next_job_runs_on.write().unwrap() = outcome.next_runs_on; if outcome.promoted > 0 { shared.state.message_notify.notify_waiters(); @@ -3703,10 +3691,6 @@ pub async fn approve_fork( .promote_ready_jobs(Some(run_id), &shared.state.environment_rules) .await .map_err(ApiError::from)?; - shared - .state - .queue_depth - .store(outcome.queue_depth, std::sync::atomic::Ordering::Release); *shared.state.next_job_runs_on.write().unwrap() = outcome.next_runs_on; if outcome.promoted > 0 { shared.state.message_notify.notify_waiters(); diff --git a/crates/preloop-runner-server/src/runtime_scheduling.rs b/crates/preloop-runner-server/src/runtime_scheduling.rs index b26fbc276..a61848685 100644 --- a/crates/preloop-runner-server/src/runtime_scheduling.rs +++ b/crates/preloop-runner-server/src/runtime_scheduling.rs @@ -730,12 +730,9 @@ pub async fn drain_expansions(shared: &Arc) -> SchedulingOutcome { } } - // Refresh node-local scheduling mirrors from committed state. + // Refresh the node-local next-job labels from committed state. The + // ready-queue depth itself now comes from the 5s sampler snapshot. if let Ok(stats) = shared.state.backend.queue_stats().await { - shared - .state - .queue_depth - .store(stats.ready, std::sync::atomic::Ordering::Release); *shared.state.next_job_runs_on.write().unwrap() = stats.next_runs_on; } } diff --git a/crates/preloop-runner-server/src/state.rs b/crates/preloop-runner-server/src/state.rs index 5e43d2226..253bb284c 100644 --- a/crates/preloop-runner-server/src/state.rs +++ b/crates/preloop-runner-server/src/state.rs @@ -734,9 +734,9 @@ pub struct AppState { pub pool_status: Arc, /// When this AppState was created (for uptime). pub started_at: std::time::Instant, - /// Jobs accepted and still waiting for a runner, refreshed whenever one - /// is claimed. A supervising runner pool reads it to decide whether the - /// work already queued outruns the runners it has left. + /// Jobs accepted and still waiting for a runner, refreshed by the 5s + /// state sampler. A supervising runner pool reads it to decide whether + /// the work already queued outruns the runners it has left. pub queue_depth: Arc, /// Plan 000 step 4 probe — broker job-lifecycle calls (`renewjob`, /// `completejob`) that authenticated with the bare runner *listen* token From 1cf035597a6d3f0208e96b20fb2bc1bccd6d3052 Mon Sep 17 00:00:00 2001 From: Bill Njoroge Date: Sat, 3 Oct 2026 15:59:56 -0400 Subject: [PATCH 3/4] Wake the state sampler early on submit bursts (debounced) The 5s sampler now feeds the pool's queue-depth gauge, but a burst could sit up to 5s before the pool noticed - 10x the 500ms snapshot-fork provision time. The sampler loop now also wakes on message_notify (which the submit path already shouts into): Notify coalesces a flurry into one wakeup, and a 1s floor caps a sustained storm at one extra grouped count per second. Same epoll-style rhythm-plus-urgency pattern as the broker long-poll. --- crates/preloop-runner-server/src/bootstrap.rs | 113 ++++++++++-------- 1 file changed, 65 insertions(+), 48 deletions(-) diff --git a/crates/preloop-runner-server/src/bootstrap.rs b/crates/preloop-runner-server/src/bootstrap.rs index f99101ba2..f6b88c667 100644 --- a/crates/preloop-runner-server/src/bootstrap.rs +++ b/crates/preloop-runner-server/src/bootstrap.rs @@ -1261,56 +1261,73 @@ async fn run_state_sampler( let mut interval = tokio::time::interval(Duration::from_secs(5)); // Immediate sample then every 5s. interval.tick().await; + // When the last tick ran: the burst-reactive early tick below must not + // turn a submit storm back into per-operation counting. + let mut last_tick = std::time::Instant::now(); loop { + let mut early = false; + let mut shutdown = false; tokio::select! { - _ = interval.tick() => { - heartbeat.beat("state_sampler"); - publish_snapshot(&shared, &store_backend, false).await; - // Record pool/queue gauges into OTel instruments so `/metrics` - // has a single exposition source (the SDK renderer). - { - let s = shared.state.status_snapshot.read(); - // The co-hosted runner pool scales off this atomic; it - // used to be refreshed by every submit/claim/complete - // (a full count(*) each). The sampler's grouped count - // is the same number at a fixed 5s cadence. - shared - .state - .queue_depth - .store(s.jobs.ready as usize, std::sync::atomic::Ordering::Release); - shared.state.pool_status.set_queue_depth(s.jobs.ready); - shared.state.observability.metrics().pool.record( - s.service.uptime_seconds, - s.pool.desired as u64, - s.pool.preparing, - s.pool.idle as u64, - s.pool.busy as u64, - s.jobs.ready as u64, - s.jobs.claimable as u64, - s.jobs.unclaimable as u64, - s.jobs.dependency_blocked as u64, - ); - // Host memory reality on the same cadence: a /proc scan - // is microseconds next to everything else on this tick. - // Recorded even when nothing else changed so OOM - // proximity is a continuous signal, not a sampled one. - shared - .state - .observability - .metrics() - .host - .record(&preloop_observability::vm_telemetry::sample_host()); - } - } - _ = shared.shutdown.cancelled() => { - // Publish one last snapshot with the shutdown flag set so - // /api/v1/status reports `overall: shutting_down` and - // `shutdown_requested: true` while /healthz//readyz already - // 503 — without this the flag would only land on the next 5s - // tick that never comes. - heartbeat.beat("state_sampler"); - publish_snapshot(&shared, &store_backend, true).await; - break; + _ = interval.tick() => {} + // A submit/claim/complete/cancel wakes the sampler early so the + // pool sees a burst in ~ms instead of at the next 5s tick. The + // Notify coalesces a flurry into one wakeup, and the 1s floor + // below caps a sustained storm at one extra count per second. + _ = shared.state.message_notify.notified() => { early = true; } + _ = shared.shutdown.cancelled() => { shutdown = true; } + } + if shutdown { + // Publish one last snapshot with the shutdown flag set so + // /api/v1/status reports `overall: shutting_down` and + // `shutdown_requested: true` while /healthz//readyz already + // 503 — without this the flag would only land on the next 5s + // tick that never comes. + heartbeat.beat("state_sampler"); + publish_snapshot(&shared, &store_backend, true).await; + break; + } + if early && last_tick.elapsed() < Duration::from_secs(1) { + continue; + } + last_tick = std::time::Instant::now(); + { + heartbeat.beat("state_sampler"); + publish_snapshot(&shared, &store_backend, false).await; + // Record pool/queue gauges into OTel instruments so `/metrics` + // has a single exposition source (the SDK renderer). + { + let s = shared.state.status_snapshot.read(); + // The co-hosted runner pool scales off this atomic; it + // used to be refreshed by every submit/claim/complete + // (a full count(*) each). The sampler's grouped count + // is the same number at a fixed cadence, plus a + // burst-reactive early tick above. + shared + .state + .queue_depth + .store(s.jobs.ready as usize, std::sync::atomic::Ordering::Release); + shared.state.pool_status.set_queue_depth(s.jobs.ready); + shared.state.observability.metrics().pool.record( + s.service.uptime_seconds, + s.pool.desired as u64, + s.pool.preparing, + s.pool.idle as u64, + s.pool.busy as u64, + s.jobs.ready as u64, + s.jobs.claimable as u64, + s.jobs.unclaimable as u64, + s.jobs.dependency_blocked as u64, + ); + // Host memory reality on the same cadence: a /proc scan + // is microseconds next to everything else on this tick. + // Recorded even when nothing else changed so OOM + // proximity is a continuous signal, not a sampled one. + shared + .state + .observability + .metrics() + .host + .record(&preloop_observability::vm_telemetry::sample_host()); } } } From bb5e9925f41bc67cefbff66b915f3ce832318d0b Mon Sep 17 00:00:00 2001 From: Bnjoroge1 Date: Tue, 6 Oct 2026 12:13:45 -0400 Subject: [PATCH 4/4] fix: isolate the sampler's wake channel; close long-poll lost-wake race Review follow-ups on the poll-write refactor: - Give the state sampler its own Notify (sampler_notify). wake_waiters fires per-job notify_one permits on message_notify; a sampler parked on the same channel could steal them and strand a runner until its window ended. Every producer now wakes both channels; the LISTEN relay too. - Register the long-poll Notified before probing (enable()), in all three loops: broker ref, broker root, disttask. A notify_waiters landing between probe and registration stored no permit and was lost. - Sampler debounce naps off the 1s floor instead of dropping the event, and MissedTickBehavior::Delay prevents a burst publish after a nap. - Clamp waitSeconds to PRELOOP_MAX_POLL_WINDOW_SECS (default 60s) and floor PRELOOP_RUNNER_LIVENESS_TIMEOUT_SECS at 3x that window, so a client-chosen window can never park a runner past its own reaping. - pg probe_poll: document that claim-enabling reads must never route to a replica (they can't see writer commits). - queue_depth tests: drive a real sampler tick via test_sample_state_once (extracted from run_state_sampler) instead of asserting the removed per-operation write. --- crates/preloop-runner-server/src/bootstrap.rs | 116 +++++++++++------- crates/preloop-runner-server/src/broker.rs | 37 ++++-- .../src/control/pg/dispatch.rs | 7 +- .../src/distributed_task.rs | 18 ++- .../src/runner_lifecycle.rs | 4 + crates/preloop-runner-server/src/runs.rs | 7 ++ crates/preloop-runner-server/src/state.rs | 101 ++++++++++++++- crates/preloop-runner-server/tests/logs.rs | 8 ++ 8 files changed, 232 insertions(+), 66 deletions(-) diff --git a/crates/preloop-runner-server/src/bootstrap.rs b/crates/preloop-runner-server/src/bootstrap.rs index f6b88c667..614b14fe6 100644 --- a/crates/preloop-runner-server/src/bootstrap.rs +++ b/crates/preloop-runner-server/src/bootstrap.rs @@ -484,6 +484,7 @@ pub async fn reap_once(shared: &Arc) { // Notify if cancellations or starvation failures occurred if cancellation_count > 0 || !starved.is_empty() || !expired_fork_approvals.is_empty() { shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); } // Surface fork-PR runs failed closed by the expired approval window. @@ -603,6 +604,7 @@ pub async fn reap_once(shared: &Arc) { } if outcome.promoted > 0 || outcome.failed > 0 { shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); } } Err(error) => warn!(?error, "environment gate promotion sweep failed"), @@ -1259,6 +1261,9 @@ async fn run_state_sampler( ); heartbeat.beat("state_sampler"); let mut interval = tokio::time::interval(Duration::from_secs(5)); + // A tick missed during a coalescing nap delays rather than bursting: + // publishing twice back-to-back reports the same counters twice. + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); // Immediate sample then every 5s. interval.tick().await; // When the last tick ran: the burst-reactive early tick below must not @@ -1270,10 +1275,11 @@ async fn run_state_sampler( tokio::select! { _ = interval.tick() => {} // A submit/claim/complete/cancel wakes the sampler early so the - // pool sees a burst in ~ms instead of at the next 5s tick. The - // Notify coalesces a flurry into one wakeup, and the 1s floor - // below caps a sustained storm at one extra count per second. - _ = shared.state.message_notify.notified() => { early = true; } + // pool sees a burst in ~ms instead of at the next 5s tick. This + // is the sampler's own channel: runner waiters on + // `message_notify` are woken with per-job `notify_one` permits + // that a parked sampler here would otherwise steal. + _ = shared.state.sampler_notify.notified() => { early = true; } _ = shared.shutdown.cancelled() => { shutdown = true; } } if shutdown { @@ -1286,53 +1292,73 @@ async fn run_state_sampler( publish_snapshot(&shared, &store_backend, true).await; break; } - if early && last_tick.elapsed() < Duration::from_secs(1) { - continue; - } - last_tick = std::time::Instant::now(); + if early + && let Some(nap) = Duration::from_secs(1).checked_sub(last_tick.elapsed()) + && !nap.is_zero() { - heartbeat.beat("state_sampler"); - publish_snapshot(&shared, &store_backend, false).await; - // Record pool/queue gauges into OTel instruments so `/metrics` - // has a single exposition source (the SDK renderer). - { - let s = shared.state.status_snapshot.read(); - // The co-hosted runner pool scales off this atomic; it - // used to be refreshed by every submit/claim/complete - // (a full count(*) each). The sampler's grouped count - // is the same number at a fixed cadence, plus a - // burst-reactive early tick above. - shared - .state - .queue_depth - .store(s.jobs.ready as usize, std::sync::atomic::Ordering::Release); - shared.state.pool_status.set_queue_depth(s.jobs.ready); - shared.state.observability.metrics().pool.record( - s.service.uptime_seconds, - s.pool.desired as u64, - s.pool.preparing, - s.pool.idle as u64, - s.pool.busy as u64, - s.jobs.ready as u64, - s.jobs.claimable as u64, - s.jobs.unclaimable as u64, - s.jobs.dependency_blocked as u64, - ); - // Host memory reality on the same cadence: a /proc scan - // is microseconds next to everything else on this tick. - // Recorded even when nothing else changed so OOM - // proximity is a continuous signal, not a sampled one. - shared - .state - .observability - .metrics() - .host - .record(&preloop_observability::vm_telemetry::sample_host()); + // Inside the 1s floor: don't drop the event — nap off the rest + // of the floor, then sample below. A `continue` here meant a + // change during the floor was invisible until the next 5s tick. + // Wakes during the nap need no flag: the sample that follows + // counts everything already. + tokio::select! { + _ = tokio::time::sleep(nap) => {} + _ = shared.shutdown.cancelled() => { + heartbeat.beat("state_sampler"); + publish_snapshot(&shared, &store_backend, true).await; + break; + } } } + last_tick = std::time::Instant::now(); + heartbeat.beat("state_sampler"); + sample_state_once(&shared, &store_backend).await; } } +/// One sampler tick: publish the snapshot, then refresh the pool/queue +/// gauges and host-memory metrics from it. The loop calls this on every +/// tick/early-wake; tests call it directly (`test_sample_state_once`) since +/// no sampler task runs under `app()`. +pub(crate) async fn sample_state_once( + shared: &Arc, + store_backend: &preloop_observability::status::StoreBackend, +) { + publish_snapshot(shared, store_backend, false).await; + // Record pool/queue gauges into OTel instruments so `/metrics` has a + // single exposition source (the SDK renderer). + let s = shared.state.status_snapshot.read(); + // The co-hosted runner pool scales off this atomic; it used to be + // refreshed by every submit/claim/complete (a full count(*) each). The + // sampler's grouped count is the same number at a fixed cadence, plus a + // burst-reactive early tick in the loop above. + shared + .state + .queue_depth + .store(s.jobs.ready as usize, std::sync::atomic::Ordering::Release); + shared.state.pool_status.set_queue_depth(s.jobs.ready); + shared.state.observability.metrics().pool.record( + s.service.uptime_seconds, + s.pool.desired as u64, + s.pool.preparing, + s.pool.idle as u64, + s.pool.busy as u64, + s.jobs.ready as u64, + s.jobs.claimable as u64, + s.jobs.unclaimable as u64, + s.jobs.dependency_blocked as u64, + ); + // Host memory reality on the same cadence: a /proc scan is microseconds + // next to everything else on this tick. Recorded even when nothing else + // changed so OOM proximity is a continuous signal, not a sampled one. + shared + .state + .observability + .metrics() + .host + .record(&preloop_observability::vm_telemetry::sample_host()); +} + fn is_routine_unix_disconnect(error: &(dyn std::error::Error + 'static)) -> bool { error .downcast_ref::() diff --git a/crates/preloop-runner-server/src/broker.rs b/crates/preloop-runner-server/src/broker.rs index 0ae1064ff..26e6d47f4 100644 --- a/crates/preloop-runner-server/src/broker.rs +++ b/crates/preloop-runner-server/src/broker.rs @@ -288,7 +288,10 @@ pub async fn next_message_broker_ref( let wait_seconds = params .get("waitSeconds") .and_then(|value| value.parse::().ok()) - .unwrap_or(50); + .unwrap_or(50) + // A client-chosen window longer than the liveness floor parks a + // healthy runner past its own reaping. + .min(crate::state::max_poll_window_secs()); let runner_busy = params .get("status") .is_some_and(|status| status.eq_ignore_ascii_case("busy")); @@ -322,6 +325,13 @@ pub async fn next_message_broker_ref( let deadline = tokio::time::Instant::now() + Duration::from_secs(wait_seconds); loop { + // Register the waiter *before* probing: a `notify_waiters` landing + // between the probe and a later `notified()` registration is lost + // forever (it stores no permit), stalling this poll until the + // window ends even though work was queued. + let notified = shared.state.message_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); let outcome = shared .state .backend @@ -388,10 +398,7 @@ pub async fn next_message_broker_ref( }; return Ok((status, Json(body)).into_response()); } - if tokio::time::timeout_at(deadline, shared.state.message_notify.notified()) - .await - .is_err() - { + if tokio::time::timeout_at(deadline, notified).await.is_err() { let status = if runner_busy { StatusCode::ACCEPTED } else { @@ -699,7 +706,10 @@ pub async fn next_message_broker_ref_root( let wait = params .get("waitSeconds") .and_then(|value| value.parse::().ok()) - .unwrap_or(50); + .unwrap_or(50) + // A client-chosen window longer than the liveness floor parks a + // healthy runner past its own reaping. + .min(crate::state::max_poll_window_secs()); // The runner may report completion before its worker process has fully // exited. GitHub keeps polling with status=Busy during that drain window; // never dispatch a successor until the runner reports Online again. @@ -710,6 +720,13 @@ pub async fn next_message_broker_ref_root( let mut deadline = std::time::Instant::now() + Duration::from_secs(wait); loop { + // Register the waiter *before* probing: a `notify_waiters` landing + // between the probe and a later `notified()` registration is lost + // forever (it stores no permit), stalling this poll until the + // window ends even though work was queued. + let notified = shared.state.message_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); // `drained`: Busy poll but nothing is still running — the session's // job has finished and the status=Busy report is stale. let mut drained = false; @@ -770,9 +787,9 @@ pub async fn next_message_broker_ref_root( // One wait per window, like `next_message_broker_ref` and the azdo // path: the poll above is a cheap reader probe now, and // `message_notify` fires on every enqueue, cancellation and - // promotion. A lost wake only delays until the window ends. - let _ = - tokio::time::timeout_at(deadline.into(), shared.state.message_notify.notified()).await; + // promotion. The waiter was registered before the probe, so a wake + // that landed mid-poll is still seen. + let _ = tokio::time::timeout_at(deadline.into(), notified).await; } } @@ -1306,6 +1323,7 @@ async fn fail_unclaimable_request(shared: &Arc, request_id: i64) { // Let a long-polling runner pick up a successor immediately rather than // waiting out its poll window behind a job that will never run. shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); } /// Reconcile job claims whose runner session did not survive a restart. @@ -1687,6 +1705,7 @@ pub async fn broker_complete_job( // Wake long-polling runners so a queued successor job is delivered promptly // after cancel/complete (concurrency release path). shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); Ok(StatusCode::NO_CONTENT) } diff --git a/crates/preloop-runner-server/src/control/pg/dispatch.rs b/crates/preloop-runner-server/src/control/pg/dispatch.rs index 4fbf319d0..2e81e4d4a 100644 --- a/crates/preloop-runner-server/src/control/pg/dispatch.rs +++ b/crates/preloop-runner-server/src/control/pg/dispatch.rs @@ -5231,7 +5231,12 @@ impl PgBackend { &self, poll: PollRequest, ) -> Result { - // Fast path first: idle polls never take a writer. + // Fast path first: idle polls never take a writer. The probe reads + // from this node's own pool — the claim-enabling probe (and every + // state read that can promote to a write) must NEVER route to a + // read replica: a replica's snapshot lags the writer's commits, so + // it could answer "empty" for a claim the writer just made, or + // claimable for one another node already took. Read your writes. let reader = self.reader().await?; if let Some(outcome) = self.probe_poll(&*reader, &poll).await? { return Ok(outcome); diff --git a/crates/preloop-runner-server/src/distributed_task.rs b/crates/preloop-runner-server/src/distributed_task.rs index 11cf9cfc1..ebc3a8b81 100644 --- a/crates/preloop-runner-server/src/distributed_task.rs +++ b/crates/preloop-runner-server/src/distributed_task.rs @@ -16,11 +16,21 @@ pub async fn next_message( let wait_seconds = params .get("waitSeconds") .and_then(|value| value.parse::().ok()) - .unwrap_or(50); + .unwrap_or(50) + // A client-chosen window longer than the liveness floor parks a + // healthy runner past its own reaping. + .min(crate::state::max_poll_window_secs()); // One window per request (see `broker::next_message_broker_ref`). let deadline = tokio::time::Instant::now() + Duration::from_secs(wait_seconds); loop { + // Register the waiter *before* probing: a `notify_waiters` landing + // between the probe and a later `notified()` registration is lost + // forever (it stores no permit), stalling this poll until the + // window ends even though work was queued. + let notified = shared.state.message_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); let outcome = shared .state .backend @@ -77,10 +87,7 @@ pub async fn next_message( if wait_seconds == 0 { return Ok((StatusCode::OK, Json(None))); } - if tokio::time::timeout_at(deadline, shared.state.message_notify.notified()) - .await - .is_err() - { + if tokio::time::timeout_at(deadline, notified).await.is_err() { return Ok((StatusCode::OK, Json(None))); } continue; @@ -852,6 +859,7 @@ pub(crate) async fn complete_job_settling( if scheduling.promoted > 0 || !cancelled_siblings.is_empty() || queue_nonempty { crate::state::wake_waiters( &shared.state.message_notify, + &shared.state.sampler_notify, scheduling.promoted.max(usize::from(queue_nonempty)), !cancelled_siblings.is_empty(), ); diff --git a/crates/preloop-runner-server/src/runner_lifecycle.rs b/crates/preloop-runner-server/src/runner_lifecycle.rs index c453f5c7e..468c7ff27 100644 --- a/crates/preloop-runner-server/src/runner_lifecycle.rs +++ b/crates/preloop-runner-server/src/runner_lifecycle.rs @@ -372,6 +372,7 @@ pub async fn delete_session( tracing::warn!(?error, "failed to persist deleted runner session"); } shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); Ok(StatusCode::NO_CONTENT) } /// DELETE /runner/server/_apis/distributedtask/pools/:pool_id/agents/:agent_id @@ -434,6 +435,7 @@ pub async fn purge_runner_identity_guarded( } } shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); Ok(()) } @@ -488,6 +490,7 @@ async fn purge_runner_identity_with_phantom_check( } } shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); true } @@ -523,6 +526,7 @@ pub async fn purge_restored_ephemeral_runners(shared: &Arc) { "purged restored ephemeral runner identities" ); shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); } /// DELETE /runner/server/_apis/distributedtask/pools/:pool_id/sessions (no session_id) diff --git a/crates/preloop-runner-server/src/runs.rs b/crates/preloop-runner-server/src/runs.rs index 9c733206a..1136e5e66 100644 --- a/crates/preloop-runner-server/src/runs.rs +++ b/crates/preloop-runner-server/src/runs.rs @@ -764,6 +764,7 @@ async fn release_webhook_run_reservation(shared: &Arc, key: (String inner.webhook_run_reservations.remove(&key); drop(inner); shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); } /// The host (and optional port) of the configured forge, for `git://`-schemed @@ -1951,6 +1952,7 @@ async fn submit_run_inner_with_webhook_delivery_unreserved( if outcome.queued_jobs > 0 || expansion.promoted > 0 { crate::state::wake_waiters( &shared.state.message_notify, + &shared.state.sampler_notify, outcome.queued_jobs + expansion.promoted, false, ); @@ -3428,6 +3430,7 @@ pub async fn cancel_run( *shared.state.next_job_runs_on.write().unwrap() = next_runs_on; if cancellation_count > 0 { shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); } shared .state @@ -3503,6 +3506,7 @@ pub async fn approve_job( *shared.state.next_job_runs_on.write().unwrap() = outcome.next_runs_on; if outcome.promoted > 0 { shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); } match outcome.result { crate::control::types::EnvironmentApprovalResult::AlreadyTerminal => { @@ -3513,6 +3517,7 @@ pub async fn approve_job( )), crate::control::types::EnvironmentApprovalResult::Expired => { shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); Err(ApiError::conflict( "approval window expired; the job was failed closed", )) @@ -3662,6 +3667,7 @@ pub async fn approve_fork( crate::fork_policy::sweep_expired_fork_approvals(&shared, now_unix_nanos).await; if !expired.is_empty() { shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); } return Err(ApiError::conflict( "approval window expired; the run was failed closed", @@ -3694,6 +3700,7 @@ pub async fn approve_fork( *shared.state.next_job_runs_on.write().unwrap() = outcome.next_runs_on; if outcome.promoted > 0 { shared.state.message_notify.notify_waiters(); + shared.state.sampler_notify.notify_waiters(); } // The expiry sweep may have won the run's terminal transition between the // guard read and the write (it clears the hold in its own transaction). diff --git a/crates/preloop-runner-server/src/state.rs b/crates/preloop-runner-server/src/state.rs index 253bb284c..589711423 100644 --- a/crates/preloop-runner-server/src/state.rs +++ b/crates/preloop-runner-server/src/state.rs @@ -306,6 +306,17 @@ impl AppState { } } + /// Test hook: run one state-sampler tick inline. The harness never + /// spawns `run_state_sampler`, so pool/queue gauges stay at their seeded + /// values unless a test samples explicitly. + #[cfg(any(test, feature = "test-support"))] + pub async fn test_sample_state_once(&self) { + crate::bootstrap::sample_state_once( + &self.shared(), + &preloop_observability::status::StoreBackend::Sqlite, + ) + .await; + } /// Mutate the control database directly. Escape hatch for seeds that no /// `ControlBackend` command expresses (forced-terminal runs, planted /// timestamps). SQLite only: production servers never see it, and the @@ -702,6 +713,13 @@ pub struct AppState { /// ([`crate::event_feed`]). pub(crate) events_dirty: Arc, pub message_notify: Arc, + /// The state sampler's dedicated wake channel. Kept separate from + /// `message_notify`: runner waiters are woken with `notify_one` per + /// ready job (a permit), and a parked sampler would consume permits + /// meant for runners — a parked runner then misses work until its poll + /// window ends. Every producer that fires `message_notify` also fires + /// this (see [`wake_waiters`]); only the sampler listens on it. + pub sampler_notify: Arc, pub webhook_queue_notify: Arc, /// Circuit breaker for durable webhook delivery work. Lifecycle check-run /// updates use a separate breaker so an ingress outage cannot suppress @@ -1214,7 +1232,13 @@ impl AppState { env::var("PRELOOP_RUNNER_LIVENESS_TIMEOUT_SECS") .ok() .and_then(|raw| raw.trim().parse().ok()) - .unwrap_or(1800), + .unwrap_or(1800) + // Floor: the liveness sweep must never outrun a + // legitimate in-flight long-poll. A runner parked for + // the full client window sends no traffic; a timeout + // below ~3x the max poll window can reap a healthy, + // still-polling runner. + .max(3 * max_poll_window_secs()), ), ..Default::default() }; @@ -1401,21 +1425,25 @@ impl AppState { ); } let message_notify = Arc::new(Notify::new()); + let sampler_notify = Arc::new(Notify::new()); // One pool status handle, shared with the backend: the pool's // advertised labels gate `runs-on` inside submit/promotion. let pool_status = Arc::new(preloop_observability::status::PoolStatus::default()); backend.set_pool_status((*pool_status).clone()); // Cross-node wake-ups (Postgres LISTEN): a job committed through any - // node wakes runners long-polling this one. + // node wakes runners long-polling this one — and nudges the local + // sampler so queue gauges track remote commits between ticks. if let Some(mut wakes) = backend.subscribe_wakes() { let notify = message_notify.clone(); + let sampler = sampler_notify.clone(); tokio::spawn(async move { loop { match wakes.recv().await { - Ok(wake) => wake_waiters(¬ify, wake.ready, wake.broadcast), + Ok(wake) => wake_waiters(¬ify, &sampler, wake.ready, wake.broadcast), // Missed signals: wake everyone once; they re-check. Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { - notify.notify_waiters() + notify.notify_waiters(); + sampler.notify_waiters(); } Err(tokio::sync::broadcast::error::RecvError::Closed) => break, } @@ -1439,6 +1467,7 @@ impl AppState { events, events_dirty, message_notify, + sampler_notify, webhook_queue_notify: Arc::new(Notify::new()), github_breaker: Arc::new(crate::github_breaker::GithubBreaker::default()), github_lifecycle_breaker: Arc::new(crate::github_breaker::GithubBreaker::default()), @@ -2181,10 +2210,29 @@ mod cluster_key_tests { } } +/// Upper bound on a runner long-poll `waitSeconds` (seconds). +/// `PRELOOP_MAX_POLL_WINDOW_SECS` overrides; the default keeps polls short +/// enough that a healthy runner always beats the liveness sweep (which is +/// floored at `3 *` this value — see `runner_liveness_timeout` seeding). +pub(crate) fn max_poll_window_secs() -> u64 { + env::var("PRELOOP_MAX_POLL_WINDOW_SECS") + .ok() + .and_then(|raw| raw.trim().parse().ok()) + .filter(|value| *value > 0) + .unwrap_or(60) +} + /// Wake long-polling runners: at most `ready` of them for newly claimable /// jobs (waking every waiter for one job is a thundering herd: all of them /// race one claim), or all of them for broadcasts such as cancellations. -pub(crate) fn wake_waiters(notify: &Notify, ready: usize, broadcast: bool) { +/// +/// The state sampler listens on a separate channel (`sampler`): the +/// `notify_one` permits below belong to runner waiters, and a parked +/// sampler parked on the same `Notify` would steal them. `notify_waiters` +/// on the sampler channel is safe — it wakes a registered waiter without +/// leaving a stale permit behind. +pub(crate) fn wake_waiters(notify: &Notify, sampler: &Notify, ready: usize, broadcast: bool) { + sampler.notify_waiters(); if broadcast { notify.notify_waiters(); return; @@ -2219,10 +2267,51 @@ mod wake_tests { })); } tokio::time::sleep(std::time::Duration::from_millis(50)).await; - super::wake_waiters(¬ify, 2, false); + let sampler = Arc::new(Notify::new()); + super::wake_waiters(¬ify, &sampler, 2, false); for waiter in waiters { waiter.await.unwrap(); } assert_eq!(woken.load(Ordering::SeqCst), 2); } + + /// A parked sampler must never consume a `notify_one` permit meant for a + /// runner: regression cover for the pre-split bug where the sampler + /// shared `message_notify` and could park on the same `Notify`. + #[tokio::test] + async fn parked_sampler_never_steals_a_runner_permit() { + let notify = Arc::new(Notify::new()); + let sampler = Arc::new(Notify::new()); + let woken = Arc::new(AtomicUsize::new(0)); + // The sampler parks on its own channel. + let (sampler_handle, woken_s) = (sampler.clone(), woken.clone()); + let sampler_task = tokio::spawn(async move { + if tokio::time::timeout( + std::time::Duration::from_millis(300), + sampler_handle.notified(), + ) + .await + .is_ok() + { + woken_s.fetch_add(1, Ordering::SeqCst); + } + }); + // One runner parks on the runner channel. + let (notify_r, woken_r) = (notify.clone(), woken.clone()); + let runner_task = tokio::spawn(async move { + if tokio::time::timeout(std::time::Duration::from_millis(300), notify_r.notified()) + .await + .is_ok() + { + woken_r.fetch_add(1, Ordering::SeqCst); + } + }); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + super::wake_waiters(¬ify, &sampler, 1, false); + sampler_task.await.unwrap(); + runner_task.await.unwrap(); + // Both wake: the sampler via notify_waiters on its channel, the + // runner via the single permit — the parked sampler can't take it. + assert_eq!(woken.load(Ordering::SeqCst), 2); + } } diff --git a/crates/preloop-runner-server/tests/logs.rs b/crates/preloop-runner-server/tests/logs.rs index 0f2c67be2..1173ab726 100644 --- a/crates/preloop-runner-server/tests/logs.rs +++ b/crates/preloop-runner-server/tests/logs.rs @@ -2812,6 +2812,9 @@ async fn cancel_run_refreshes_runner_pool_queue_metadata() { }), ) .await; + // Queue depth is sampled by the 5s sampler, not updated per operation; + // run one tick inline (the test harness spawns no sampler task). + state.test_sample_state_once().await; assert_eq!( *state.next_job_runs_on.read().unwrap(), vec!["ubuntu-22.04"] @@ -2828,6 +2831,7 @@ async fn cancel_run_refreshes_runner_pool_queue_metadata() { Value::Null, ) .await; + state.test_sample_state_once().await; assert_eq!(cancelled["conclusion"], "cancelled"); assert_eq!( @@ -2864,6 +2868,9 @@ async fn disttask_claim_refreshes_runner_pool_queue_metadata() { ) .await; } + // Queue depth is sampled by the 5s sampler, not updated per operation; + // run one tick inline (the test harness spawns no sampler task). + state.test_sample_state_once().await; assert_eq!( state.queue_depth.load(std::sync::atomic::Ordering::Acquire), 2 @@ -2895,6 +2902,7 @@ async fn disttask_claim_refreshes_runner_pool_queue_metadata() { message["messageType"], azdo::message_type::PIPELINE_AGENT_JOB_REQUEST ); + state.test_sample_state_once().await; assert_eq!( state.queue_depth.load(std::sync::atomic::Ordering::Acquire),