Skip to main content

faucet_cli/serve/
runner.rs

1//! The run lifecycle: validate + queue a submission (`submit`), then run it under
2//! a permit. A cancel / timeout / shutdown trigger cooperatively cancels the
3//! pipeline (so a buffered sink flushes at its next page boundary, #146 H16) and
4//! grants a bounded flush grace before hard-dropping it; the task then finalizes
5//! an authoritative terminal status. See spec §7 + §20.
6
7use crate::auth_catalog::build_auth_catalog;
8use crate::executor::{ExecuteOptions, RunSummary, run_expanded};
9use crate::registry::build_source;
10use crate::serve::error::ServeError;
11use crate::serve::history::{Claim, InvocationRecord, RunRecord, RunStatus};
12use crate::serve::history::{ClaimedShard, ShardInsert};
13use crate::serve::load::{ConfigFormat, LoadedSubmission, load_submission};
14use crate::serve::rbac::AuthContext;
15use crate::serve::state::ServerState;
16use crate::serve::{idempotency, metrics};
17use chrono::{DateTime, FixedOffset, Utc};
18use serde::{Deserialize, Serialize};
19use std::collections::BTreeMap;
20use std::time::Duration;
21use tokio_util::sync::CancellationToken;
22use tracing::Instrument;
23
24/// `Retry-After` advertised when the queue is full.
25const QUEUE_FULL_RETRY_AFTER_SECS: u64 = 5;
26
27/// Grace granted to a cancelled / timed-out / shutting-down run to flush
28/// buffered sink output cooperatively before its future is hard-dropped (which
29/// aborts the pipeline's task set, the backstop for a run stuck mid-write so a
30/// hung run can't wedge shutdown). Generous enough for an S3 multipart
31/// completion (#146 H16).
32pub(crate) const RUN_FLUSH_GRACE: Duration = Duration::from_secs(30);
33
34/// `POST /v1/runs` request body.
35#[derive(Debug, Deserialize)]
36pub struct SubmitRequest {
37    pub config: String,
38    #[serde(default)]
39    pub config_format: ConfigFormatWire,
40    pub name: Option<String>,
41    #[serde(default)]
42    pub labels: BTreeMap<String, String>,
43    pub timeout_secs: Option<u64>,
44    #[serde(default)]
45    pub doctor_first: bool,
46    pub idempotency_key: Option<String>,
47    pub clock: Option<String>,
48    /// Optional completion callback fired when this run reaches a terminal
49    /// state (#481). Validated at submit time so a bad destination is a 422 on
50    /// the submission rather than a silent no-op when the run finishes.
51    #[serde(default)]
52    pub callback: Option<crate::serve::callback::CallbackSpec>,
53}
54
55/// Wire enum mirroring `load::ConfigFormat` with serde rename.
56#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)]
57#[serde(rename_all = "lowercase")]
58pub enum ConfigFormatWire {
59    #[default]
60    Yaml,
61    Json,
62}
63
64impl From<ConfigFormatWire> for ConfigFormat {
65    fn from(w: ConfigFormatWire) -> Self {
66        match w {
67            ConfigFormatWire::Yaml => ConfigFormat::Yaml,
68            ConfigFormatWire::Json => ConfigFormat::Json,
69        }
70    }
71}
72
73/// `POST /v1/runs` success body (202).
74#[derive(Debug, Serialize)]
75pub struct SubmitResponse {
76    pub run_id: String,
77    pub status: RunStatus,
78    pub submitted_at: DateTime<Utc>,
79}
80
81/// Re-run a claimed Pending run on this instance (cluster mode). Reconstructs the
82/// execution inputs from the persisted record (re-resolving config with this
83/// instance's own env/credentials), acquires a permit, and runs the shared tail.
84pub fn resume_claimed_run(state: ServerState, rec: RunRecord) {
85    tokio::spawn(async move {
86        let run_id = rec.run_id.clone();
87        let Some(body) = rec.config_body.as_deref() else {
88            tracing::error!(run_id, "claimed run has no stored config; failing it");
89            finalize(
90                &state,
91                &run_id,
92                rec.submitted_at,
93                Terminal::Failed {
94                    reason: "claimed run record missing config_body".into(),
95                    records: 0,
96                    invs: Vec::new(),
97                },
98            )
99            .await;
100            return;
101        };
102        let format = rec.config_format.unwrap_or_default();
103        let loaded = match load_submission(body, format, state.default_base().as_ref()).await {
104            Ok(l) => l,
105            Err(e) => {
106                finalize(
107                    &state,
108                    &run_id,
109                    rec.submitted_at,
110                    Terminal::Failed {
111                        reason: format!(
112                            "re-loading claimed config: {}",
113                            e.api_error().error.message
114                        ),
115                        records: 0,
116                        invs: Vec::new(),
117                    },
118                )
119                .await;
120                return;
121            }
122        };
123
124        // Mode B (#230): a sharded run is expanded into shard rows here — the
125        // claiming instance acts as the (ephemeral) coordinator — and is NOT
126        // executed as a whole. Enumeration + insert is idempotent, so two
127        // instances both claiming + coordinating converge on the same shard set.
128        if let Some(sh) = loaded.cfg.shard.clone()
129            && sh.count >= 2
130        {
131            match coordinate_sharded_run(&state, &run_id, &loaded, sh.count).await {
132                Ok(true) => return, // expanded into shards — shard loop runs them
133                Ok(false) => {}     // not shardable → fall through, run the whole run
134                Err(e) => {
135                    finalize(
136                        &state,
137                        &run_id,
138                        rec.submitted_at,
139                        Terminal::Failed {
140                            reason: format!("sharding: {e}"),
141                            records: 0,
142                            invs: Vec::new(),
143                        },
144                    )
145                    .await;
146                    return;
147                }
148            }
149        }
150
151        // The claim loop only claims up to available_permits and is the sole
152        // permit consumer, so this acquire returns immediately.
153        let _permit = state
154            .semaphore()
155            .acquire_owned()
156            .await
157            .expect("semaphore not closed");
158        // Register a local cancel token so a cross-instance cancel (the claim loop
159        // calling registry.cancel) reaches this run.
160        let run_token = CancellationToken::new();
161        state.registry().register(run_id.clone(), run_token.clone());
162        execute_run(
163            state.clone(),
164            loaded,
165            run_id,
166            run_token,
167            rec.submitted_at,
168            rec.timeout_secs,
169            rec.clock.clone(),
170            false,
171        )
172        .await;
173    });
174}
175
176/// Coordinator step (Mode B): expand a sharded run into `faucet_serve_shards`
177/// rows. Returns `Ok(true)` when the run was sharded (caller must not execute it
178/// as a whole), `Ok(false)` when it isn't shardable (caller runs it whole).
179///
180/// Idempotent: enumeration is deterministic and the insert is
181/// `ON CONFLICT DO NOTHING`, so a re-coordinated run (e.g. after the coordinator
182/// crashed and the Pending run was requeued) converges on the same shard set.
183async fn coordinate_sharded_run(
184    state: &ServerState,
185    run_id: &str,
186    loaded: &LoadedSubmission,
187    count: usize,
188) -> crate::error::CliResult<bool> {
189    use crate::error::CliError;
190
191    // Sharding applies to a single-source pipeline; a matrix fan-out is not
192    // shardable (each row is already an independent unit — use Mode A).
193    if loaded.nodes.len() != 1 {
194        tracing::warn!(
195            run_id,
196            nodes = loaded.nodes.len(),
197            "shard requested but the run is not a single-node pipeline; running it whole"
198        );
199        return Ok(false);
200    }
201    let node = &loaded.nodes[0];
202    let auth = build_auth_catalog(loaded.cfg.auth.as_ref())
203        .map_err(|e| CliError::Internal(format!("auth catalog: {e}")))?;
204    let source = build_source(&node.source.kind, node.source.config.clone(), &auth, None).await?;
205    if !source.is_shardable() {
206        tracing::warn!(
207            run_id,
208            kind = %node.source.kind,
209            "source is not shardable; running the run whole"
210        );
211        return Ok(false);
212    }
213
214    // Mark the parent run Sharded BEFORE inserting any shard rows (F27).
215    // Orphan recovery (`recover_orphans` / `reclaim_orphans`) only reclaims
216    // `running` runs; once the parent is `Sharded` it is immune to a false
217    // fail. Doing this first guarantees the invariant that **no shard row ever
218    // exists while the parent is still `running`** — otherwise a coordinator
219    // crash between `insert_shards` and the status flip would let orphan
220    // recovery mark the run Failed while its already-inserted shards are
221    // claimed, execute, and write data (a run the user believes failed and may
222    // resubmit → duplicate writes). If the flip itself fails, the run stays
223    // `running` + owned by us with no shards inserted, so a crash here is
224    // safely *requeued* by reclaim, not falsely failed. (A crash after the flip
225    // but before `insert_shards` leaves a 0-shard `Sharded` run, which
226    // `all_terminal()` never finalizes — an availability leak, not a
227    // correctness/duplication bug; the desired trade.)
228    {
229        let mut r = state
230            .history()
231            .get(run_id)
232            .await
233            .map_err(|e| CliError::Internal(e.to_string()))?
234            .ok_or_else(|| CliError::Internal(format!("run {run_id} vanished before sharding")))?;
235        r.status = RunStatus::Sharded;
236        state
237            .history()
238            .upsert(&r)
239            .await
240            .map_err(|e| CliError::Internal(e.to_string()))?;
241    }
242
243    let shards = source
244        .enumerate_shards(count)
245        .await
246        .map_err(|e| CliError::Internal(format!("enumerate_shards: {e}")))?;
247    let inserts: Vec<ShardInsert> = shards
248        .iter()
249        .map(|s| ShardInsert {
250            shard_id: s.id.clone(),
251            descriptor: s.descriptor.clone(),
252            size_estimate: s.size_estimate,
253        })
254        .collect();
255    let inserted = state
256        .history()
257        .insert_shards(run_id, &inserts)
258        .await
259        .map_err(|e| CliError::Internal(e.to_string()))?;
260    tracing::info!(
261        run_id,
262        shards = inserts.len(),
263        inserted,
264        "expanded run into shards (Mode B)"
265    );
266
267    // Wake the local claim loop so it picks up the freshly-inserted shards.
268    state.cluster().kick();
269    Ok(true)
270}
271
272/// Execute one claimed shard (Mode B): rebuild + narrow the source to the shard,
273/// run it under a permit, owner-fenced-finalize the shard, then finalize the
274/// parent run once every shard is terminal.
275pub fn resume_claimed_shard(state: ServerState, claimed: ClaimedShard) {
276    tokio::spawn(async move {
277        let ClaimedShard {
278            run_id,
279            shard_id,
280            descriptor,
281            run,
282        } = claimed;
283
284        let Some(body) = run.config_body.clone() else {
285            tracing::error!(run_id, shard_id, "claimed shard's run has no stored config");
286            let _ = state
287                .history()
288                .finalize_shard(&run_id, &shard_id, false)
289                .await;
290            maybe_finalize_parent(&state, &run_id).await;
291            return;
292        };
293        let format = run.config_format.unwrap_or_default();
294        let loaded = match load_submission(&body, format, state.default_base().as_ref()).await {
295            Ok(l) => l,
296            Err(e) => {
297                tracing::error!(
298                    run_id,
299                    shard_id,
300                    error = %e.api_error().error.message,
301                    "re-loading shard config failed"
302                );
303                let _ = state
304                    .history()
305                    .finalize_shard(&run_id, &shard_id, false)
306                    .await;
307                maybe_finalize_parent(&state, &run_id).await;
308                return;
309            }
310        };
311
312        let _permit = state
313            .semaphore()
314            .acquire_owned()
315            .await
316            .expect("semaphore not closed");
317
318        let shard = faucet_core::ShardSpec {
319            id: shard_id.clone(),
320            descriptor,
321            size_estimate: None,
322        };
323        // Register a per-shard cooperative-cancel token BEFORE running so a
324        // cross-instance cancel reaches this shard: `cancel_run` flags the parent
325        // → `pending_shard_cancellations` → the claim loop calls
326        // `registry().cancel_run_shards(run_id)`, which fires this token. Removed
327        // on return so a terminated shard never leaks a token (shard accounting is
328        // separate from the run's `in_flight`, so a plain `deregister` is used).
329        let coop = CancellationToken::new();
330        state
331            .registry()
332            .register_shard(&run_id, &shard_id, coop.clone());
333        // Count the shard in `in_flight` so the shutdown drain waits for it and
334        // fires the cooperative-flush cancel (audit #321 H5). The guard releases
335        // the slot + token on EVERY return path (including panic), so a shard can
336        // never leak the counter and wedge graceful shutdown.
337        state.registry().mark_shard_running();
338        let _shard_guard = ShardInFlightGuard {
339            state: state.clone(),
340            run_id: run_id.clone(),
341            shard_id: shard_id.clone(),
342        };
343        let success = execute_shard(
344            &state,
345            loaded,
346            &run_id,
347            &shard_id,
348            shard,
349            coop,
350            run.timeout_secs,
351            run.clock.clone(),
352            run.submitted_at,
353        )
354        .await;
355
356        match state
357            .history()
358            .finalize_shard(&run_id, &shard_id, success)
359            .await
360        {
361            Ok(true) => {}
362            Ok(false) => tracing::warn!(
363                run_id,
364                shard_id,
365                "shard was reclaimed by another instance; discarding result"
366            ),
367            Err(e) => tracing::error!(run_id, shard_id, error = %e, "finalize_shard failed"),
368        }
369        maybe_finalize_parent(&state, &run_id).await;
370    });
371}
372
373/// Run one shard's pipeline (single node, source narrowed via `opts.shard`).
374/// Returns `true` on clean completion. Does not touch the parent run record —
375/// the caller finalizes the shard and the parent.
376#[allow(clippy::too_many_arguments)]
377async fn execute_shard(
378    state: &ServerState,
379    loaded: LoadedSubmission,
380    run_id: &str,
381    shard_id: &str,
382    shard: faucet_core::ShardSpec,
383    coop: CancellationToken,
384    timeout_secs: Option<u64>,
385    clock_flag: Option<String>,
386    submitted_at: DateTime<Utc>,
387) -> bool {
388    let LoadedSubmission { cfg, nodes } = loaded;
389    let pipeline_name = cfg.name.clone().unwrap_or_else(|| "serve".to_string());
390
391    let auth = match build_auth_catalog(cfg.auth.as_ref()) {
392        Ok(a) => a,
393        Err(e) => {
394            tracing::error!(run_id, shard_id, "shard auth catalog: {e}");
395            return false;
396        }
397    };
398    let clock = match resolve_clock(clock_flag.as_deref(), submitted_at) {
399        Ok(c) => c,
400        Err(e) => {
401            tracing::error!(
402                run_id,
403                shard_id,
404                "shard clock: {}",
405                e.api_error().error.message
406            );
407            return false;
408        }
409    };
410    let resilience = match &cfg.resilience {
411        Some(spec) => match spec.to_policy() {
412            Ok(p) => Some(p),
413            Err(e) => {
414                tracing::error!(run_id, shard_id, "shard resilience: {e}");
415                return false;
416            }
417        },
418        None => None,
419    };
420    #[cfg(feature = "lineage")]
421    let lineage = match crate::lineage_glue::build_emitter(cfg.lineage.as_ref()) {
422        Ok(l) => l,
423        Err(e) => {
424            tracing::error!(run_id, shard_id, "shard lineage: {e}");
425            return false;
426        }
427    };
428
429    // `coop` is the per-shard cancel token registered by `resume_claimed_shard`;
430    // a cross-instance cancel (F10), a server shutdown, or a timeout all fire it
431    // so the shard's pipeline flushes at its next page boundary.
432    let opts = ExecuteOptions {
433        pipeline_name,
434        // Correlate notifications to the submitted run (#480). Every shard of a
435        // sharded run reports the same `run_id`; they differ by `invocation_id`.
436        run_id: Some(run_id.to_string()),
437        execution: cfg.execution.clone(),
438        dry_run: false,
439        limit: None,
440        state_path_override: None,
441        shard: Some(shard),
442        auth,
443        clock,
444        cancel: Some(coop.clone()),
445        resilience,
446        // Inert under `shard: Some(..)` (a shard's volume is not the row's) —
447        // carried for uniformity with the whole-run path below.
448        sla: cfg.sla.clone(),
449        reconcile: cfg.reconcile.clone(),
450        #[cfg(feature = "lineage")]
451        lineage,
452        #[cfg(feature = "lineage")]
453        lineage_cfg: cfg.lineage.clone(),
454        #[cfg(feature = "notify")]
455        notifier: None,
456        // Inert under `shard: Some(..)` — a shard's records/URIs are a slice
457        // of the row's; the catalog records whole runs only.
458        #[cfg(feature = "catalog")]
459        catalog: None,
460    };
461
462    let server_shutdown = state.shutdown_token();
463    let span = tracing::info_span!("faucet.serve.shard", serve_run_id = %run_id, shard = %shard_id);
464    let work = async move { classify_run(run_expanded(nodes, opts).await) }.instrument(span);
465    tokio::pin!(work);
466    let timeout_fut = async {
467        match timeout_secs {
468            Some(s) => tokio::time::sleep(Duration::from_secs(s)).await,
469            None => std::future::pending::<()>().await,
470        }
471    };
472    tokio::pin!(timeout_fut);
473
474    // Cancel triggers (shutdown / timeout) cooperatively cancel + flush, like
475    // execute_run. A failed shard simply returns false → its lease eventually
476    // reassigns it (or it poisons after max_attempts).
477    // A cross-instance cancel (F10) flows in via the shard's registered coop
478    // token: `cancel_run` flags the parent → `pending_cancellations` →
479    // `claim_loop` fires `registry().cancel(run_id)`. The token is registered
480    // under the run id by `resume_claimed_shard` before this runs.
481    let terminal = tokio::select! {
482        biased;
483        t = &mut work => t,
484        _ = coop.cancelled() => {
485            // Fired by the claim loop for a remote cancel. Flush within the grace
486            // window; a cooperative cancel returns Ok(partial), so the shard is
487            // Cancelled — but a flush that FAILS must surface, not be masked.
488            match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
489                Ok(failed @ Terminal::Failed { .. }) => failed,
490                Ok(_) | Err(_) => Terminal::Cancelled,
491            }
492        }
493        _ = server_shutdown.cancelled() => {
494            coop.cancel();
495            match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
496                Ok(failed @ Terminal::Failed { .. }) => failed,
497                Ok(_) | Err(_) => Terminal::ShutdownFailed,
498            }
499        }
500        _ = &mut timeout_fut => {
501            coop.cancel();
502            match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
503                Ok(failed @ Terminal::Failed { .. }) => failed,
504                Ok(_) | Err(_) => Terminal::Timeout { secs: timeout_secs.unwrap_or(0) },
505            }
506        }
507    };
508    matches!(terminal, Terminal::Completed { .. })
509}
510
511/// Finalize a `Sharded` parent run once all its shards are terminal. The last
512/// shard to finish always observes `all_terminal` (its own `finalize_shard`
513/// committed first), so the run never lingers `Sharded`. A benign double-finalize
514/// (two shards finishing simultaneously) writes the same terminal status twice.
515async fn maybe_finalize_parent(state: &ServerState, run_id: &str) {
516    let progress = match state.history().shard_progress(run_id).await {
517        Ok(p) => p,
518        Err(e) => {
519            tracing::warn!(run_id, error = %e, "shard_progress failed");
520            return;
521        }
522    };
523    if !progress.all_terminal() {
524        return;
525    }
526    let success = progress.failed == 0;
527    let status = if success {
528        RunStatus::Completed
529    } else {
530        RunStatus::Failed
531    };
532    let error =
533        (!success).then(|| format!("{}/{} shard(s) failed", progress.failed, progress.total));
534    // Status-fenced finalize: transitions the parent only while it is still
535    // `Sharded`, and does NOT re-stamp owner/lease on the terminal record, so a
536    // near-simultaneous double-finalize from two instances has a single winner
537    // (F45). The loser (and any non-`Sharded` parent) is a no-op.
538    match state
539        .history()
540        .finalize_sharded_parent(run_id, status, Utc::now(), error)
541        .await
542    {
543        Ok(true) => {
544            metrics::record_run_finished(status, if success { "ok" } else { "error" });
545            tracing::info!(
546                run_id,
547                shards = progress.total,
548                failed = progress.failed,
549                "sharded run finalized"
550            );
551            // Fire the parent's completion callback exactly once — `Ok(true)` is
552            // the status-fenced winner, so a near-simultaneous double-finalize
553            // from two instances delivers a single callback (#481). Re-read the
554            // record: `finalize_sharded_parent` takes fields, not the record, so
555            // this is the only way to get the persisted terminal state.
556            match state.history().get(run_id).await {
557                Ok(Some(rec)) => crate::serve::callback::fire(&rec).await,
558                Ok(None) => tracing::warn!(
559                    run_id,
560                    "sharded parent vanished before its callback could fire"
561                ),
562                Err(e) => tracing::warn!(
563                    run_id,
564                    error = %e,
565                    "could not read sharded parent for its completion callback"
566                ),
567            }
568        }
569        Ok(false) => {} // already finalized by another shard/instance — nothing to do
570        Err(e) => {
571            tracing::error!(run_id, error = %e, "finalizing sharded parent run failed");
572        }
573    }
574}
575
576/// Warn at submit when a clustered or source-sharded run is at-least-once with
577/// a non-idempotent destination (F26/F39). Cluster failover and shard reclaim
578/// are at-least-once by construction: an owner whose lease lapses while it is
579/// still alive (slow, paused), or a reassigned shard whose source re-reads from
580/// the start (postgres PK-range sharding yields no resumable bookmark), can
581/// re-execute work and write **duplicate** rows to an append-mode sink. The
582/// safe configuration is `write_mode: upsert`/`delete` (keyed, so re-writes are
583/// idempotent) or `delivery: exactly_once`. We surface the risk loudly rather
584/// than silently duplicating — the repo's #1 worst bug class.
585/// Pure decision for [`warn_if_cluster_at_least_once`]: the sink kinds that
586/// would write non-idempotently under cluster failover / shard reclaim. Empty
587/// when the run is neither clustered nor sharded, is `exactly_once`, or every
588/// sink is keyed (`upsert`/`delete`).
589fn at_least_once_risky_sinks(
590    loaded: &LoadedSubmission,
591    clustered: bool,
592    sharded: bool,
593) -> Vec<&str> {
594    if !(clustered || sharded) || loaded.cfg.delivery == faucet_core::DeliveryMode::ExactlyOnce {
595        return Vec::new();
596    }
597    loaded
598        .nodes
599        .iter()
600        .filter(|n| {
601            !matches!(
602                n.sink
603                    .config
604                    .get("write_mode")
605                    .and_then(|v| v.as_str())
606                    .unwrap_or("append"),
607                "upsert" | "delete"
608            )
609        })
610        .map(|n| n.sink.kind.as_str())
611        .collect()
612}
613
614fn warn_if_cluster_at_least_once(loaded: &LoadedSubmission, clustered: bool, sharded: bool) {
615    let risky = at_least_once_risky_sinks(loaded, clustered, sharded);
616    if !risky.is_empty() {
617        let scope = if sharded {
618            "source-sharded (Mode B)"
619        } else {
620            "clustered"
621        };
622        tracing::warn!(
623            sinks = ?risky,
624            "{scope} execution is at-least-once: a failover or shard reclaim can re-run work \
625             and write duplicate rows to an append-mode sink. Set `write_mode: upsert` (or \
626             `delivery: exactly_once`) on the destination to make re-execution idempotent (F26/F39)."
627        );
628    }
629}
630
631/// Best-effort release of an idempotency claim whose paired run-record upsert
632/// failed (F21). No-op when the request carried no key; the release itself is
633/// scoped by `run_id` (a newer run that re-claimed the key keeps its claim).
634async fn release_orphaned_claim(state: &ServerState, req: &SubmitRequest, run_id: &str) {
635    if req.idempotency_key.is_some()
636        && let Err(e) = state.history().release_idempotency(run_id).await
637    {
638        tracing::warn!(
639            run_id,
640            error = %e,
641            "failed to release idempotency claim after a run-record write error; \
642             a replay of the key may 404 until the claim self-expires"
643        );
644    }
645}
646
647/// Validate, idempotency-claim, queue, and spawn a submission.
648pub async fn submit(
649    state: ServerState,
650    req: SubmitRequest,
651    actor: AuthContext,
652) -> Result<SubmitResponse, ServeError> {
653    let format: ConfigFormat = req.config_format.into();
654    let loaded = load_submission(&req.config, format, state.default_base().as_ref()).await?;
655
656    // At-least-once duplicate-write warning for clustered / source-sharded runs
657    // with an append-mode destination (F26/F39).
658    let sharded = loaded.cfg.shard.as_ref().is_some_and(|s| s.count >= 2);
659    warn_if_cluster_at_least_once(&loaded, state.cluster().enabled(), sharded);
660
661    // Completion-callback validation (#481) — before the queue reservation, so a
662    // malformed or refused destination costs nothing and surfaces as a 422 on
663    // the submission rather than as a silent no-op when the run finishes.
664    if let Some(cb) = req.callback.as_ref() {
665        cb.validate(state.callback_allow_hosts())
666            .map_err(|message| ServeError::Unprocessable {
667                message,
668                details: None,
669            })?;
670        cb.reject_secrets_in_cluster(state.cluster().enabled())
671            .map_err(|message| ServeError::Unprocessable {
672                message,
673                details: None,
674            })?;
675    }
676
677    // Reserve a queue slot first, so a Fresh idempotency claim is always followed
678    // by a spawned run (no orphaned claims — spec §20.2).
679    if !state.registry().try_reserve() {
680        return Err(ServeError::QueueFull {
681            retry_after_secs: QUEUE_FULL_RETRY_AFTER_SECS,
682        });
683    }
684    // Releases the reservation on ANY early return below (doctor_first 422 /
685    // replay / conflict / claim or upsert error). Defused just before spawn.
686    let reservation = ReservationGuard::new(state.clone());
687
688    // doctor_first preflight — run BEHIND the reservation so concurrent preflight
689    // probing is bounded by `max_queued_runs` rather than running unthrottled
690    // before any limit applies (#146 R). On failure the guard releases the slot
691    // via the early `?`. The (redacted) report is stored on the run record below
692    // so `GET /v1/runs/{id}` exposes it (#146 R: doctor_report was never set).
693    let doctor_report = if req.doctor_first {
694        Some(run_doctor_first(&state, &loaded).await?)
695    } else {
696        None
697    };
698
699    let run_id = uuid::Uuid::now_v7().to_string();
700
701    // Config fingerprint (sha256) — the idempotency identity AND the value
702    // recorded on the `run.submit` audit entry (#205).
703    let merged = serde_json::to_value(&loaded.cfg).unwrap_or(serde_json::Value::Null);
704    let fp_config = idempotency::fingerprint(&merged, loaded.cfg.name.as_deref());
705
706    // Idempotency claim (if a key was supplied).
707    if let Some(key) = &req.idempotency_key {
708        // Fold the run-affecting request fields (clock / timeout_secs / labels)
709        // into the fingerprint, not just the config — so a key replayed with a
710        // different backfill `clock` is a 409, not a replay of the original
711        // run's window (#146 M7).
712        let fp = idempotency::request_fingerprint(
713            &fp_config,
714            req.clock.as_deref(),
715            req.timeout_secs,
716            &req.labels,
717        );
718        match state
719            .history()
720            .claim_idempotency(key, &fp, &run_id, state.idempotency_retention())
721            .await
722            .map_err(|e| match e {
723                // Degraded backend can't safely honor idempotency → 503, retry.
724                crate::serve::history::HistoryError::Degraded(m) => ServeError::Unavailable(m),
725                other => ServeError::Internal(other.to_string()),
726            })? {
727            Claim::Fresh => {}
728            Claim::Replay(existing) => {
729                metrics::record_idempotency_hit();
730                return replay_response(&state, &existing).await;
731            }
732            Claim::Conflict => {
733                return Err(ServeError::Conflict(
734                    "idempotency key reused with a different payload".into(),
735                ));
736            }
737        }
738        // A `Fresh` claim is recorded BEFORE the record upsert below. The memory
739        // backend's `upsert` is infallible, but a SQL backend's can fail — so if
740        // the upsert fails, `release_orphaned_claim` drops the claim (F21) rather
741        // than leaving a replay 404-ing until the claim self-expires.
742    }
743
744    let submitted_at = Utc::now();
745    let mut rec = RunRecord::queued(
746        run_id.clone(),
747        req.name.clone(),
748        req.labels.clone(),
749        req.idempotency_key.clone(),
750        submitted_at,
751    );
752    rec.doctor_report = doctor_report;
753    rec.callback = req.callback.clone();
754
755    if state.cluster().enabled() {
756        // A degraded (DB-unreachable) backend can't coordinate a cluster: the
757        // claim loop's claim_pending is a no-op on the in-memory fallback, so a
758        // Pending run would never be claimed. Fail closed with a retryable 503
759        // rather than silently orphaning the run (#197 spec §9).
760        if state.history().degraded() {
761            return Err(ServeError::Unavailable(
762                "clustered run-history backend is degraded; runs cannot be claimed \
763                 by any instance — retry once it recovers"
764                    .into(),
765            ));
766        }
767        // Cluster mode: persist the RAW config so any instance can re-resolve +
768        // run it, mark the run Pending, and wake the local claim loop. No local
769        // queue slot / spawn — the claim loop owns execution.
770        rec.status = RunStatus::Pending;
771        rec.config_body = Some(req.config.clone());
772        rec.config_format = Some(req.config_format.into());
773        rec.timeout_secs = req.timeout_secs;
774        rec.clock = req.clock.clone();
775        if let Err(e) = state.history().upsert(&rec).await {
776            // The record write that should follow a `Fresh` claim failed — release
777            // the orphaned claim so a replay starts fresh instead of 404-ing for
778            // the whole retention window (F21). Best-effort.
779            release_orphaned_claim(&state, &req, &run_id).await;
780            return Err(ServeError::Internal(e.to_string()));
781        }
782        // Release the local queue reservation (cluster runs are bounded by the
783        // claim loop + semaphore, not the submit-side queue).
784        drop(reservation);
785        state.cluster().kick();
786        crate::serve::audit::write(
787            &state,
788            &actor,
789            "run.submit",
790            Some(run_id.clone()),
791            Some(fp_config.clone()),
792            "ok",
793        )
794        .await;
795        return Ok(SubmitResponse {
796            run_id,
797            status: RunStatus::Pending,
798            submitted_at,
799        });
800    }
801
802    if let Err(e) = state.history().upsert(&rec).await {
803        // See the cluster path above: release the orphaned claim (F21).
804        release_orphaned_claim(&state, &req, &run_id).await;
805        return Err(ServeError::Internal(e.to_string()));
806    }
807
808    let run_token = CancellationToken::new();
809    state.registry().register(run_id.clone(), run_token.clone());
810    metrics::set_run_gauges(&state);
811
812    // The spawned task now owns the queued→running→finished lifecycle.
813    reservation.defuse();
814    spawn_run(
815        state.clone(),
816        loaded,
817        req,
818        run_id.clone(),
819        run_token,
820        submitted_at,
821    );
822
823    crate::serve::audit::write(
824        &state,
825        &actor,
826        "run.submit",
827        Some(run_id.clone()),
828        Some(fp_config.clone()),
829        "ok",
830    )
831    .await;
832
833    Ok(SubmitResponse {
834        run_id,
835        status: RunStatus::Queued,
836        submitted_at,
837    })
838}
839
840/// Run the `doctor_first` probes; on any failure return 422 with the report.
841/// Run the `doctor_first` probes. On success returns the (redacted) report so
842/// the caller can store it on the run record (`doctor_report`); on any probe
843/// failure returns 422 with the same redacted report as `details`.
844pub(crate) async fn run_doctor_first(
845    state: &ServerState,
846    loaded: &LoadedSubmission,
847) -> Result<serde_json::Value, ServeError> {
848    use faucet_core::check::CheckContext;
849    let auth =
850        build_auth_catalog(loaded.cfg.auth.as_ref()).map_err(|e| ServeError::Unprocessable {
851            message: e.to_string(),
852            details: None,
853        })?;
854    let ctx = CheckContext {
855        timeout: state.probe_timeout(),
856    };
857    // Same pipeline-name derivation as the run path above, so SLA probes read
858    // the state keys the executor writes.
859    let pipeline_name = loaded
860        .cfg
861        .name
862        .clone()
863        .unwrap_or_else(|| "serve".to_string());
864    let mut invs = crate::commands::doctor::probe_roots(
865        &loaded.nodes,
866        &auth,
867        &ctx,
868        loaded.cfg.sla.as_ref(),
869        &pipeline_name,
870    )
871    .await;
872    let failed = crate::commands::doctor::count_failures(&invs);
873    // Redact regardless of outcome — the report is surfaced either way (as the
874    // 422 `details` on failure, or stored on the run record on success).
875    crate::commands::doctor::redact_invocations(&mut invs);
876    let report = serde_json::json!({ "invocations": invs });
877    if failed > 0 {
878        return Err(ServeError::Unprocessable {
879            message: format!("doctor_first preflight failed: {failed} probe(s) failed"),
880            details: Some(report),
881        });
882    }
883    Ok(report)
884}
885
886/// Build the replay response for an idempotency hit (the existing run's status).
887async fn replay_response(state: &ServerState, run_id: &str) -> Result<SubmitResponse, ServeError> {
888    let rec = state
889        .history()
890        .get(run_id)
891        .await
892        .map_err(|e| ServeError::Internal(e.to_string()))?
893        .ok_or(ServeError::NotFound)?;
894    Ok(SubmitResponse {
895        run_id: rec.run_id,
896        status: rec.status,
897        submitted_at: rec.submitted_at,
898    })
899}
900
901/// Releases a queue reservation on drop unless [`Self::defuse`]d. Guarantees the
902/// `queued` counter is balanced on every early-return path (replay / conflict /
903/// claim-or-upsert error) without a manual `release_reservation` at each site.
904/// Defused once the run is handed to the spawned task, which then owns the
905/// queued→running transition via `mark_running`.
906struct ReservationGuard {
907    state: Option<ServerState>,
908}
909
910impl ReservationGuard {
911    fn new(state: ServerState) -> Self {
912        Self { state: Some(state) }
913    }
914
915    /// Hand the reservation off to the spawned task (no release on drop).
916    fn defuse(mut self) {
917        self.state = None;
918    }
919}
920
921impl Drop for ReservationGuard {
922    fn drop(&mut self) {
923        if let Some(state) = self.state.take() {
924            state.registry().release_reservation();
925            metrics::set_run_gauges(&state);
926        }
927    }
928}
929
930/// Releases the in-flight slot (decrement `in_flight`, drop the cancel token, wake
931/// the shutdown drain) on drop — on EVERY path including panic. Without this, a
932/// panic between `mark_running` and a manual `mark_finished` would leak the
933/// counter and hang graceful shutdown forever.
934struct InFlightGuard {
935    state: ServerState,
936    run_id: String,
937}
938
939impl Drop for InFlightGuard {
940    fn drop(&mut self) {
941        self.state.registry().mark_finished(&self.run_id);
942        metrics::set_run_gauges(&self.state);
943    }
944}
945
946/// Releases a claimed shard's in-flight slot (decrement `in_flight`, drop the
947/// shard cancel token, wake the shutdown drain) on drop — on EVERY path
948/// including panic. Mirrors [`InFlightGuard`] but keys on the per-shard token so
949/// a Mode-B shard participates in the graceful-shutdown drain (audit #321 H5).
950struct ShardInFlightGuard {
951    state: ServerState,
952    run_id: String,
953    shard_id: String,
954}
955
956impl Drop for ShardInFlightGuard {
957    fn drop(&mut self) {
958        self.state
959            .registry()
960            .mark_shard_finished(&self.run_id, &self.shard_id);
961        metrics::set_run_gauges(&self.state);
962    }
963}
964
965/// Terminal classification of a run task.
966enum Terminal {
967    Completed {
968        records: u64,
969        invs: Vec<InvocationRecord>,
970    },
971    Failed {
972        reason: String,
973        records: u64,
974        invs: Vec<InvocationRecord>,
975    },
976    Timeout {
977        secs: u64,
978    },
979    Cancelled,
980    ShutdownFailed,
981}
982
983impl Terminal {
984    /// (status, metric reason label, records, invocations, error message)
985    fn into_parts(
986        self,
987    ) -> (
988        RunStatus,
989        &'static str,
990        u64,
991        Vec<InvocationRecord>,
992        Option<String>,
993    ) {
994        match self {
995            Terminal::Completed { records, invs } => {
996                (RunStatus::Completed, "ok", records, invs, None)
997            }
998            Terminal::Failed {
999                reason,
1000                records,
1001                invs,
1002            } => (RunStatus::Failed, "error", records, invs, Some(reason)),
1003            Terminal::Timeout { secs } => (
1004                RunStatus::Failed,
1005                "timeout",
1006                0,
1007                Vec::new(),
1008                Some(format!("run exceeded timeout_secs ({secs}s)")),
1009            ),
1010            Terminal::Cancelled => (RunStatus::Cancelled, "cancelled", 0, Vec::new(), None),
1011            Terminal::ShutdownFailed => (
1012                RunStatus::Failed,
1013                "server_shutdown",
1014                0,
1015                Vec::new(),
1016                Some("server shutdown before the run finished".into()),
1017            ),
1018        }
1019    }
1020}
1021
1022/// Classify a finished `run_expanded` result into a `Terminal`.
1023fn classify_run(result: crate::error::CliResult<RunSummary>) -> Terminal {
1024    match result {
1025        Ok(summary) => {
1026            let records: u64 = summary
1027                .invocations
1028                .iter()
1029                .map(|i| i.records_written as u64)
1030                .sum();
1031            let invs: Vec<InvocationRecord> = summary
1032                .invocations
1033                .iter()
1034                .map(InvocationRecord::from)
1035                .collect();
1036            if summary.had_failures() {
1037                Terminal::Failed {
1038                    reason: format!("{} invocation(s) failed", summary.failure_count()),
1039                    records,
1040                    invs,
1041                }
1042            } else {
1043                Terminal::Completed { records, invs }
1044            }
1045        }
1046        Err(e) => Terminal::Failed {
1047            reason: e.to_string(),
1048            records: 0,
1049            invs: Vec::new(),
1050        },
1051    }
1052}
1053
1054/// Parse the optional request `clock` (RFC3339), defaulting to `submitted_at`.
1055fn resolve_clock(
1056    flag: Option<&str>,
1057    default: DateTime<Utc>,
1058) -> Result<DateTime<FixedOffset>, ServeError> {
1059    match flag {
1060        None => Ok(default.fixed_offset()),
1061        Some(s) => DateTime::parse_from_rfc3339(s)
1062            .map_err(|_| ServeError::BadConfig(format!("clock '{s}' is not RFC3339"))),
1063    }
1064}
1065
1066/// Spawn the detached run task: acquire a permit, run under the 3-arm select,
1067/// finalize the terminal status.
1068fn spawn_run(
1069    state: ServerState,
1070    loaded: LoadedSubmission,
1071    req: SubmitRequest,
1072    run_id: String,
1073    run_token: CancellationToken,
1074    submitted_at: DateTime<Utc>,
1075) {
1076    let server_shutdown = state.shutdown_token();
1077    tokio::spawn(async move {
1078        // Race the permit acquisition against cancel / shutdown so a run
1079        // cancelled while STILL QUEUED (before any permit frees) is finalized
1080        // immediately, instead of only after it eventually acquires a permit
1081        // (#146 R). `biased` prefers the cancel/shutdown signals over a
1082        // simultaneously-available permit.
1083        let _permit = tokio::select! {
1084            biased;
1085            _ = run_token.cancelled() => {
1086                finalize_queued_cancel(&state, &run_id, submitted_at, Terminal::Cancelled).await;
1087                return;
1088            }
1089            _ = server_shutdown.cancelled() => {
1090                finalize_queued_cancel(&state, &run_id, submitted_at, Terminal::ShutdownFailed).await;
1091                return;
1092            }
1093            permit = state.semaphore().acquire_owned() => permit.expect("semaphore not closed"),
1094        };
1095        execute_run(
1096            state,
1097            loaded,
1098            run_id,
1099            run_token,
1100            submitted_at,
1101            req.timeout_secs,
1102            req.clock,
1103            true,
1104        )
1105        .await;
1106        // `_permit` drops here.
1107    });
1108}
1109
1110/// The queued→running→finalize execution tail, shared by the submit path
1111/// (`spawn_run`) and the cluster claim path (`resume_claimed_run`). Assumes the
1112/// caller already holds an execution permit and has registered `run_token`.
1113/// `from_queue` is `true` when the run consumed a local queue slot (submit path)
1114/// and `false` for a cluster claim-path run that never reserved one (#228).
1115#[allow(clippy::too_many_arguments)]
1116async fn execute_run(
1117    state: ServerState,
1118    loaded: LoadedSubmission,
1119    run_id: String,
1120    run_token: CancellationToken,
1121    submitted_at: DateTime<Utc>,
1122    timeout_secs: Option<u64>,
1123    clock_flag: Option<String>,
1124    from_queue: bool,
1125) {
1126    let server_shutdown = state.shutdown_token();
1127    let LoadedSubmission { cfg, nodes } = loaded;
1128
1129    // Queued → running. From here the guard guarantees `mark_finished` (and a
1130    // gauge refresh) on EVERY exit, including early returns and panics.
1131    // A submit-path run consumed a local queue slot (Queued→Running); a cluster
1132    // claim-path run never did (#228) — only bump in_flight for it.
1133    if from_queue {
1134        state.registry().mark_running();
1135    } else {
1136        state.registry().mark_running_unqueued();
1137    }
1138    let _guard = InFlightGuard {
1139        state: state.clone(),
1140        run_id: run_id.clone(),
1141    };
1142    let started = Utc::now();
1143    if let Ok(Some(mut rec)) = state.history().get(&run_id).await {
1144        rec.status = RunStatus::Running;
1145        rec.started_at = Some(started);
1146        let _ = state.history().upsert(&rec).await;
1147    }
1148    metrics::set_run_gauges(&state);
1149
1150    // Build execution options (auth/clock failures finalize as Failed).
1151    let pipeline_name = cfg.name.clone().unwrap_or_else(|| "serve".to_string());
1152    let auth = match build_auth_catalog(cfg.auth.as_ref()) {
1153        Ok(a) => a,
1154        Err(e) => {
1155            finalize(
1156                &state,
1157                &run_id,
1158                started,
1159                Terminal::Failed {
1160                    reason: format!("auth catalog: {e}"),
1161                    records: 0,
1162                    invs: Vec::new(),
1163                },
1164            )
1165            .await;
1166            return;
1167        }
1168    };
1169    let clock = match resolve_clock(clock_flag.as_deref(), submitted_at) {
1170        Ok(c) => c,
1171        Err(e) => {
1172            finalize(
1173                &state,
1174                &run_id,
1175                started,
1176                Terminal::Failed {
1177                    reason: e.api_error().error.message,
1178                    records: 0,
1179                    invs: Vec::new(),
1180                },
1181            )
1182            .await;
1183            return;
1184        }
1185    };
1186
1187    // Cooperative-cancel token the pipeline observes so it flushes buffered
1188    // output (e.g. a Parquet footer, an S3 multipart upload) at its next
1189    // page boundary on cancel / timeout / shutdown — instead of having its
1190    // future hard-dropped, which flushes nothing (#146 H16).
1191    let coop = CancellationToken::new();
1192    // Resilience policy from the (merged) submitted config. A malformed
1193    // `resilience:` block finalizes the run as Failed, mirroring the
1194    // auth/clock failure handling above.
1195    let resilience = match &cfg.resilience {
1196        Some(spec) => match spec.to_policy() {
1197            Ok(p) => Some(p),
1198            Err(e) => {
1199                finalize(
1200                    &state,
1201                    &run_id,
1202                    started,
1203                    Terminal::Failed {
1204                        reason: format!("resilience: {e}"),
1205                        records: 0,
1206                        invs: Vec::new(),
1207                    },
1208                )
1209                .await;
1210                return;
1211            }
1212        },
1213        None => None,
1214    };
1215    // Build the per-run OpenLineage emitter from the (merged) submitted
1216    // config. A malformed `lineage:` block finalizes the run as Failed,
1217    // mirroring the auth/clock failure handling above.
1218    #[cfg(feature = "lineage")]
1219    let lineage = match crate::lineage_glue::build_emitter(cfg.lineage.as_ref()) {
1220        Ok(l) => l,
1221        Err(e) => {
1222            finalize(
1223                &state,
1224                &run_id,
1225                started,
1226                Terminal::Failed {
1227                    reason: format!("lineage: {e}"),
1228                    records: 0,
1229                    invs: Vec::new(),
1230                },
1231            )
1232            .await;
1233            return;
1234        }
1235    };
1236    // Notifier from the submitted config's `notifications:` block. A malformed
1237    // block disables notifications for this run (logged) rather than failing an
1238    // already-accepted run.
1239    #[cfg(feature = "notify")]
1240    let notifier = crate::notify::Notifier::from_specs(&cfg.notifications).unwrap_or_else(|e| {
1241        tracing::error!(%run_id, "notifications config invalid, disabling: {e}");
1242        None
1243    });
1244    let opts = ExecuteOptions {
1245        pipeline_name,
1246        // The id returned by `POST /v1/runs`, so a completion notification can be
1247        // matched back to the submission (#480).
1248        run_id: Some(run_id.clone()),
1249        execution: cfg.execution.clone(),
1250        dry_run: false,
1251        limit: None,
1252        state_path_override: None,
1253        shard: None,
1254        auth,
1255        clock,
1256        cancel: Some(coop.clone()),
1257        resilience,
1258        sla: cfg.sla.clone(),
1259        reconcile: cfg.reconcile.clone(),
1260        #[cfg(feature = "lineage")]
1261        lineage,
1262        #[cfg(feature = "lineage")]
1263        lineage_cfg: cfg.lineage.clone(),
1264        #[cfg(feature = "notify")]
1265        notifier,
1266        // Record every serve run into the Data Movement Catalog (#279),
1267        // stored alongside the run history (`--history` backend) and
1268        // attributed to this serve run id.
1269        #[cfg(feature = "catalog")]
1270        catalog: Some(crate::catalog::CatalogHandle {
1271            store: state.history(),
1272            run_id: Some(run_id.clone()),
1273            sample_records: crate::catalog::DEFAULT_SAMPLE_RECORDS,
1274        }),
1275    };
1276
1277    let span = tracing::info_span!("faucet.serve.run", serve_run_id = %run_id);
1278    let work = async move {
1279        // Emitted inside the run span so it is captured by the SSE log layer
1280        // (and gives every `/logs` reader at least one line to anchor on).
1281        tracing::info!("pipeline run starting");
1282        classify_run(run_expanded(nodes, opts).await)
1283    }
1284    .instrument(span);
1285    tokio::pin!(work);
1286
1287    // The run timeout is modelled as a cancel trigger (not a hard
1288    // `tokio::time::timeout` drop) so a timed-out run still flushes.
1289    let timeout_fut = async {
1290        match timeout_secs {
1291            Some(s) => tokio::time::sleep(Duration::from_secs(s)).await,
1292            None => std::future::pending::<()>().await,
1293        }
1294    };
1295    tokio::pin!(timeout_fut);
1296
1297    enum Trigger {
1298        Done(Terminal),
1299        Cancel,
1300        Shutdown,
1301        Timeout(u64),
1302    }
1303
1304    // Phase 1: run to natural completion, or until a cancel trigger fires.
1305    // `biased` prefers a just-completed run over a simultaneous trigger.
1306    let trigger = tokio::select! {
1307        biased;
1308        t = &mut work => Trigger::Done(t),
1309        _ = run_token.cancelled() => Trigger::Cancel,
1310        _ = server_shutdown.cancelled() => Trigger::Shutdown,
1311        _ = &mut timeout_fut => Trigger::Timeout(timeout_secs.unwrap_or(0)),
1312    };
1313
1314    let terminal = match trigger {
1315        Trigger::Done(t) => t,
1316        triggered => {
1317            // Phase 2: a trigger fired. Cancel cooperatively and give the
1318            // pipeline a bounded grace to flush at its next page boundary,
1319            // then hard-drop it (drops the JoinSet, aborting any pipeline
1320            // genuinely stuck mid-write) so a hung run can't wedge shutdown.
1321            coop.cancel();
1322            let trigger_terminal = match triggered {
1323                Trigger::Cancel => Terminal::Cancelled,
1324                Trigger::Shutdown => Terminal::ShutdownFailed,
1325                Trigger::Timeout(secs) => Terminal::Timeout { secs },
1326                Trigger::Done(_) => unreachable!("matched in the outer arm"),
1327            };
1328            // A cooperative cancel makes `run_stream` flush and return Ok(partial),
1329            // so the trigger label (Cancelled/Timeout/ShutdownFailed) is the
1330            // correct status for that path. But if the flush itself FAILS within
1331            // the grace window, surface that real failure — never mask it behind
1332            // the trigger label, which would hide a data error / partial write.
1333            match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
1334                Ok(failed @ Terminal::Failed { .. }) => failed,
1335                Ok(_) | Err(_) => trigger_terminal,
1336            }
1337        }
1338    };
1339
1340    finalize(&state, &run_id, started, terminal).await;
1341    // Signal `/logs` readers the run is done, then drop the buffer after a
1342    // drain window so a late fetcher can still replay it (spec §12).
1343    state.log_hub().finish(&run_id);
1344    schedule_log_drop(state.clone(), run_id.clone());
1345    // `_guard` drops here → mark_finished + gauge refresh.
1346}
1347
1348/// Write the authoritative terminal record + the run-finished metric.
1349async fn finalize(state: &ServerState, run_id: &str, started: DateTime<Utc>, term: Terminal) {
1350    let finished = Utc::now();
1351    let elapsed = (finished - started).to_std().ok().map(|d| d.as_secs_f64());
1352    let (status, reason, records, invs, error) = term.into_parts();
1353    // Read-modify-write the existing record to preserve its metadata
1354    // (name / labels / idempotency_key / submitted_at). If it can't be read —
1355    // the backend errored, or the record was purged / landed in another store
1356    // under degraded fallback — DON'T silently drop the terminal status (#146
1357    // M6): reconstruct a minimal terminal record and upsert it, so the run
1358    // never lingers non-terminal while `record_run_finished` has already fired.
1359    let mut rec = match state.history().get(run_id).await {
1360        Ok(Some(rec)) => rec,
1361        Ok(None) => {
1362            tracing::warn!(
1363                run_id,
1364                "finalize: run record not found; writing a fresh terminal record"
1365            );
1366            RunRecord::queued(run_id.to_string(), None, BTreeMap::new(), None, started)
1367        }
1368        Err(e) => {
1369            tracing::warn!(
1370                run_id,
1371                error = %e,
1372                "finalize: failed to read run record; writing a fresh terminal record"
1373            );
1374            RunRecord::queued(run_id.to_string(), None, BTreeMap::new(), None, started)
1375        }
1376    };
1377    rec.status = status;
1378    rec.started_at.get_or_insert(started);
1379    rec.finished_at = Some(finished);
1380    rec.elapsed_secs = elapsed;
1381    rec.records_written = records;
1382    rec.invocations = invs;
1383    rec.error = error;
1384    if state.cluster().enabled() {
1385        // Owner-fenced: if another instance reclaimed this run (lease expired),
1386        // our write is a no-op and we discard our result — the reclaimer is now
1387        // authoritative (#197).
1388        match state.history().finalize_owned(&rec).await {
1389            Ok(true) => {
1390                metrics::record_run_finished(status, reason);
1391                // Only the instance that won the owner-fenced write fires the
1392                // callback, so a reclaimed run cannot deliver twice (#481).
1393                crate::serve::callback::fire(&rec).await;
1394            }
1395            Ok(false) => tracing::warn!(
1396                run_id,
1397                "finalize: run was reclaimed by another instance; discarding result"
1398            ),
1399            Err(e) => {
1400                tracing::error!(run_id, error = %e, "finalize: owner-fenced write failed")
1401            }
1402        }
1403    } else {
1404        if let Err(e) = state.history().upsert(&rec).await {
1405            tracing::error!(
1406                run_id,
1407                error = %e,
1408                "finalize: failed to persist terminal run record"
1409            );
1410        }
1411        metrics::record_run_finished(status, reason);
1412        crate::serve::callback::fire(&rec).await;
1413    }
1414}
1415
1416/// Finalize a run that was cancelled / hit shutdown while still QUEUED (before
1417/// it acquired an execution permit, so it never became in-flight and has no
1418/// `InFlightGuard`). Releases the queue slot, writes the terminal record, closes
1419/// the log buffer, and refreshes the gauges — the queued-path analogue of the
1420/// normal `finalize` + `InFlightGuard`-drop cleanup.
1421async fn finalize_queued_cancel(
1422    state: &ServerState,
1423    run_id: &str,
1424    submitted_at: DateTime<Utc>,
1425    term: Terminal,
1426) {
1427    state.registry().mark_queued_cancelled(run_id);
1428    finalize(state, run_id, submitted_at, term).await;
1429    state.log_hub().finish(run_id);
1430    schedule_log_drop(state.clone(), run_id.to_string());
1431    metrics::set_run_gauges(state);
1432}
1433
1434/// Spawn a detached timer that drops a finished run's log buffer after the drain
1435/// window, freeing its ring once late `/logs` fetchers have had a chance to read.
1436fn schedule_log_drop(state: ServerState, run_id: String) {
1437    tokio::spawn(async move {
1438        tokio::time::sleep(crate::serve::logs::LOG_DRAIN).await;
1439        state.log_hub().drop_run(&run_id);
1440    });
1441}
1442
1443#[cfg(test)]
1444mod tests {
1445    use super::*;
1446
1447    /// A stand-in admin actor for the submit() tests.
1448    fn admin_actor() -> AuthContext {
1449        AuthContext {
1450            principal: "test".into(),
1451            role: crate::serve::rbac::Role::Admin,
1452            source_ip: None,
1453        }
1454    }
1455
1456    #[test]
1457    fn classify_ok_no_failures_is_completed() {
1458        let summary = RunSummary {
1459            invocations: vec![crate::executor::InvocationOutcome {
1460                row_id: "r".into(),
1461                parent_record_key: None,
1462                records_written: 3,
1463                error: None,
1464                metrics: None,
1465            }],
1466        };
1467        let (status, reason, records, _, error) = classify_run(Ok(summary)).into_parts();
1468        assert_eq!(status, RunStatus::Completed);
1469        assert_eq!(reason, "ok");
1470        assert_eq!(records, 3);
1471        assert!(error.is_none());
1472    }
1473
1474    #[test]
1475    fn classify_ok_with_failures_is_failed() {
1476        let summary = RunSummary {
1477            invocations: vec![crate::executor::InvocationOutcome {
1478                row_id: "r".into(),
1479                parent_record_key: None,
1480                records_written: 0,
1481                error: Some("boom".into()),
1482                metrics: None,
1483            }],
1484        };
1485        let (status, reason, _, _, error) = classify_run(Ok(summary)).into_parts();
1486        assert_eq!(status, RunStatus::Failed);
1487        assert_eq!(reason, "error");
1488        assert!(error.unwrap().contains("invocation(s) failed"));
1489    }
1490
1491    #[test]
1492    fn timeout_maps_to_failed_with_timeout_reason() {
1493        let (status, reason, _, _, error) = Terminal::Timeout { secs: 30 }.into_parts();
1494        assert_eq!(status, RunStatus::Failed);
1495        assert_eq!(reason, "timeout");
1496        assert!(error.unwrap().contains("30s"));
1497    }
1498
1499    #[test]
1500    fn resolve_clock_defaults_and_parses() {
1501        let default = Utc::now();
1502        assert_eq!(
1503            resolve_clock(None, default).unwrap(),
1504            default.fixed_offset()
1505        );
1506        assert!(resolve_clock(Some("2026-01-31T00:00:00Z"), default).is_ok());
1507        assert!(resolve_clock(Some("not-a-time"), default).is_err());
1508    }
1509
1510    #[tokio::test]
1511    async fn conflict_releases_reservation() {
1512        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1513        use crate::serve::history::RunHistory;
1514        use crate::serve::history::memory::MemoryHistory;
1515        use crate::serve::state::ServerState;
1516        use std::sync::Arc;
1517        use tokio_util::sync::CancellationToken;
1518
1519        let cfg = ServeConfig {
1520            listen: "127.0.0.1:0".parse().unwrap(),
1521            auth: AuthMode::None,
1522            max_concurrent_runs: 4,
1523            max_queued_runs: 4,
1524            default_config_path: None,
1525            history: HistoryBackendSpec::Memory,
1526            cors_origins: vec![],
1527            body_limit_bytes: 1_048_576,
1528            shutdown_grace: Duration::from_secs(60),
1529            retain_terminal_runs: Duration::from_secs(60),
1530            idempotency_retention: Duration::from_secs(60),
1531            log_retention: std::time::Duration::from_secs(0),
1532            log_max_lines_per_run: 100_000,
1533            lease_ttl: Duration::from_secs(30),
1534            probe_timeout: Duration::from_secs(10),
1535            env_file: None,
1536            no_env_file: false,
1537            log_level: "info".into(),
1538            ui_enabled: true,
1539            cluster: crate::serve::cluster::ClusterConfig::disabled(),
1540            triggers_path: None,
1541            callback_allow_hosts: Vec::new(),
1542        };
1543        let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1544        let state = ServerState::new(
1545            &cfg,
1546            None,
1547            CancellationToken::new(),
1548            history,
1549            crate::serve::logs::LogHub::new(),
1550            None,
1551            #[cfg(feature = "triggers")]
1552            crate::serve::triggers::health::TriggersHandle::empty(),
1553        );
1554
1555        // Pre-claim the key with a DIFFERENT fingerprint so submit() hits Conflict.
1556        state
1557            .history()
1558            .claim_idempotency("k", "different-fp", "prior", Duration::from_secs(60))
1559            .await
1560            .unwrap();
1561
1562        let req = SubmitRequest {
1563            config: "version: 1\npipeline:\n  source: { type: csv, config: { path: x.csv } }\n  sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1564            config_format: ConfigFormatWire::Yaml,
1565            name: None,
1566            labels: BTreeMap::new(),
1567            timeout_secs: None,
1568            doctor_first: false,
1569            callback: None,
1570            idempotency_key: Some("k".into()),
1571            clock: None,
1572        };
1573
1574        let err = submit(state.clone(), req, admin_actor()).await.unwrap_err();
1575        assert!(
1576            matches!(err, ServeError::Conflict(_)),
1577            "expected Conflict, got {err:?}"
1578        );
1579        // The reservation taken before the claim must have been released by the guard.
1580        assert_eq!(state.registry().queued(), 0);
1581    }
1582
1583    #[tokio::test]
1584    async fn cluster_submit_writes_pending_with_config_and_does_not_spawn() {
1585        use crate::serve::cluster::ClusterConfig;
1586        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1587        use crate::serve::history::RunHistory;
1588        use crate::serve::history::memory::MemoryHistory;
1589        use crate::serve::state::ServerState;
1590        use std::sync::Arc;
1591        use tokio_util::sync::CancellationToken;
1592
1593        let mut cluster = ClusterConfig::disabled();
1594        cluster.enabled = true;
1595        let cfg = ServeConfig {
1596            listen: "127.0.0.1:0".parse().unwrap(),
1597            auth: AuthMode::None,
1598            max_concurrent_runs: 4,
1599            max_queued_runs: 4,
1600            default_config_path: None,
1601            history: HistoryBackendSpec::Memory,
1602            cors_origins: vec![],
1603            body_limit_bytes: 1_048_576,
1604            shutdown_grace: Duration::from_secs(60),
1605            retain_terminal_runs: Duration::from_secs(60),
1606            idempotency_retention: Duration::from_secs(60),
1607            log_retention: std::time::Duration::from_secs(0),
1608            log_max_lines_per_run: 100_000,
1609            lease_ttl: Duration::from_secs(30),
1610            probe_timeout: Duration::from_secs(10),
1611            env_file: None,
1612            no_env_file: false,
1613            log_level: "info".into(),
1614            ui_enabled: true,
1615            cluster,
1616            triggers_path: None,
1617            callback_allow_hosts: Vec::new(),
1618        };
1619        let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1620        let state = ServerState::new(
1621            &cfg,
1622            None,
1623            CancellationToken::new(),
1624            history,
1625            crate::serve::logs::LogHub::new(),
1626            None,
1627            #[cfg(feature = "triggers")]
1628            crate::serve::triggers::health::TriggersHandle::empty(),
1629        );
1630
1631        let req = SubmitRequest {
1632            config: "version: 1\npipeline:\n  source: { type: csv, config: { path: x.csv } }\n  sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1633            config_format: ConfigFormatWire::Yaml,
1634            name: Some("n".into()),
1635            labels: BTreeMap::new(),
1636            timeout_secs: Some(99),
1637            doctor_first: false,
1638            callback: None,
1639            idempotency_key: None,
1640            clock: None,
1641        };
1642        let resp = submit(state.clone(), req, admin_actor()).await.unwrap();
1643        assert_eq!(resp.status, RunStatus::Pending);
1644        // No local queue slot was consumed (cluster runs don't queue locally).
1645        assert_eq!(state.registry().queued(), 0);
1646        let rec = state.history().get(&resp.run_id).await.unwrap().unwrap();
1647        assert_eq!(rec.status, RunStatus::Pending);
1648        assert!(rec.config_body.as_deref().unwrap().contains("version: 1"));
1649        assert_eq!(rec.timeout_secs, Some(99));
1650    }
1651
1652    /// A `ServerState` backed by an in-memory history, for finalize tests.
1653    fn memory_state() -> crate::serve::state::ServerState {
1654        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1655        use crate::serve::history::RunHistory;
1656        use crate::serve::history::memory::MemoryHistory;
1657        use crate::serve::state::ServerState;
1658        use std::sync::Arc;
1659        use tokio_util::sync::CancellationToken;
1660
1661        let cfg = ServeConfig {
1662            listen: "127.0.0.1:0".parse().unwrap(),
1663            auth: AuthMode::None,
1664            max_concurrent_runs: 4,
1665            max_queued_runs: 4,
1666            default_config_path: None,
1667            history: HistoryBackendSpec::Memory,
1668            cors_origins: vec![],
1669            body_limit_bytes: 1_048_576,
1670            shutdown_grace: Duration::from_secs(60),
1671            retain_terminal_runs: Duration::from_secs(60),
1672            idempotency_retention: Duration::from_secs(60),
1673            log_retention: std::time::Duration::from_secs(0),
1674            log_max_lines_per_run: 100_000,
1675            lease_ttl: Duration::from_secs(30),
1676            probe_timeout: Duration::from_secs(10),
1677            env_file: None,
1678            no_env_file: false,
1679            log_level: "info".into(),
1680            ui_enabled: true,
1681            cluster: crate::serve::cluster::ClusterConfig::disabled(),
1682            triggers_path: None,
1683            callback_allow_hosts: Vec::new(),
1684        };
1685        let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1686        ServerState::new(
1687            &cfg,
1688            None,
1689            CancellationToken::new(),
1690            history,
1691            crate::serve::logs::LogHub::new(),
1692            None,
1693            #[cfg(feature = "triggers")]
1694            crate::serve::triggers::health::TriggersHandle::empty(),
1695        )
1696    }
1697
1698    #[tokio::test]
1699    async fn finalize_writes_terminal_record_when_record_is_missing() {
1700        // M6 (#146): if the run record can't be read at finalize time (purged,
1701        // or split to another store under degraded fallback), the terminal
1702        // status must NOT be silently dropped — a fresh terminal record is
1703        // written so the run never lingers non-terminal while the run-finished
1704        // metric has already fired.
1705        let state = memory_state();
1706        let started = Utc::now();
1707        finalize(
1708            &state,
1709            "ghost",
1710            started,
1711            Terminal::Failed {
1712                reason: "boom".into(),
1713                records: 0,
1714                invs: Vec::new(),
1715            },
1716        )
1717        .await;
1718        let rec = state
1719            .history()
1720            .get("ghost")
1721            .await
1722            .unwrap()
1723            .expect("finalize must create a terminal record even when none existed");
1724        assert_eq!(rec.status, RunStatus::Failed);
1725        assert!(rec.finished_at.is_some());
1726        assert!(rec.started_at.is_some());
1727        assert_eq!(rec.error.as_deref(), Some("boom"));
1728    }
1729
1730    #[tokio::test]
1731    async fn finalize_preserves_metadata_of_existing_record() {
1732        // The happy path still read-modify-writes, preserving name/labels/key.
1733        let state = memory_state();
1734        let started = Utc::now();
1735        let mut rec = RunRecord::queued(
1736            "r1".into(),
1737            Some("nightly".into()),
1738            BTreeMap::new(),
1739            Some("idem-k".into()),
1740            started,
1741        );
1742        rec.status = RunStatus::Running;
1743        rec.started_at = Some(started);
1744        state.history().upsert(&rec).await.unwrap();
1745
1746        finalize(
1747            &state,
1748            "r1",
1749            started,
1750            Terminal::Completed {
1751                records: 5,
1752                invs: Vec::new(),
1753            },
1754        )
1755        .await;
1756        let got = state.history().get("r1").await.unwrap().unwrap();
1757        assert_eq!(got.status, RunStatus::Completed);
1758        assert_eq!(got.records_written, 5);
1759        assert_eq!(got.name.as_deref(), Some("nightly"));
1760        assert_eq!(got.idempotency_key.as_deref(), Some("idem-k"));
1761    }
1762
1763    #[cfg(any(feature = "serve-history-sqlite", feature = "serve-history-postgres"))]
1764    #[tokio::test]
1765    async fn cluster_submit_503s_when_history_degraded() {
1766        // #197 spec §9: a degraded backend can't coordinate a cluster, so submit
1767        // must fail closed with 503 rather than orphan a never-claimable Pending run.
1768        use crate::serve::cluster::ClusterConfig;
1769        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1770        use crate::serve::history::RunHistory;
1771        use crate::serve::history::fallback::FallbackHistory;
1772        use crate::serve::state::ServerState;
1773        use std::sync::Arc;
1774        use tokio_util::sync::CancellationToken;
1775
1776        let mut cluster = ClusterConfig::disabled();
1777        cluster.enabled = true;
1778        let cfg = ServeConfig {
1779            listen: "127.0.0.1:0".parse().unwrap(),
1780            auth: AuthMode::None,
1781            max_concurrent_runs: 4,
1782            max_queued_runs: 4,
1783            default_config_path: None,
1784            history: HistoryBackendSpec::Memory,
1785            cors_origins: vec![],
1786            body_limit_bytes: 1_048_576,
1787            shutdown_grace: Duration::from_secs(60),
1788            retain_terminal_runs: Duration::from_secs(60),
1789            idempotency_retention: Duration::from_secs(60),
1790            log_retention: std::time::Duration::from_secs(0),
1791            log_max_lines_per_run: 100_000,
1792            lease_ttl: Duration::from_secs(30),
1793            probe_timeout: Duration::from_secs(10),
1794            env_file: None,
1795            no_env_file: false,
1796            log_level: "info".into(),
1797            ui_enabled: true,
1798            cluster,
1799            triggers_path: None,
1800            callback_allow_hosts: Vec::new(),
1801        };
1802        // A backend that is degraded from startup (primary unreachable).
1803        let history = Arc::new(FallbackHistory::degraded_at_startup(
1804            Duration::from_secs(60),
1805            "test",
1806        )) as Arc<dyn RunHistory>;
1807        assert!(history.degraded());
1808        let state = ServerState::new(
1809            &cfg,
1810            None,
1811            CancellationToken::new(),
1812            history,
1813            crate::serve::logs::LogHub::new(),
1814            None,
1815            #[cfg(feature = "triggers")]
1816            crate::serve::triggers::health::TriggersHandle::empty(),
1817        );
1818        let req = SubmitRequest {
1819            config: "version: 1\npipeline:\n  source: { type: csv, config: { path: x.csv } }\n  sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1820            config_format: ConfigFormatWire::Yaml,
1821            name: None,
1822            labels: BTreeMap::new(),
1823            timeout_secs: None,
1824            doctor_first: false,
1825            callback: None,
1826            idempotency_key: None,
1827            clock: None,
1828        };
1829        let err = submit(state.clone(), req, admin_actor()).await.unwrap_err();
1830        assert!(
1831            matches!(err, ServeError::Unavailable(_)),
1832            "expected 503 Unavailable, got {err:?}"
1833        );
1834        // The queue reservation must have been released (no leak).
1835        assert_eq!(state.registry().queued(), 0);
1836    }
1837
1838    // ── Mode B coverage: coordinator / parent finalize / shard execution ─────
1839    // SQLite-backed so the shard RunHistory methods are live (the memory backend
1840    // is inert for shards). No Docker: the S3 source builds offline and its
1841    // enumerate_shards is pure (hash-modulo); csv→jsonl runs entirely on temp
1842    // files.
1843    #[cfg(feature = "serve-history-sqlite")]
1844    mod shards {
1845        use super::*;
1846        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1847        use crate::serve::history::RunHistory;
1848        use crate::serve::history::sqlite::SqliteHistory;
1849        use crate::serve::load::{ConfigFormat, load_submission};
1850        use crate::serve::state::ServerState;
1851        use faucet_core::ShardSpec;
1852        use std::collections::BTreeMap;
1853        use std::sync::Arc;
1854        use tokio_util::sync::CancellationToken;
1855
1856        async fn sqlite_state(dir: &std::path::Path) -> ServerState {
1857            let url = format!("sqlite://{}/h.db", dir.display());
1858            let history = Arc::new(
1859                SqliteHistory::connect(
1860                    &url,
1861                    Duration::from_secs(300),
1862                    Duration::from_secs(300),
1863                    "inst-test".into(),
1864                )
1865                .await
1866                .expect("sqlite history"),
1867            ) as Arc<dyn RunHistory>;
1868            let cfg = ServeConfig {
1869                listen: "127.0.0.1:0".parse().unwrap(),
1870                auth: AuthMode::None,
1871                max_concurrent_runs: 4,
1872                max_queued_runs: 4,
1873                default_config_path: None,
1874                history: HistoryBackendSpec::Memory,
1875                cors_origins: vec![],
1876                body_limit_bytes: 1_048_576,
1877                shutdown_grace: Duration::from_secs(60),
1878                retain_terminal_runs: Duration::from_secs(60),
1879                idempotency_retention: Duration::from_secs(60),
1880                log_retention: std::time::Duration::from_secs(0),
1881                log_max_lines_per_run: 100_000,
1882                lease_ttl: Duration::from_secs(30),
1883                probe_timeout: Duration::from_secs(10),
1884                env_file: None,
1885                no_env_file: false,
1886                log_level: "info".into(),
1887                ui_enabled: true,
1888                cluster: crate::serve::cluster::ClusterConfig::disabled(),
1889                triggers_path: None,
1890                callback_allow_hosts: Vec::new(),
1891            };
1892            ServerState::new(
1893                &cfg,
1894                None,
1895                CancellationToken::new(),
1896                history,
1897                crate::serve::logs::LogHub::new(),
1898                None,
1899                #[cfg(feature = "triggers")]
1900                crate::serve::triggers::health::TriggersHandle::empty(),
1901            )
1902        }
1903
1904        async fn loaded(yaml: &str) -> LoadedSubmission {
1905            load_submission(yaml, ConfigFormat::Yaml, None)
1906                .await
1907                .expect("load submission")
1908        }
1909
1910        async fn seed_run(state: &ServerState, run_id: &str, status: RunStatus) {
1911            let mut rec = RunRecord::queued(run_id.into(), None, BTreeMap::new(), None, Utc::now());
1912            rec.status = status;
1913            rec.config_body = Some("version: 1".into());
1914            state.history().upsert(&rec).await.expect("seed run");
1915        }
1916
1917        #[tokio::test]
1918        async fn coordinate_matrix_run_is_not_shardable() {
1919            // A matrix expands to >1 node → not shardable → Ok(false), no build.
1920            let dir = tempfile::tempdir().unwrap();
1921            let state = sqlite_state(dir.path()).await;
1922            let l = loaded(
1923                "version: 1\nname: m\nmatrix:\n  - id: a\n  - id: b\npipeline:\n  \
1924                 source: { type: rest, config: { url: \"http://localhost/x\" } }\n  \
1925                 sink: { type: stdout, config: {} }\n",
1926            )
1927            .await;
1928            assert!(!coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1929        }
1930
1931        #[tokio::test]
1932        async fn coordinate_non_shardable_source_runs_whole() {
1933            // A csv source is not shardable → Ok(false) (built offline).
1934            let dir = tempfile::tempdir().unwrap();
1935            let state = sqlite_state(dir.path()).await;
1936            let input = dir.path().join("in.csv");
1937            std::fs::write(&input, "id\n1\n").unwrap();
1938            let l = loaded(&format!(
1939                "version: 1\npipeline:\n  \
1940                 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n  \
1941                 sink: {{ type: stdout, config: {{}} }}\n",
1942                input.display()
1943            ))
1944            .await;
1945            assert!(!coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1946        }
1947
1948        #[tokio::test]
1949        async fn coordinate_s3_source_inserts_shards_and_marks_sharded() {
1950            let dir = tempfile::tempdir().unwrap();
1951            let state = sqlite_state(dir.path()).await;
1952            seed_run(&state, "r", RunStatus::Running).await;
1953            let l = loaded(
1954                "version: 1\npipeline:\n  \
1955                 source: { type: s3, config: { bucket: my-bucket, prefix: null, \
1956                 region: null, endpoint_url: null, file_format: json_lines, \
1957                 max_objects: null, concurrency: 10 } }\n  \
1958                 sink: { type: stdout, config: {} }\n",
1959            )
1960            .await;
1961            assert!(coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1962            // 4 shards inserted; parent flipped to Sharded.
1963            let prog = state.history().shard_progress("r").await.unwrap();
1964            assert_eq!(prog.total, 4);
1965            assert_eq!(prog.pending, 4);
1966            assert_eq!(
1967                state.history().get("r").await.unwrap().unwrap().status,
1968                RunStatus::Sharded
1969            );
1970        }
1971
1972        #[tokio::test]
1973        async fn at_least_once_risky_sinks_flags_append_cluster_and_shard() {
1974            // F26/F39: a clustered or sharded run with an append-mode sink is
1975            // at-least-once → flagged. Neither flag → not flagged.
1976            let append = loaded(
1977                "version: 1\npipeline:\n  \
1978                 source: { type: rest, config: { url: \"http://localhost/x\" } }\n  \
1979                 sink: { type: stdout, config: {} }\n",
1980            )
1981            .await;
1982            // Not clustered, not sharded → no warning.
1983            assert!(at_least_once_risky_sinks(&append, false, false).is_empty());
1984            // Clustered append → flagged.
1985            assert_eq!(
1986                at_least_once_risky_sinks(&append, true, false),
1987                vec!["stdout"]
1988            );
1989            // Sharded append → flagged.
1990            assert_eq!(
1991                at_least_once_risky_sinks(&append, false, true),
1992                vec!["stdout"]
1993            );
1994        }
1995
1996        #[tokio::test]
1997        async fn at_least_once_risky_sinks_safe_for_upsert_and_exactly_once() {
1998            // Upsert sink → keyed re-writes are idempotent → not flagged.
1999            let upsert = loaded(
2000                "version: 1\npipeline:\n  \
2001                 source: { type: postgres, config: { connection_url: \"postgres://x\", \
2002                 query: \"select 1\" } }\n  \
2003                 sink: { type: postgres, config: { connection_url: \"postgres://y\", \
2004                 table_name: t, column_mapping: auto_map, write_mode: upsert, key: [id] } }\n",
2005            )
2006            .await;
2007            assert!(at_least_once_risky_sinks(&upsert, true, true).is_empty());
2008
2009            // exactly_once delivery → not flagged even with an append sink.
2010            let eo = loaded(
2011                "version: 1\ndelivery: exactly_once\npipeline:\n  \
2012                 source: { type: postgres-cdc, config: {} }\n  \
2013                 sink: { type: sqlite, config: {} }\n  \
2014                 state: { type: file, config: { path: \"/tmp/x.json\" } }\n",
2015            )
2016            .await;
2017            assert!(at_least_once_risky_sinks(&eo, true, false).is_empty());
2018        }
2019
2020        async fn seed_sharded_with_shards(state: &ServerState, run_id: &str, n: usize) {
2021            use crate::serve::history::ShardInsert;
2022            seed_run(state, run_id, RunStatus::Sharded).await;
2023            let shards: Vec<ShardInsert> = (0..n)
2024                .map(|i| ShardInsert {
2025                    shard_id: i.to_string(),
2026                    descriptor: serde_json::json!({ "i": i }),
2027                    size_estimate: None,
2028                })
2029                .collect();
2030            state
2031                .history()
2032                .insert_shards(run_id, &shards)
2033                .await
2034                .unwrap();
2035            // Claim them so they are 'running' and finalizable by this instance.
2036            let claimed = state.history().claim_shards(n).await.unwrap();
2037            assert_eq!(claimed.len(), n);
2038        }
2039
2040        #[tokio::test]
2041        async fn maybe_finalize_parent_completes_when_all_shards_succeed() {
2042            let dir = tempfile::tempdir().unwrap();
2043            let state = sqlite_state(dir.path()).await;
2044            seed_sharded_with_shards(&state, "r", 3).await;
2045            for i in 0..3 {
2046                state
2047                    .history()
2048                    .finalize_shard("r", &i.to_string(), true)
2049                    .await
2050                    .unwrap();
2051            }
2052            maybe_finalize_parent(&state, "r").await;
2053            assert_eq!(
2054                state.history().get("r").await.unwrap().unwrap().status,
2055                RunStatus::Completed
2056            );
2057        }
2058
2059        #[tokio::test]
2060        async fn maybe_finalize_parent_fails_when_a_shard_fails() {
2061            let dir = tempfile::tempdir().unwrap();
2062            let state = sqlite_state(dir.path()).await;
2063            seed_sharded_with_shards(&state, "r", 2).await;
2064            state
2065                .history()
2066                .finalize_shard("r", "0", true)
2067                .await
2068                .unwrap();
2069            state
2070                .history()
2071                .finalize_shard("r", "1", false)
2072                .await
2073                .unwrap();
2074            maybe_finalize_parent(&state, "r").await;
2075            assert_eq!(
2076                state.history().get("r").await.unwrap().unwrap().status,
2077                RunStatus::Failed
2078            );
2079        }
2080
2081        #[tokio::test]
2082        async fn maybe_finalize_parent_keeps_sharded_until_all_terminal() {
2083            let dir = tempfile::tempdir().unwrap();
2084            let state = sqlite_state(dir.path()).await;
2085            seed_sharded_with_shards(&state, "r", 2).await;
2086            // Only one shard finalized → run stays Sharded.
2087            state
2088                .history()
2089                .finalize_shard("r", "0", true)
2090                .await
2091                .unwrap();
2092            maybe_finalize_parent(&state, "r").await;
2093            assert_eq!(
2094                state.history().get("r").await.unwrap().unwrap().status,
2095                RunStatus::Sharded
2096            );
2097        }
2098
2099        #[tokio::test]
2100        async fn finalize_sharded_parent_is_status_fenced_and_idempotent() {
2101            // F45: the status-fenced finalize transitions the parent exactly
2102            // once. A concurrent second finalize (e.g. two shards finishing on
2103            // two instances) is a no-op, and a non-Sharded run is never touched.
2104            let dir = tempfile::tempdir().unwrap();
2105            let state = sqlite_state(dir.path()).await;
2106            seed_sharded_with_shards(&state, "r", 2).await;
2107
2108            let first = state
2109                .history()
2110                .finalize_sharded_parent("r", RunStatus::Completed, Utc::now(), None)
2111                .await
2112                .unwrap();
2113            assert!(first, "first finalize wins");
2114            assert_eq!(
2115                state.history().get("r").await.unwrap().unwrap().status,
2116                RunStatus::Completed
2117            );
2118
2119            // Second finalize sees a non-Sharded parent → no-op, status unchanged.
2120            let second = state
2121                .history()
2122                .finalize_sharded_parent("r", RunStatus::Failed, Utc::now(), Some("late".into()))
2123                .await
2124                .unwrap();
2125            assert!(!second, "second finalize is a no-op");
2126            let r = state.history().get("r").await.unwrap().unwrap();
2127            assert_eq!(r.status, RunStatus::Completed, "status not overwritten");
2128            assert!(r.error.is_none(), "late error must not be stamped");
2129
2130            // An unknown / never-sharded run id is not finalized.
2131            let missing = state
2132                .history()
2133                .finalize_sharded_parent("does-not-exist", RunStatus::Completed, Utc::now(), None)
2134                .await
2135                .unwrap();
2136            assert!(!missing);
2137        }
2138
2139        #[tokio::test]
2140        async fn execute_shard_runs_a_csv_to_jsonl_shard() {
2141            let dir = tempfile::tempdir().unwrap();
2142            let state = sqlite_state(dir.path()).await;
2143            let input = dir.path().join("in.csv");
2144            std::fs::write(&input, "id,name\n1,alice\n2,bob\n").unwrap();
2145            let output = dir.path().join("out.jsonl");
2146            let yaml = format!(
2147                "version: 1\npipeline:\n  \
2148                 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n  \
2149                 sink: {{ type: jsonl, config: {{ path: \"{}\" }} }}\n",
2150                input.display(),
2151                output.display()
2152            );
2153            let l = loaded(&yaml).await;
2154            // The whole-dataset shard is a no-op for the (non-shardable) csv source,
2155            // exercising the apply_shard call + per-shard state-key path end-to-end.
2156            let ok = execute_shard(
2157                &state,
2158                l,
2159                "r",
2160                "0",
2161                ShardSpec::whole(),
2162                CancellationToken::new(),
2163                None,
2164                None,
2165                Utc::now(),
2166            )
2167            .await;
2168            assert!(ok, "csv→jsonl shard should complete");
2169            let written = std::fs::read_to_string(&output).unwrap();
2170            assert_eq!(written.lines().count(), 2, "both rows written");
2171            assert!(written.contains("alice") && written.contains("bob"));
2172        }
2173
2174        #[tokio::test]
2175        async fn resume_claimed_shard_executes_and_finalizes_parent() {
2176            // End-to-end per-shard entry point: claim a shard whose run is a
2177            // csv→jsonl pipeline, dispatch it, and confirm the shard runs, is
2178            // finalized, and the parent run flips to Completed.
2179            let dir = tempfile::tempdir().unwrap();
2180            let state = sqlite_state(dir.path()).await;
2181            let input = dir.path().join("in.csv");
2182            std::fs::write(&input, "id,name\n1,alice\n").unwrap();
2183            let output = dir.path().join("out.jsonl");
2184            let yaml = format!(
2185                "version: 1\npipeline:\n  \
2186                 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n  \
2187                 sink: {{ type: jsonl, config: {{ path: \"{}\" }} }}\n",
2188                input.display(),
2189                output.display()
2190            );
2191            // Seed the parent Sharded run carrying the pipeline config, + one shard.
2192            let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2193            rec.status = RunStatus::Sharded;
2194            rec.config_body = Some(yaml);
2195            state.history().upsert(&rec).await.unwrap();
2196            use crate::serve::history::ShardInsert;
2197            state
2198                .history()
2199                .insert_shards(
2200                    "r",
2201                    &[ShardInsert {
2202                        shard_id: "0".into(),
2203                        descriptor: serde_json::Value::Null,
2204                        size_estimate: None,
2205                    }],
2206                )
2207                .await
2208                .unwrap();
2209            let claimed = state.history().claim_shards(1).await.unwrap();
2210            assert_eq!(claimed.len(), 1);
2211
2212            resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2213
2214            // Poll until the parent run reaches a terminal state (the spawned
2215            // task runs the shard, finalizes it, then finalizes the parent).
2216            let mut status = RunStatus::Sharded;
2217            for _ in 0..100 {
2218                tokio::time::sleep(Duration::from_millis(50)).await;
2219                status = state.history().get("r").await.unwrap().unwrap().status;
2220                if status.is_terminal() {
2221                    break;
2222                }
2223            }
2224            assert_eq!(status, RunStatus::Completed, "shard ran → parent completed");
2225            assert!(output.exists(), "shard wrote its output");
2226        }
2227
2228        #[tokio::test]
2229        async fn resume_claimed_shard_with_no_config_fails_the_shard() {
2230            // A claimed shard whose parent run has no config_body can't run →
2231            // the shard is finalized failed and the parent run fails.
2232            let dir = tempfile::tempdir().unwrap();
2233            let state = sqlite_state(dir.path()).await;
2234            let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2235            rec.status = RunStatus::Sharded; // config_body intentionally None
2236            state.history().upsert(&rec).await.unwrap();
2237            use crate::serve::history::ShardInsert;
2238            state
2239                .history()
2240                .insert_shards(
2241                    "r",
2242                    &[ShardInsert {
2243                        shard_id: "0".into(),
2244                        descriptor: serde_json::Value::Null,
2245                        size_estimate: None,
2246                    }],
2247                )
2248                .await
2249                .unwrap();
2250            let claimed = state.history().claim_shards(1).await.unwrap();
2251            resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2252
2253            let mut status = RunStatus::Sharded;
2254            for _ in 0..100 {
2255                tokio::time::sleep(Duration::from_millis(50)).await;
2256                status = state.history().get("r").await.unwrap().unwrap().status;
2257                if status.is_terminal() {
2258                    break;
2259                }
2260            }
2261            assert_eq!(status, RunStatus::Failed, "no-config shard → parent failed");
2262        }
2263
2264        #[tokio::test]
2265        async fn coordinate_returns_err_when_source_build_fails() {
2266            // A shardable-looking source with an invalid config fails to build →
2267            // coordinate_sharded_run surfaces the error (the caller fails the run).
2268            let dir = tempfile::tempdir().unwrap();
2269            let state = sqlite_state(dir.path()).await;
2270            // s3 missing the required `file_format` field → build_source errors.
2271            let l = loaded(
2272                "version: 1\npipeline:\n  \
2273                 source: { type: s3, config: { bucket: b } }\n  \
2274                 sink: { type: stdout, config: {} }\n",
2275            )
2276            .await;
2277            assert!(coordinate_sharded_run(&state, "r", &l, 4).await.is_err());
2278        }
2279
2280        #[tokio::test]
2281        async fn execute_shard_returns_false_on_malformed_resilience() {
2282            // A resilience block that parses but fails to compile makes
2283            // execute_shard fail fast (false) before running the pipeline.
2284            let dir = tempfile::tempdir().unwrap();
2285            let state = sqlite_state(dir.path()).await;
2286            let input = dir.path().join("in.csv");
2287            std::fs::write(&input, "id\n1\n").unwrap();
2288            let yaml = format!(
2289                "version: 1\nresilience:\n  retry:\n    max_attempts: 0\npipeline:\n  \
2290                 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n  \
2291                 sink: {{ type: stdout, config: {{}} }}\n",
2292                input.display()
2293            );
2294            let l = loaded(&yaml).await;
2295            let ok = execute_shard(
2296                &state,
2297                l,
2298                "r",
2299                "0",
2300                ShardSpec::whole(),
2301                CancellationToken::new(),
2302                None,
2303                None,
2304                Utc::now(),
2305            )
2306            .await;
2307            assert!(!ok, "malformed resilience → shard fails fast");
2308        }
2309
2310        #[tokio::test]
2311        async fn resume_claimed_shard_with_unloadable_config_fails_the_shard() {
2312            // A claimed shard whose run config fails to re-load (malformed body)
2313            // exercises resume_claimed_shard's load-error branch → shard failed.
2314            let dir = tempfile::tempdir().unwrap();
2315            let state = sqlite_state(dir.path()).await;
2316            let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2317            rec.status = RunStatus::Sharded;
2318            rec.config_body = Some("this: is: not: valid: yaml: [".into());
2319            state.history().upsert(&rec).await.unwrap();
2320            use crate::serve::history::ShardInsert;
2321            state
2322                .history()
2323                .insert_shards(
2324                    "r",
2325                    &[ShardInsert {
2326                        shard_id: "0".into(),
2327                        descriptor: serde_json::Value::Null,
2328                        size_estimate: None,
2329                    }],
2330                )
2331                .await
2332                .unwrap();
2333            let claimed = state.history().claim_shards(1).await.unwrap();
2334            resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2335
2336            let mut status = RunStatus::Sharded;
2337            for _ in 0..100 {
2338                tokio::time::sleep(Duration::from_millis(50)).await;
2339                status = state.history().get("r").await.unwrap().unwrap().status;
2340                if status.is_terminal() {
2341                    break;
2342                }
2343            }
2344            assert_eq!(
2345                status,
2346                RunStatus::Failed,
2347                "unloadable config → parent failed"
2348            );
2349        }
2350
2351        // ── F10: cross-instance cancel of a Sharded run ──────────────────────
2352
2353        #[tokio::test]
2354        async fn request_cancel_flags_a_sharded_parent() {
2355            // F10: a Sharded parent (status 'sharded') must accept request_cancel
2356            // (the guard was broadened from 'running' to ('running','sharded')).
2357            let dir = tempfile::tempdir().unwrap();
2358            let state = sqlite_state(dir.path()).await;
2359            seed_run(&state, "r", RunStatus::Sharded).await;
2360            // request_cancel is fire-and-forget; verify it took effect by reading
2361            // back via pending_shard_cancellations after a shard is claimed below.
2362            state.history().request_cancel("r").await.unwrap();
2363            // Insert + claim a shard so this instance owns a running shard under r.
2364            use crate::serve::history::ShardInsert;
2365            state
2366                .history()
2367                .insert_shards(
2368                    "r",
2369                    &[ShardInsert {
2370                        shard_id: "0".into(),
2371                        descriptor: serde_json::Value::Null,
2372                        size_estimate: None,
2373                    }],
2374                )
2375                .await
2376                .unwrap();
2377            let claimed = state.history().claim_shards(1).await.unwrap();
2378            assert_eq!(claimed.len(), 1, "shard claimed (running, owned)");
2379
2380            let flagged = state.history().pending_shard_cancellations().await.unwrap();
2381            assert_eq!(
2382                flagged,
2383                vec!["r".to_string()],
2384                "the flagged sharded parent's run id is returned for its running shard"
2385            );
2386        }
2387
2388        #[tokio::test]
2389        async fn pending_shard_cancellations_filters_unflagged_and_pending_shards() {
2390            let dir = tempfile::tempdir().unwrap();
2391            let state = sqlite_state(dir.path()).await;
2392            use crate::serve::history::ShardInsert;
2393            let one = |id: &str| {
2394                vec![ShardInsert {
2395                    shard_id: id.into(),
2396                    descriptor: serde_json::Value::Null,
2397                    size_estimate: None,
2398                }]
2399            };
2400
2401            // Runs A (flagged) and C (NOT flagged) each get one shard; claim both
2402            // so they are running+owned here.
2403            seed_run(&state, "A", RunStatus::Sharded).await;
2404            state.history().request_cancel("A").await.unwrap();
2405            state.history().insert_shards("A", &one("0")).await.unwrap();
2406            seed_run(&state, "C", RunStatus::Sharded).await;
2407            state.history().insert_shards("C", &one("0")).await.unwrap();
2408            let claimed = state.history().claim_shards(8).await.unwrap();
2409            assert_eq!(claimed.len(), 2, "A and C shards claimed (running)");
2410
2411            // Run B: flagged, but its shard stays PENDING (inserted after the
2412            // claim, never claimed) → must be excluded (the join requires a
2413            // 'running' shard owned by this instance).
2414            seed_run(&state, "B", RunStatus::Sharded).await;
2415            state.history().request_cancel("B").await.unwrap();
2416            state.history().insert_shards("B", &one("0")).await.unwrap();
2417
2418            let flagged = state.history().pending_shard_cancellations().await.unwrap();
2419            assert_eq!(
2420                flagged,
2421                vec!["A".to_string()],
2422                "only A (flagged + a running owned shard); B pending-shard, C unflagged"
2423            );
2424        }
2425
2426        // ── F11: orphaned-Sharded-parent sweep ───────────────────────────────
2427
2428        #[tokio::test]
2429        async fn finalize_sweep_completes_an_all_success_sharded_parent() {
2430            let dir = tempfile::tempdir().unwrap();
2431            let state = sqlite_state(dir.path()).await;
2432            seed_sharded_with_shards(&state, "r", 3).await;
2433            for i in 0..3 {
2434                state
2435                    .history()
2436                    .finalize_shard("r", &i.to_string(), true)
2437                    .await
2438                    .unwrap();
2439            }
2440            // No inline maybe_finalize_parent — the sweep must finalize it.
2441            let n = state
2442                .history()
2443                .finalize_completed_sharded_parents()
2444                .await
2445                .unwrap();
2446            assert_eq!(n, 1, "one sharded parent finalized");
2447            let rec = state.history().get("r").await.unwrap().unwrap();
2448            assert_eq!(rec.status, RunStatus::Completed);
2449            assert!(rec.finished_at.is_some());
2450            assert!(rec.error.is_none());
2451
2452            // Idempotent: a second sweep finalizes nothing.
2453            assert_eq!(
2454                state
2455                    .history()
2456                    .finalize_completed_sharded_parents()
2457                    .await
2458                    .unwrap(),
2459                0,
2460                "already-terminal parent is not re-finalized"
2461            );
2462        }
2463
2464        #[tokio::test]
2465        async fn finalize_sweep_fails_a_parent_with_a_failed_shard() {
2466            let dir = tempfile::tempdir().unwrap();
2467            let state = sqlite_state(dir.path()).await;
2468            seed_sharded_with_shards(&state, "r", 3).await;
2469            state
2470                .history()
2471                .finalize_shard("r", "0", true)
2472                .await
2473                .unwrap();
2474            state
2475                .history()
2476                .finalize_shard("r", "1", false)
2477                .await
2478                .unwrap();
2479            state
2480                .history()
2481                .finalize_shard("r", "2", true)
2482                .await
2483                .unwrap();
2484            let n = state
2485                .history()
2486                .finalize_completed_sharded_parents()
2487                .await
2488                .unwrap();
2489            assert_eq!(n, 1);
2490            let rec = state.history().get("r").await.unwrap().unwrap();
2491            assert_eq!(rec.status, RunStatus::Failed);
2492            assert!(rec.finished_at.is_some());
2493            assert_eq!(rec.error.as_deref(), Some("1/3 shard(s) failed"));
2494        }
2495
2496        #[tokio::test]
2497        async fn finalize_sweep_leaves_a_not_all_terminal_parent_sharded() {
2498            let dir = tempfile::tempdir().unwrap();
2499            let state = sqlite_state(dir.path()).await;
2500            seed_sharded_with_shards(&state, "r", 2).await;
2501            // Only one shard terminal → the parent stays sharded.
2502            state
2503                .history()
2504                .finalize_shard("r", "0", true)
2505                .await
2506                .unwrap();
2507            let n = state
2508                .history()
2509                .finalize_completed_sharded_parents()
2510                .await
2511                .unwrap();
2512            assert_eq!(n, 0, "parent with a still-running shard is not finalized");
2513            assert_eq!(
2514                state.history().get("r").await.unwrap().unwrap().status,
2515                RunStatus::Sharded
2516            );
2517        }
2518    }
2519}