1use 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::state::ServerState;
15use crate::serve::{idempotency, metrics};
16use chrono::{DateTime, FixedOffset, Utc};
17use serde::{Deserialize, Serialize};
18use std::collections::BTreeMap;
19use std::time::Duration;
20use tokio_util::sync::CancellationToken;
21use tracing::Instrument;
22
23const QUEUE_FULL_RETRY_AFTER_SECS: u64 = 5;
25
26const RUN_FLUSH_GRACE: Duration = Duration::from_secs(30);
32
33#[derive(Debug, Deserialize)]
35pub struct SubmitRequest {
36 pub config: String,
37 #[serde(default)]
38 pub config_format: ConfigFormatWire,
39 pub name: Option<String>,
40 #[serde(default)]
41 pub labels: BTreeMap<String, String>,
42 pub timeout_secs: Option<u64>,
43 #[serde(default)]
44 pub doctor_first: bool,
45 pub idempotency_key: Option<String>,
46 pub clock: Option<String>,
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)]
51#[serde(rename_all = "lowercase")]
52pub enum ConfigFormatWire {
53 #[default]
54 Yaml,
55 Json,
56}
57
58impl From<ConfigFormatWire> for ConfigFormat {
59 fn from(w: ConfigFormatWire) -> Self {
60 match w {
61 ConfigFormatWire::Yaml => ConfigFormat::Yaml,
62 ConfigFormatWire::Json => ConfigFormat::Json,
63 }
64 }
65}
66
67#[derive(Debug, Serialize)]
69pub struct SubmitResponse {
70 pub run_id: String,
71 pub status: RunStatus,
72 pub submitted_at: DateTime<Utc>,
73}
74
75pub fn resume_claimed_run(state: ServerState, rec: RunRecord) {
79 tokio::spawn(async move {
80 let run_id = rec.run_id.clone();
81 let Some(body) = rec.config_body.as_deref() else {
82 tracing::error!(run_id, "claimed run has no stored config; failing it");
83 finalize(
84 &state,
85 &run_id,
86 rec.submitted_at,
87 Terminal::Failed {
88 reason: "claimed run record missing config_body".into(),
89 records: 0,
90 invs: Vec::new(),
91 },
92 )
93 .await;
94 return;
95 };
96 let format = rec.config_format.unwrap_or_default();
97 let loaded = match load_submission(body, format, state.default_base()).await {
98 Ok(l) => l,
99 Err(e) => {
100 finalize(
101 &state,
102 &run_id,
103 rec.submitted_at,
104 Terminal::Failed {
105 reason: format!(
106 "re-loading claimed config: {}",
107 e.api_error().error.message
108 ),
109 records: 0,
110 invs: Vec::new(),
111 },
112 )
113 .await;
114 return;
115 }
116 };
117
118 if let Some(sh) = loaded.cfg.shard.clone()
123 && sh.count >= 2
124 {
125 match coordinate_sharded_run(&state, &run_id, &loaded, sh.count).await {
126 Ok(true) => return, Ok(false) => {} Err(e) => {
129 finalize(
130 &state,
131 &run_id,
132 rec.submitted_at,
133 Terminal::Failed {
134 reason: format!("sharding: {e}"),
135 records: 0,
136 invs: Vec::new(),
137 },
138 )
139 .await;
140 return;
141 }
142 }
143 }
144
145 let _permit = state
148 .semaphore()
149 .acquire_owned()
150 .await
151 .expect("semaphore not closed");
152 let run_token = CancellationToken::new();
155 state.registry().register(run_id.clone(), run_token.clone());
156 execute_run(
157 state.clone(),
158 loaded,
159 run_id,
160 run_token,
161 rec.submitted_at,
162 rec.timeout_secs,
163 rec.clock.clone(),
164 false,
165 )
166 .await;
167 });
168}
169
170async fn coordinate_sharded_run(
178 state: &ServerState,
179 run_id: &str,
180 loaded: &LoadedSubmission,
181 count: usize,
182) -> crate::error::CliResult<bool> {
183 use crate::error::CliError;
184
185 if loaded.nodes.len() != 1 {
188 tracing::warn!(
189 run_id,
190 nodes = loaded.nodes.len(),
191 "shard requested but the run is not a single-node pipeline; running it whole"
192 );
193 return Ok(false);
194 }
195 let node = &loaded.nodes[0];
196 let auth = build_auth_catalog(loaded.cfg.auth.as_ref())
197 .map_err(|e| CliError::Internal(format!("auth catalog: {e}")))?;
198 let source = build_source(&node.source.kind, node.source.config.clone(), &auth, None).await?;
199 if !source.is_shardable() {
200 tracing::warn!(
201 run_id,
202 kind = %node.source.kind,
203 "source is not shardable; running the run whole"
204 );
205 return Ok(false);
206 }
207
208 {
223 let mut r = state
224 .history()
225 .get(run_id)
226 .await
227 .map_err(|e| CliError::Internal(e.to_string()))?
228 .ok_or_else(|| CliError::Internal(format!("run {run_id} vanished before sharding")))?;
229 r.status = RunStatus::Sharded;
230 state
231 .history()
232 .upsert(&r)
233 .await
234 .map_err(|e| CliError::Internal(e.to_string()))?;
235 }
236
237 let shards = source
238 .enumerate_shards(count)
239 .await
240 .map_err(|e| CliError::Internal(format!("enumerate_shards: {e}")))?;
241 let inserts: Vec<ShardInsert> = shards
242 .iter()
243 .map(|s| ShardInsert {
244 shard_id: s.id.clone(),
245 descriptor: s.descriptor.clone(),
246 size_estimate: s.size_estimate,
247 })
248 .collect();
249 let inserted = state
250 .history()
251 .insert_shards(run_id, &inserts)
252 .await
253 .map_err(|e| CliError::Internal(e.to_string()))?;
254 tracing::info!(
255 run_id,
256 shards = inserts.len(),
257 inserted,
258 "expanded run into shards (Mode B)"
259 );
260
261 state.cluster().kick();
263 Ok(true)
264}
265
266pub fn resume_claimed_shard(state: ServerState, claimed: ClaimedShard) {
270 tokio::spawn(async move {
271 let ClaimedShard {
272 run_id,
273 shard_id,
274 descriptor,
275 run,
276 } = claimed;
277
278 let Some(body) = run.config_body.clone() else {
279 tracing::error!(run_id, shard_id, "claimed shard's run has no stored config");
280 let _ = state
281 .history()
282 .finalize_shard(&run_id, &shard_id, false)
283 .await;
284 maybe_finalize_parent(&state, &run_id).await;
285 return;
286 };
287 let format = run.config_format.unwrap_or_default();
288 let loaded = match load_submission(&body, format, state.default_base()).await {
289 Ok(l) => l,
290 Err(e) => {
291 tracing::error!(
292 run_id,
293 shard_id,
294 error = %e.api_error().error.message,
295 "re-loading shard config failed"
296 );
297 let _ = state
298 .history()
299 .finalize_shard(&run_id, &shard_id, false)
300 .await;
301 maybe_finalize_parent(&state, &run_id).await;
302 return;
303 }
304 };
305
306 let _permit = state
307 .semaphore()
308 .acquire_owned()
309 .await
310 .expect("semaphore not closed");
311
312 let shard = faucet_core::ShardSpec {
313 id: shard_id.clone(),
314 descriptor,
315 size_estimate: None,
316 };
317 let coop = CancellationToken::new();
324 state
325 .registry()
326 .register_shard(&run_id, &shard_id, coop.clone());
327 let success = execute_shard(
328 &state,
329 loaded,
330 &run_id,
331 &shard_id,
332 shard,
333 coop,
334 run.timeout_secs,
335 run.clock.clone(),
336 run.submitted_at,
337 )
338 .await;
339 state.registry().deregister_shard(&run_id, &shard_id);
340
341 match state
342 .history()
343 .finalize_shard(&run_id, &shard_id, success)
344 .await
345 {
346 Ok(true) => {}
347 Ok(false) => tracing::warn!(
348 run_id,
349 shard_id,
350 "shard was reclaimed by another instance; discarding result"
351 ),
352 Err(e) => tracing::error!(run_id, shard_id, error = %e, "finalize_shard failed"),
353 }
354 maybe_finalize_parent(&state, &run_id).await;
355 });
356}
357
358#[allow(clippy::too_many_arguments)]
362async fn execute_shard(
363 state: &ServerState,
364 loaded: LoadedSubmission,
365 run_id: &str,
366 shard_id: &str,
367 shard: faucet_core::ShardSpec,
368 coop: CancellationToken,
369 timeout_secs: Option<u64>,
370 clock_flag: Option<String>,
371 submitted_at: DateTime<Utc>,
372) -> bool {
373 let LoadedSubmission { cfg, nodes } = loaded;
374 let pipeline_name = cfg.name.clone().unwrap_or_else(|| "serve".to_string());
375
376 let auth = match build_auth_catalog(cfg.auth.as_ref()) {
377 Ok(a) => a,
378 Err(e) => {
379 tracing::error!(run_id, shard_id, "shard auth catalog: {e}");
380 return false;
381 }
382 };
383 let clock = match resolve_clock(clock_flag.as_deref(), submitted_at) {
384 Ok(c) => c,
385 Err(e) => {
386 tracing::error!(
387 run_id,
388 shard_id,
389 "shard clock: {}",
390 e.api_error().error.message
391 );
392 return false;
393 }
394 };
395 let resilience = match &cfg.resilience {
396 Some(spec) => match spec.to_policy() {
397 Ok(p) => Some(p),
398 Err(e) => {
399 tracing::error!(run_id, shard_id, "shard resilience: {e}");
400 return false;
401 }
402 },
403 None => None,
404 };
405 #[cfg(feature = "lineage")]
406 let lineage = match crate::lineage_glue::build_emitter(cfg.lineage.as_ref()) {
407 Ok(l) => l,
408 Err(e) => {
409 tracing::error!(run_id, shard_id, "shard lineage: {e}");
410 return false;
411 }
412 };
413
414 let opts = ExecuteOptions {
418 pipeline_name,
419 execution: cfg.execution.clone(),
420 dry_run: false,
421 limit: None,
422 state_path_override: None,
423 shard: Some(shard),
424 auth,
425 clock,
426 cancel: Some(coop.clone()),
427 resilience,
428 #[cfg(feature = "lineage")]
429 lineage,
430 #[cfg(feature = "lineage")]
431 lineage_cfg: cfg.lineage.clone(),
432 };
433
434 let server_shutdown = state.shutdown_token();
435 let span = tracing::info_span!("faucet.serve.shard", serve_run_id = %run_id, shard = %shard_id);
436 let work = async move { classify_run(run_expanded(nodes, opts).await) }.instrument(span);
437 tokio::pin!(work);
438 let timeout_fut = async {
439 match timeout_secs {
440 Some(s) => tokio::time::sleep(Duration::from_secs(s)).await,
441 None => std::future::pending::<()>().await,
442 }
443 };
444 tokio::pin!(timeout_fut);
445
446 let terminal = tokio::select! {
454 biased;
455 t = &mut work => t,
456 _ = coop.cancelled() => {
457 match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
461 Ok(failed @ Terminal::Failed { .. }) => failed,
462 Ok(_) | Err(_) => Terminal::Cancelled,
463 }
464 }
465 _ = server_shutdown.cancelled() => {
466 coop.cancel();
467 match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
468 Ok(failed @ Terminal::Failed { .. }) => failed,
469 Ok(_) | Err(_) => Terminal::ShutdownFailed,
470 }
471 }
472 _ = &mut timeout_fut => {
473 coop.cancel();
474 match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
475 Ok(failed @ Terminal::Failed { .. }) => failed,
476 Ok(_) | Err(_) => Terminal::Timeout { secs: timeout_secs.unwrap_or(0) },
477 }
478 }
479 };
480 matches!(terminal, Terminal::Completed { .. })
481}
482
483async fn maybe_finalize_parent(state: &ServerState, run_id: &str) {
488 let progress = match state.history().shard_progress(run_id).await {
489 Ok(p) => p,
490 Err(e) => {
491 tracing::warn!(run_id, error = %e, "shard_progress failed");
492 return;
493 }
494 };
495 if !progress.all_terminal() {
496 return;
497 }
498 let success = progress.failed == 0;
499 let status = if success {
500 RunStatus::Completed
501 } else {
502 RunStatus::Failed
503 };
504 let error =
505 (!success).then(|| format!("{}/{} shard(s) failed", progress.failed, progress.total));
506 match state
511 .history()
512 .finalize_sharded_parent(run_id, status, Utc::now(), error)
513 .await
514 {
515 Ok(true) => {
516 metrics::record_run_finished(status, if success { "ok" } else { "error" });
517 tracing::info!(
518 run_id,
519 shards = progress.total,
520 failed = progress.failed,
521 "sharded run finalized"
522 );
523 }
524 Ok(false) => {} Err(e) => {
526 tracing::error!(run_id, error = %e, "finalizing sharded parent run failed");
527 }
528 }
529}
530
531fn at_least_once_risky_sinks(
545 loaded: &LoadedSubmission,
546 clustered: bool,
547 sharded: bool,
548) -> Vec<&str> {
549 if !(clustered || sharded) || loaded.cfg.delivery == faucet_core::DeliveryMode::ExactlyOnce {
550 return Vec::new();
551 }
552 loaded
553 .nodes
554 .iter()
555 .filter(|n| {
556 !matches!(
557 n.sink
558 .config
559 .get("write_mode")
560 .and_then(|v| v.as_str())
561 .unwrap_or("append"),
562 "upsert" | "delete"
563 )
564 })
565 .map(|n| n.sink.kind.as_str())
566 .collect()
567}
568
569fn warn_if_cluster_at_least_once(loaded: &LoadedSubmission, clustered: bool, sharded: bool) {
570 let risky = at_least_once_risky_sinks(loaded, clustered, sharded);
571 if !risky.is_empty() {
572 let scope = if sharded {
573 "source-sharded (Mode B)"
574 } else {
575 "clustered"
576 };
577 tracing::warn!(
578 sinks = ?risky,
579 "{scope} execution is at-least-once: a failover or shard reclaim can re-run work \
580 and write duplicate rows to an append-mode sink. Set `write_mode: upsert` (or \
581 `delivery: exactly_once`) on the destination to make re-execution idempotent (F26/F39)."
582 );
583 }
584}
585
586async fn release_orphaned_claim(state: &ServerState, req: &SubmitRequest, run_id: &str) {
590 if req.idempotency_key.is_some()
591 && let Err(e) = state.history().release_idempotency(run_id).await
592 {
593 tracing::warn!(
594 run_id,
595 error = %e,
596 "failed to release idempotency claim after a run-record write error; \
597 a replay of the key may 404 until the claim self-expires"
598 );
599 }
600}
601
602pub async fn submit(state: ServerState, req: SubmitRequest) -> Result<SubmitResponse, ServeError> {
604 let format: ConfigFormat = req.config_format.into();
605 let loaded = load_submission(&req.config, format, state.default_base()).await?;
606
607 let sharded = loaded.cfg.shard.as_ref().is_some_and(|s| s.count >= 2);
610 warn_if_cluster_at_least_once(&loaded, state.cluster().enabled(), sharded);
611
612 if !state.registry().try_reserve() {
615 return Err(ServeError::QueueFull {
616 retry_after_secs: QUEUE_FULL_RETRY_AFTER_SECS,
617 });
618 }
619 let reservation = ReservationGuard::new(state.clone());
622
623 let doctor_report = if req.doctor_first {
629 Some(run_doctor_first(&state, &loaded).await?)
630 } else {
631 None
632 };
633
634 let run_id = uuid::Uuid::now_v7().to_string();
635
636 if let Some(key) = &req.idempotency_key {
638 let merged = serde_json::to_value(&loaded.cfg).unwrap_or(serde_json::Value::Null);
639 let fp_config = idempotency::fingerprint(&merged, loaded.cfg.name.as_deref());
644 let fp = idempotency::request_fingerprint(
645 &fp_config,
646 req.clock.as_deref(),
647 req.timeout_secs,
648 &req.labels,
649 );
650 match state
651 .history()
652 .claim_idempotency(key, &fp, &run_id, state.idempotency_retention())
653 .await
654 .map_err(|e| match e {
655 crate::serve::history::HistoryError::Degraded(m) => ServeError::Unavailable(m),
657 other => ServeError::Internal(other.to_string()),
658 })? {
659 Claim::Fresh => {}
660 Claim::Replay(existing) => {
661 metrics::record_idempotency_hit();
662 return replay_response(&state, &existing).await;
663 }
664 Claim::Conflict => {
665 return Err(ServeError::Conflict(
666 "idempotency key reused with a different payload".into(),
667 ));
668 }
669 }
670 }
675
676 let submitted_at = Utc::now();
677 let mut rec = RunRecord::queued(
678 run_id.clone(),
679 req.name.clone(),
680 req.labels.clone(),
681 req.idempotency_key.clone(),
682 submitted_at,
683 );
684 rec.doctor_report = doctor_report;
685
686 if state.cluster().enabled() {
687 if state.history().degraded() {
692 return Err(ServeError::Unavailable(
693 "clustered run-history backend is degraded; runs cannot be claimed \
694 by any instance — retry once it recovers"
695 .into(),
696 ));
697 }
698 rec.status = RunStatus::Pending;
702 rec.config_body = Some(req.config.clone());
703 rec.config_format = Some(req.config_format.into());
704 rec.timeout_secs = req.timeout_secs;
705 rec.clock = req.clock.clone();
706 if let Err(e) = state.history().upsert(&rec).await {
707 release_orphaned_claim(&state, &req, &run_id).await;
711 return Err(ServeError::Internal(e.to_string()));
712 }
713 drop(reservation);
716 state.cluster().kick();
717 return Ok(SubmitResponse {
718 run_id,
719 status: RunStatus::Pending,
720 submitted_at,
721 });
722 }
723
724 if let Err(e) = state.history().upsert(&rec).await {
725 release_orphaned_claim(&state, &req, &run_id).await;
727 return Err(ServeError::Internal(e.to_string()));
728 }
729
730 let run_token = CancellationToken::new();
731 state.registry().register(run_id.clone(), run_token.clone());
732 metrics::set_run_gauges(&state);
733
734 reservation.defuse();
736 spawn_run(
737 state.clone(),
738 loaded,
739 req,
740 run_id.clone(),
741 run_token,
742 submitted_at,
743 );
744
745 Ok(SubmitResponse {
746 run_id,
747 status: RunStatus::Queued,
748 submitted_at,
749 })
750}
751
752pub(crate) async fn run_doctor_first(
757 state: &ServerState,
758 loaded: &LoadedSubmission,
759) -> Result<serde_json::Value, ServeError> {
760 use faucet_core::check::CheckContext;
761 let auth =
762 build_auth_catalog(loaded.cfg.auth.as_ref()).map_err(|e| ServeError::Unprocessable {
763 message: e.to_string(),
764 details: None,
765 })?;
766 let ctx = CheckContext {
767 timeout: state.probe_timeout(),
768 };
769 let mut invs = crate::commands::doctor::probe_roots(&loaded.nodes, &auth, &ctx).await;
770 let failed = crate::commands::doctor::count_failures(&invs);
771 crate::commands::doctor::redact_invocations(&mut invs);
774 let report = serde_json::json!({ "invocations": invs });
775 if failed > 0 {
776 return Err(ServeError::Unprocessable {
777 message: format!("doctor_first preflight failed: {failed} probe(s) failed"),
778 details: Some(report),
779 });
780 }
781 Ok(report)
782}
783
784async fn replay_response(state: &ServerState, run_id: &str) -> Result<SubmitResponse, ServeError> {
786 let rec = state
787 .history()
788 .get(run_id)
789 .await
790 .map_err(|e| ServeError::Internal(e.to_string()))?
791 .ok_or(ServeError::NotFound)?;
792 Ok(SubmitResponse {
793 run_id: rec.run_id,
794 status: rec.status,
795 submitted_at: rec.submitted_at,
796 })
797}
798
799struct ReservationGuard {
805 state: Option<ServerState>,
806}
807
808impl ReservationGuard {
809 fn new(state: ServerState) -> Self {
810 Self { state: Some(state) }
811 }
812
813 fn defuse(mut self) {
815 self.state = None;
816 }
817}
818
819impl Drop for ReservationGuard {
820 fn drop(&mut self) {
821 if let Some(state) = self.state.take() {
822 state.registry().release_reservation();
823 metrics::set_run_gauges(&state);
824 }
825 }
826}
827
828struct InFlightGuard {
833 state: ServerState,
834 run_id: String,
835}
836
837impl Drop for InFlightGuard {
838 fn drop(&mut self) {
839 self.state.registry().mark_finished(&self.run_id);
840 metrics::set_run_gauges(&self.state);
841 }
842}
843
844enum Terminal {
846 Completed {
847 records: u64,
848 invs: Vec<InvocationRecord>,
849 },
850 Failed {
851 reason: String,
852 records: u64,
853 invs: Vec<InvocationRecord>,
854 },
855 Timeout {
856 secs: u64,
857 },
858 Cancelled,
859 ShutdownFailed,
860}
861
862impl Terminal {
863 fn into_parts(
865 self,
866 ) -> (
867 RunStatus,
868 &'static str,
869 u64,
870 Vec<InvocationRecord>,
871 Option<String>,
872 ) {
873 match self {
874 Terminal::Completed { records, invs } => {
875 (RunStatus::Completed, "ok", records, invs, None)
876 }
877 Terminal::Failed {
878 reason,
879 records,
880 invs,
881 } => (RunStatus::Failed, "error", records, invs, Some(reason)),
882 Terminal::Timeout { secs } => (
883 RunStatus::Failed,
884 "timeout",
885 0,
886 Vec::new(),
887 Some(format!("run exceeded timeout_secs ({secs}s)")),
888 ),
889 Terminal::Cancelled => (RunStatus::Cancelled, "cancelled", 0, Vec::new(), None),
890 Terminal::ShutdownFailed => (
891 RunStatus::Failed,
892 "server_shutdown",
893 0,
894 Vec::new(),
895 Some("server shutdown before the run finished".into()),
896 ),
897 }
898 }
899}
900
901fn classify_run(result: crate::error::CliResult<RunSummary>) -> Terminal {
903 match result {
904 Ok(summary) => {
905 let records: u64 = summary
906 .invocations
907 .iter()
908 .map(|i| i.records_written as u64)
909 .sum();
910 let invs: Vec<InvocationRecord> = summary
911 .invocations
912 .iter()
913 .map(InvocationRecord::from)
914 .collect();
915 if summary.had_failures() {
916 Terminal::Failed {
917 reason: format!("{} invocation(s) failed", summary.failure_count()),
918 records,
919 invs,
920 }
921 } else {
922 Terminal::Completed { records, invs }
923 }
924 }
925 Err(e) => Terminal::Failed {
926 reason: e.to_string(),
927 records: 0,
928 invs: Vec::new(),
929 },
930 }
931}
932
933fn resolve_clock(
935 flag: Option<&str>,
936 default: DateTime<Utc>,
937) -> Result<DateTime<FixedOffset>, ServeError> {
938 match flag {
939 None => Ok(default.fixed_offset()),
940 Some(s) => DateTime::parse_from_rfc3339(s)
941 .map_err(|_| ServeError::BadConfig(format!("clock '{s}' is not RFC3339"))),
942 }
943}
944
945fn spawn_run(
948 state: ServerState,
949 loaded: LoadedSubmission,
950 req: SubmitRequest,
951 run_id: String,
952 run_token: CancellationToken,
953 submitted_at: DateTime<Utc>,
954) {
955 let server_shutdown = state.shutdown_token();
956 tokio::spawn(async move {
957 let _permit = tokio::select! {
963 biased;
964 _ = run_token.cancelled() => {
965 finalize_queued_cancel(&state, &run_id, submitted_at, Terminal::Cancelled).await;
966 return;
967 }
968 _ = server_shutdown.cancelled() => {
969 finalize_queued_cancel(&state, &run_id, submitted_at, Terminal::ShutdownFailed).await;
970 return;
971 }
972 permit = state.semaphore().acquire_owned() => permit.expect("semaphore not closed"),
973 };
974 execute_run(
975 state,
976 loaded,
977 run_id,
978 run_token,
979 submitted_at,
980 req.timeout_secs,
981 req.clock,
982 true,
983 )
984 .await;
985 });
987}
988
989#[allow(clippy::too_many_arguments)]
995async fn execute_run(
996 state: ServerState,
997 loaded: LoadedSubmission,
998 run_id: String,
999 run_token: CancellationToken,
1000 submitted_at: DateTime<Utc>,
1001 timeout_secs: Option<u64>,
1002 clock_flag: Option<String>,
1003 from_queue: bool,
1004) {
1005 let server_shutdown = state.shutdown_token();
1006 let LoadedSubmission { cfg, nodes } = loaded;
1007
1008 if from_queue {
1013 state.registry().mark_running();
1014 } else {
1015 state.registry().mark_running_unqueued();
1016 }
1017 let _guard = InFlightGuard {
1018 state: state.clone(),
1019 run_id: run_id.clone(),
1020 };
1021 let started = Utc::now();
1022 if let Ok(Some(mut rec)) = state.history().get(&run_id).await {
1023 rec.status = RunStatus::Running;
1024 rec.started_at = Some(started);
1025 let _ = state.history().upsert(&rec).await;
1026 }
1027 metrics::set_run_gauges(&state);
1028
1029 let pipeline_name = cfg.name.clone().unwrap_or_else(|| "serve".to_string());
1031 let auth = match build_auth_catalog(cfg.auth.as_ref()) {
1032 Ok(a) => a,
1033 Err(e) => {
1034 finalize(
1035 &state,
1036 &run_id,
1037 started,
1038 Terminal::Failed {
1039 reason: format!("auth catalog: {e}"),
1040 records: 0,
1041 invs: Vec::new(),
1042 },
1043 )
1044 .await;
1045 return;
1046 }
1047 };
1048 let clock = match resolve_clock(clock_flag.as_deref(), submitted_at) {
1049 Ok(c) => c,
1050 Err(e) => {
1051 finalize(
1052 &state,
1053 &run_id,
1054 started,
1055 Terminal::Failed {
1056 reason: e.api_error().error.message,
1057 records: 0,
1058 invs: Vec::new(),
1059 },
1060 )
1061 .await;
1062 return;
1063 }
1064 };
1065
1066 let coop = CancellationToken::new();
1071 let resilience = match &cfg.resilience {
1075 Some(spec) => match spec.to_policy() {
1076 Ok(p) => Some(p),
1077 Err(e) => {
1078 finalize(
1079 &state,
1080 &run_id,
1081 started,
1082 Terminal::Failed {
1083 reason: format!("resilience: {e}"),
1084 records: 0,
1085 invs: Vec::new(),
1086 },
1087 )
1088 .await;
1089 return;
1090 }
1091 },
1092 None => None,
1093 };
1094 #[cfg(feature = "lineage")]
1098 let lineage = match crate::lineage_glue::build_emitter(cfg.lineage.as_ref()) {
1099 Ok(l) => l,
1100 Err(e) => {
1101 finalize(
1102 &state,
1103 &run_id,
1104 started,
1105 Terminal::Failed {
1106 reason: format!("lineage: {e}"),
1107 records: 0,
1108 invs: Vec::new(),
1109 },
1110 )
1111 .await;
1112 return;
1113 }
1114 };
1115 let opts = ExecuteOptions {
1116 pipeline_name,
1117 execution: cfg.execution.clone(),
1118 dry_run: false,
1119 limit: None,
1120 state_path_override: None,
1121 shard: None,
1122 auth,
1123 clock,
1124 cancel: Some(coop.clone()),
1125 resilience,
1126 #[cfg(feature = "lineage")]
1127 lineage,
1128 #[cfg(feature = "lineage")]
1129 lineage_cfg: cfg.lineage.clone(),
1130 };
1131
1132 let span = tracing::info_span!("faucet.serve.run", serve_run_id = %run_id);
1133 let work = async move {
1134 tracing::info!("pipeline run starting");
1137 classify_run(run_expanded(nodes, opts).await)
1138 }
1139 .instrument(span);
1140 tokio::pin!(work);
1141
1142 let timeout_fut = async {
1145 match timeout_secs {
1146 Some(s) => tokio::time::sleep(Duration::from_secs(s)).await,
1147 None => std::future::pending::<()>().await,
1148 }
1149 };
1150 tokio::pin!(timeout_fut);
1151
1152 enum Trigger {
1153 Done(Terminal),
1154 Cancel,
1155 Shutdown,
1156 Timeout(u64),
1157 }
1158
1159 let trigger = tokio::select! {
1162 biased;
1163 t = &mut work => Trigger::Done(t),
1164 _ = run_token.cancelled() => Trigger::Cancel,
1165 _ = server_shutdown.cancelled() => Trigger::Shutdown,
1166 _ = &mut timeout_fut => Trigger::Timeout(timeout_secs.unwrap_or(0)),
1167 };
1168
1169 let terminal = match trigger {
1170 Trigger::Done(t) => t,
1171 triggered => {
1172 coop.cancel();
1177 let trigger_terminal = match triggered {
1178 Trigger::Cancel => Terminal::Cancelled,
1179 Trigger::Shutdown => Terminal::ShutdownFailed,
1180 Trigger::Timeout(secs) => Terminal::Timeout { secs },
1181 Trigger::Done(_) => unreachable!("matched in the outer arm"),
1182 };
1183 match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
1189 Ok(failed @ Terminal::Failed { .. }) => failed,
1190 Ok(_) | Err(_) => trigger_terminal,
1191 }
1192 }
1193 };
1194
1195 finalize(&state, &run_id, started, terminal).await;
1196 state.log_hub().finish(&run_id);
1199 schedule_log_drop(state.clone(), run_id.clone());
1200 }
1202
1203async fn finalize(state: &ServerState, run_id: &str, started: DateTime<Utc>, term: Terminal) {
1205 let finished = Utc::now();
1206 let elapsed = (finished - started).to_std().ok().map(|d| d.as_secs_f64());
1207 let (status, reason, records, invs, error) = term.into_parts();
1208 let mut rec = match state.history().get(run_id).await {
1215 Ok(Some(rec)) => rec,
1216 Ok(None) => {
1217 tracing::warn!(
1218 run_id,
1219 "finalize: run record not found; writing a fresh terminal record"
1220 );
1221 RunRecord::queued(run_id.to_string(), None, BTreeMap::new(), None, started)
1222 }
1223 Err(e) => {
1224 tracing::warn!(
1225 run_id,
1226 error = %e,
1227 "finalize: failed to read run record; writing a fresh terminal record"
1228 );
1229 RunRecord::queued(run_id.to_string(), None, BTreeMap::new(), None, started)
1230 }
1231 };
1232 rec.status = status;
1233 rec.started_at.get_or_insert(started);
1234 rec.finished_at = Some(finished);
1235 rec.elapsed_secs = elapsed;
1236 rec.records_written = records;
1237 rec.invocations = invs;
1238 rec.error = error;
1239 if state.cluster().enabled() {
1240 match state.history().finalize_owned(&rec).await {
1244 Ok(true) => metrics::record_run_finished(status, reason),
1245 Ok(false) => tracing::warn!(
1246 run_id,
1247 "finalize: run was reclaimed by another instance; discarding result"
1248 ),
1249 Err(e) => {
1250 tracing::error!(run_id, error = %e, "finalize: owner-fenced write failed")
1251 }
1252 }
1253 } else {
1254 if let Err(e) = state.history().upsert(&rec).await {
1255 tracing::error!(
1256 run_id,
1257 error = %e,
1258 "finalize: failed to persist terminal run record"
1259 );
1260 }
1261 metrics::record_run_finished(status, reason);
1262 }
1263}
1264
1265async fn finalize_queued_cancel(
1271 state: &ServerState,
1272 run_id: &str,
1273 submitted_at: DateTime<Utc>,
1274 term: Terminal,
1275) {
1276 state.registry().mark_queued_cancelled(run_id);
1277 finalize(state, run_id, submitted_at, term).await;
1278 state.log_hub().finish(run_id);
1279 schedule_log_drop(state.clone(), run_id.to_string());
1280 metrics::set_run_gauges(state);
1281}
1282
1283fn schedule_log_drop(state: ServerState, run_id: String) {
1286 tokio::spawn(async move {
1287 tokio::time::sleep(crate::serve::logs::LOG_DRAIN).await;
1288 state.log_hub().drop_run(&run_id);
1289 });
1290}
1291
1292#[cfg(test)]
1293mod tests {
1294 use super::*;
1295
1296 #[test]
1297 fn classify_ok_no_failures_is_completed() {
1298 let summary = RunSummary {
1299 invocations: vec![crate::executor::InvocationOutcome {
1300 row_id: "r".into(),
1301 parent_record_key: None,
1302 records_written: 3,
1303 error: None,
1304 }],
1305 };
1306 let (status, reason, records, _, error) = classify_run(Ok(summary)).into_parts();
1307 assert_eq!(status, RunStatus::Completed);
1308 assert_eq!(reason, "ok");
1309 assert_eq!(records, 3);
1310 assert!(error.is_none());
1311 }
1312
1313 #[test]
1314 fn classify_ok_with_failures_is_failed() {
1315 let summary = RunSummary {
1316 invocations: vec![crate::executor::InvocationOutcome {
1317 row_id: "r".into(),
1318 parent_record_key: None,
1319 records_written: 0,
1320 error: Some("boom".into()),
1321 }],
1322 };
1323 let (status, reason, _, _, error) = classify_run(Ok(summary)).into_parts();
1324 assert_eq!(status, RunStatus::Failed);
1325 assert_eq!(reason, "error");
1326 assert!(error.unwrap().contains("invocation(s) failed"));
1327 }
1328
1329 #[test]
1330 fn timeout_maps_to_failed_with_timeout_reason() {
1331 let (status, reason, _, _, error) = Terminal::Timeout { secs: 30 }.into_parts();
1332 assert_eq!(status, RunStatus::Failed);
1333 assert_eq!(reason, "timeout");
1334 assert!(error.unwrap().contains("30s"));
1335 }
1336
1337 #[test]
1338 fn resolve_clock_defaults_and_parses() {
1339 let default = Utc::now();
1340 assert_eq!(
1341 resolve_clock(None, default).unwrap(),
1342 default.fixed_offset()
1343 );
1344 assert!(resolve_clock(Some("2026-01-31T00:00:00Z"), default).is_ok());
1345 assert!(resolve_clock(Some("not-a-time"), default).is_err());
1346 }
1347
1348 #[tokio::test]
1349 async fn conflict_releases_reservation() {
1350 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1351 use crate::serve::history::RunHistory;
1352 use crate::serve::history::memory::MemoryHistory;
1353 use crate::serve::state::ServerState;
1354 use std::sync::Arc;
1355 use tokio_util::sync::CancellationToken;
1356
1357 let cfg = ServeConfig {
1358 listen: "127.0.0.1:0".parse().unwrap(),
1359 auth: AuthMode::None,
1360 max_concurrent_runs: 4,
1361 max_queued_runs: 4,
1362 default_config_path: None,
1363 history: HistoryBackendSpec::Memory,
1364 cors_origins: vec![],
1365 body_limit_bytes: 1_048_576,
1366 shutdown_grace: Duration::from_secs(60),
1367 retain_terminal_runs: Duration::from_secs(60),
1368 idempotency_retention: Duration::from_secs(60),
1369 lease_ttl: Duration::from_secs(30),
1370 probe_timeout: Duration::from_secs(10),
1371 env_file: None,
1372 no_env_file: false,
1373 log_level: "info".into(),
1374 ui_enabled: true,
1375 cluster: crate::serve::cluster::ClusterConfig::disabled(),
1376 triggers_path: None,
1377 };
1378 let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1379 let state = ServerState::new(
1380 &cfg,
1381 None,
1382 CancellationToken::new(),
1383 history,
1384 crate::serve::logs::LogHub::new(),
1385 None,
1386 #[cfg(feature = "triggers")]
1387 crate::serve::triggers::health::TriggersHandle::empty(),
1388 );
1389
1390 state
1392 .history()
1393 .claim_idempotency("k", "different-fp", "prior", Duration::from_secs(60))
1394 .await
1395 .unwrap();
1396
1397 let req = SubmitRequest {
1398 config: "version: 1\npipeline:\n source: { type: csv, config: { path: x.csv } }\n sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1399 config_format: ConfigFormatWire::Yaml,
1400 name: None,
1401 labels: BTreeMap::new(),
1402 timeout_secs: None,
1403 doctor_first: false,
1404 idempotency_key: Some("k".into()),
1405 clock: None,
1406 };
1407
1408 let err = submit(state.clone(), req).await.unwrap_err();
1409 assert!(
1410 matches!(err, ServeError::Conflict(_)),
1411 "expected Conflict, got {err:?}"
1412 );
1413 assert_eq!(state.registry().queued(), 0);
1415 }
1416
1417 #[tokio::test]
1418 async fn cluster_submit_writes_pending_with_config_and_does_not_spawn() {
1419 use crate::serve::cluster::ClusterConfig;
1420 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1421 use crate::serve::history::RunHistory;
1422 use crate::serve::history::memory::MemoryHistory;
1423 use crate::serve::state::ServerState;
1424 use std::sync::Arc;
1425 use tokio_util::sync::CancellationToken;
1426
1427 let mut cluster = ClusterConfig::disabled();
1428 cluster.enabled = true;
1429 let cfg = ServeConfig {
1430 listen: "127.0.0.1:0".parse().unwrap(),
1431 auth: AuthMode::None,
1432 max_concurrent_runs: 4,
1433 max_queued_runs: 4,
1434 default_config_path: None,
1435 history: HistoryBackendSpec::Memory,
1436 cors_origins: vec![],
1437 body_limit_bytes: 1_048_576,
1438 shutdown_grace: Duration::from_secs(60),
1439 retain_terminal_runs: Duration::from_secs(60),
1440 idempotency_retention: Duration::from_secs(60),
1441 lease_ttl: Duration::from_secs(30),
1442 probe_timeout: Duration::from_secs(10),
1443 env_file: None,
1444 no_env_file: false,
1445 log_level: "info".into(),
1446 ui_enabled: true,
1447 cluster,
1448 triggers_path: None,
1449 };
1450 let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1451 let state = ServerState::new(
1452 &cfg,
1453 None,
1454 CancellationToken::new(),
1455 history,
1456 crate::serve::logs::LogHub::new(),
1457 None,
1458 #[cfg(feature = "triggers")]
1459 crate::serve::triggers::health::TriggersHandle::empty(),
1460 );
1461
1462 let req = SubmitRequest {
1463 config: "version: 1\npipeline:\n source: { type: csv, config: { path: x.csv } }\n sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1464 config_format: ConfigFormatWire::Yaml,
1465 name: Some("n".into()),
1466 labels: BTreeMap::new(),
1467 timeout_secs: Some(99),
1468 doctor_first: false,
1469 idempotency_key: None,
1470 clock: None,
1471 };
1472 let resp = submit(state.clone(), req).await.unwrap();
1473 assert_eq!(resp.status, RunStatus::Pending);
1474 assert_eq!(state.registry().queued(), 0);
1476 let rec = state.history().get(&resp.run_id).await.unwrap().unwrap();
1477 assert_eq!(rec.status, RunStatus::Pending);
1478 assert!(rec.config_body.as_deref().unwrap().contains("version: 1"));
1479 assert_eq!(rec.timeout_secs, Some(99));
1480 }
1481
1482 fn memory_state() -> crate::serve::state::ServerState {
1484 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1485 use crate::serve::history::RunHistory;
1486 use crate::serve::history::memory::MemoryHistory;
1487 use crate::serve::state::ServerState;
1488 use std::sync::Arc;
1489 use tokio_util::sync::CancellationToken;
1490
1491 let cfg = ServeConfig {
1492 listen: "127.0.0.1:0".parse().unwrap(),
1493 auth: AuthMode::None,
1494 max_concurrent_runs: 4,
1495 max_queued_runs: 4,
1496 default_config_path: None,
1497 history: HistoryBackendSpec::Memory,
1498 cors_origins: vec![],
1499 body_limit_bytes: 1_048_576,
1500 shutdown_grace: Duration::from_secs(60),
1501 retain_terminal_runs: Duration::from_secs(60),
1502 idempotency_retention: Duration::from_secs(60),
1503 lease_ttl: Duration::from_secs(30),
1504 probe_timeout: Duration::from_secs(10),
1505 env_file: None,
1506 no_env_file: false,
1507 log_level: "info".into(),
1508 ui_enabled: true,
1509 cluster: crate::serve::cluster::ClusterConfig::disabled(),
1510 triggers_path: None,
1511 };
1512 let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1513 ServerState::new(
1514 &cfg,
1515 None,
1516 CancellationToken::new(),
1517 history,
1518 crate::serve::logs::LogHub::new(),
1519 None,
1520 #[cfg(feature = "triggers")]
1521 crate::serve::triggers::health::TriggersHandle::empty(),
1522 )
1523 }
1524
1525 #[tokio::test]
1526 async fn finalize_writes_terminal_record_when_record_is_missing() {
1527 let state = memory_state();
1533 let started = Utc::now();
1534 finalize(
1535 &state,
1536 "ghost",
1537 started,
1538 Terminal::Failed {
1539 reason: "boom".into(),
1540 records: 0,
1541 invs: Vec::new(),
1542 },
1543 )
1544 .await;
1545 let rec = state
1546 .history()
1547 .get("ghost")
1548 .await
1549 .unwrap()
1550 .expect("finalize must create a terminal record even when none existed");
1551 assert_eq!(rec.status, RunStatus::Failed);
1552 assert!(rec.finished_at.is_some());
1553 assert!(rec.started_at.is_some());
1554 assert_eq!(rec.error.as_deref(), Some("boom"));
1555 }
1556
1557 #[tokio::test]
1558 async fn finalize_preserves_metadata_of_existing_record() {
1559 let state = memory_state();
1561 let started = Utc::now();
1562 let mut rec = RunRecord::queued(
1563 "r1".into(),
1564 Some("nightly".into()),
1565 BTreeMap::new(),
1566 Some("idem-k".into()),
1567 started,
1568 );
1569 rec.status = RunStatus::Running;
1570 rec.started_at = Some(started);
1571 state.history().upsert(&rec).await.unwrap();
1572
1573 finalize(
1574 &state,
1575 "r1",
1576 started,
1577 Terminal::Completed {
1578 records: 5,
1579 invs: Vec::new(),
1580 },
1581 )
1582 .await;
1583 let got = state.history().get("r1").await.unwrap().unwrap();
1584 assert_eq!(got.status, RunStatus::Completed);
1585 assert_eq!(got.records_written, 5);
1586 assert_eq!(got.name.as_deref(), Some("nightly"));
1587 assert_eq!(got.idempotency_key.as_deref(), Some("idem-k"));
1588 }
1589
1590 #[cfg(any(feature = "serve-history-sqlite", feature = "serve-history-postgres"))]
1591 #[tokio::test]
1592 async fn cluster_submit_503s_when_history_degraded() {
1593 use crate::serve::cluster::ClusterConfig;
1596 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1597 use crate::serve::history::RunHistory;
1598 use crate::serve::history::fallback::FallbackHistory;
1599 use crate::serve::state::ServerState;
1600 use std::sync::Arc;
1601 use tokio_util::sync::CancellationToken;
1602
1603 let mut cluster = ClusterConfig::disabled();
1604 cluster.enabled = true;
1605 let cfg = ServeConfig {
1606 listen: "127.0.0.1:0".parse().unwrap(),
1607 auth: AuthMode::None,
1608 max_concurrent_runs: 4,
1609 max_queued_runs: 4,
1610 default_config_path: None,
1611 history: HistoryBackendSpec::Memory,
1612 cors_origins: vec![],
1613 body_limit_bytes: 1_048_576,
1614 shutdown_grace: Duration::from_secs(60),
1615 retain_terminal_runs: Duration::from_secs(60),
1616 idempotency_retention: Duration::from_secs(60),
1617 lease_ttl: Duration::from_secs(30),
1618 probe_timeout: Duration::from_secs(10),
1619 env_file: None,
1620 no_env_file: false,
1621 log_level: "info".into(),
1622 ui_enabled: true,
1623 cluster,
1624 triggers_path: None,
1625 };
1626 let history = Arc::new(FallbackHistory::degraded_at_startup(
1628 Duration::from_secs(60),
1629 "test",
1630 )) as Arc<dyn RunHistory>;
1631 assert!(history.degraded());
1632 let state = ServerState::new(
1633 &cfg,
1634 None,
1635 CancellationToken::new(),
1636 history,
1637 crate::serve::logs::LogHub::new(),
1638 None,
1639 #[cfg(feature = "triggers")]
1640 crate::serve::triggers::health::TriggersHandle::empty(),
1641 );
1642 let req = SubmitRequest {
1643 config: "version: 1\npipeline:\n source: { type: csv, config: { path: x.csv } }\n sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1644 config_format: ConfigFormatWire::Yaml,
1645 name: None,
1646 labels: BTreeMap::new(),
1647 timeout_secs: None,
1648 doctor_first: false,
1649 idempotency_key: None,
1650 clock: None,
1651 };
1652 let err = submit(state.clone(), req).await.unwrap_err();
1653 assert!(
1654 matches!(err, ServeError::Unavailable(_)),
1655 "expected 503 Unavailable, got {err:?}"
1656 );
1657 assert_eq!(state.registry().queued(), 0);
1659 }
1660
1661 #[cfg(feature = "serve-history-sqlite")]
1667 mod shards {
1668 use super::*;
1669 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1670 use crate::serve::history::RunHistory;
1671 use crate::serve::history::sqlite::SqliteHistory;
1672 use crate::serve::load::{ConfigFormat, load_submission};
1673 use crate::serve::state::ServerState;
1674 use faucet_core::ShardSpec;
1675 use std::collections::BTreeMap;
1676 use std::sync::Arc;
1677 use tokio_util::sync::CancellationToken;
1678
1679 async fn sqlite_state(dir: &std::path::Path) -> ServerState {
1680 let url = format!("sqlite://{}/h.db", dir.display());
1681 let history = Arc::new(
1682 SqliteHistory::connect(
1683 &url,
1684 Duration::from_secs(300),
1685 Duration::from_secs(300),
1686 "inst-test".into(),
1687 )
1688 .await
1689 .expect("sqlite history"),
1690 ) as Arc<dyn RunHistory>;
1691 let cfg = ServeConfig {
1692 listen: "127.0.0.1:0".parse().unwrap(),
1693 auth: AuthMode::None,
1694 max_concurrent_runs: 4,
1695 max_queued_runs: 4,
1696 default_config_path: None,
1697 history: HistoryBackendSpec::Memory,
1698 cors_origins: vec![],
1699 body_limit_bytes: 1_048_576,
1700 shutdown_grace: Duration::from_secs(60),
1701 retain_terminal_runs: Duration::from_secs(60),
1702 idempotency_retention: Duration::from_secs(60),
1703 lease_ttl: Duration::from_secs(30),
1704 probe_timeout: Duration::from_secs(10),
1705 env_file: None,
1706 no_env_file: false,
1707 log_level: "info".into(),
1708 ui_enabled: true,
1709 cluster: crate::serve::cluster::ClusterConfig::disabled(),
1710 triggers_path: None,
1711 };
1712 ServerState::new(
1713 &cfg,
1714 None,
1715 CancellationToken::new(),
1716 history,
1717 crate::serve::logs::LogHub::new(),
1718 None,
1719 #[cfg(feature = "triggers")]
1720 crate::serve::triggers::health::TriggersHandle::empty(),
1721 )
1722 }
1723
1724 async fn loaded(yaml: &str) -> LoadedSubmission {
1725 load_submission(yaml, ConfigFormat::Yaml, None)
1726 .await
1727 .expect("load submission")
1728 }
1729
1730 async fn seed_run(state: &ServerState, run_id: &str, status: RunStatus) {
1731 let mut rec = RunRecord::queued(run_id.into(), None, BTreeMap::new(), None, Utc::now());
1732 rec.status = status;
1733 rec.config_body = Some("version: 1".into());
1734 state.history().upsert(&rec).await.expect("seed run");
1735 }
1736
1737 #[tokio::test]
1738 async fn coordinate_matrix_run_is_not_shardable() {
1739 let dir = tempfile::tempdir().unwrap();
1741 let state = sqlite_state(dir.path()).await;
1742 let l = loaded(
1743 "version: 1\nname: m\nmatrix:\n - id: a\n - id: b\npipeline:\n \
1744 source: { type: rest, config: { url: \"http://localhost/x\" } }\n \
1745 sink: { type: stdout, config: {} }\n",
1746 )
1747 .await;
1748 assert!(!coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1749 }
1750
1751 #[tokio::test]
1752 async fn coordinate_non_shardable_source_runs_whole() {
1753 let dir = tempfile::tempdir().unwrap();
1755 let state = sqlite_state(dir.path()).await;
1756 let input = dir.path().join("in.csv");
1757 std::fs::write(&input, "id\n1\n").unwrap();
1758 let l = loaded(&format!(
1759 "version: 1\npipeline:\n \
1760 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n \
1761 sink: {{ type: stdout, config: {{}} }}\n",
1762 input.display()
1763 ))
1764 .await;
1765 assert!(!coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1766 }
1767
1768 #[tokio::test]
1769 async fn coordinate_s3_source_inserts_shards_and_marks_sharded() {
1770 let dir = tempfile::tempdir().unwrap();
1771 let state = sqlite_state(dir.path()).await;
1772 seed_run(&state, "r", RunStatus::Running).await;
1773 let l = loaded(
1774 "version: 1\npipeline:\n \
1775 source: { type: s3, config: { bucket: my-bucket, prefix: null, \
1776 region: null, endpoint_url: null, file_format: json_lines, \
1777 max_objects: null, concurrency: 10 } }\n \
1778 sink: { type: stdout, config: {} }\n",
1779 )
1780 .await;
1781 assert!(coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1782 let prog = state.history().shard_progress("r").await.unwrap();
1784 assert_eq!(prog.total, 4);
1785 assert_eq!(prog.pending, 4);
1786 assert_eq!(
1787 state.history().get("r").await.unwrap().unwrap().status,
1788 RunStatus::Sharded
1789 );
1790 }
1791
1792 #[tokio::test]
1793 async fn at_least_once_risky_sinks_flags_append_cluster_and_shard() {
1794 let append = loaded(
1797 "version: 1\npipeline:\n \
1798 source: { type: rest, config: { url: \"http://localhost/x\" } }\n \
1799 sink: { type: stdout, config: {} }\n",
1800 )
1801 .await;
1802 assert!(at_least_once_risky_sinks(&append, false, false).is_empty());
1804 assert_eq!(
1806 at_least_once_risky_sinks(&append, true, false),
1807 vec!["stdout"]
1808 );
1809 assert_eq!(
1811 at_least_once_risky_sinks(&append, false, true),
1812 vec!["stdout"]
1813 );
1814 }
1815
1816 #[tokio::test]
1817 async fn at_least_once_risky_sinks_safe_for_upsert_and_exactly_once() {
1818 let upsert = loaded(
1820 "version: 1\npipeline:\n \
1821 source: { type: postgres, config: { connection_url: \"postgres://x\", \
1822 query: \"select 1\" } }\n \
1823 sink: { type: postgres, config: { connection_url: \"postgres://y\", \
1824 table_name: t, column_mapping: auto_map, write_mode: upsert, key: [id] } }\n",
1825 )
1826 .await;
1827 assert!(at_least_once_risky_sinks(&upsert, true, true).is_empty());
1828
1829 let eo = loaded(
1831 "version: 1\ndelivery: exactly_once\npipeline:\n \
1832 source: { type: postgres-cdc, config: {} }\n \
1833 sink: { type: sqlite, config: {} }\n \
1834 state: { type: file, config: { path: \"/tmp/x.json\" } }\n",
1835 )
1836 .await;
1837 assert!(at_least_once_risky_sinks(&eo, true, false).is_empty());
1838 }
1839
1840 async fn seed_sharded_with_shards(state: &ServerState, run_id: &str, n: usize) {
1841 use crate::serve::history::ShardInsert;
1842 seed_run(state, run_id, RunStatus::Sharded).await;
1843 let shards: Vec<ShardInsert> = (0..n)
1844 .map(|i| ShardInsert {
1845 shard_id: i.to_string(),
1846 descriptor: serde_json::json!({ "i": i }),
1847 size_estimate: None,
1848 })
1849 .collect();
1850 state
1851 .history()
1852 .insert_shards(run_id, &shards)
1853 .await
1854 .unwrap();
1855 let claimed = state.history().claim_shards(n).await.unwrap();
1857 assert_eq!(claimed.len(), n);
1858 }
1859
1860 #[tokio::test]
1861 async fn maybe_finalize_parent_completes_when_all_shards_succeed() {
1862 let dir = tempfile::tempdir().unwrap();
1863 let state = sqlite_state(dir.path()).await;
1864 seed_sharded_with_shards(&state, "r", 3).await;
1865 for i in 0..3 {
1866 state
1867 .history()
1868 .finalize_shard("r", &i.to_string(), true)
1869 .await
1870 .unwrap();
1871 }
1872 maybe_finalize_parent(&state, "r").await;
1873 assert_eq!(
1874 state.history().get("r").await.unwrap().unwrap().status,
1875 RunStatus::Completed
1876 );
1877 }
1878
1879 #[tokio::test]
1880 async fn maybe_finalize_parent_fails_when_a_shard_fails() {
1881 let dir = tempfile::tempdir().unwrap();
1882 let state = sqlite_state(dir.path()).await;
1883 seed_sharded_with_shards(&state, "r", 2).await;
1884 state
1885 .history()
1886 .finalize_shard("r", "0", true)
1887 .await
1888 .unwrap();
1889 state
1890 .history()
1891 .finalize_shard("r", "1", false)
1892 .await
1893 .unwrap();
1894 maybe_finalize_parent(&state, "r").await;
1895 assert_eq!(
1896 state.history().get("r").await.unwrap().unwrap().status,
1897 RunStatus::Failed
1898 );
1899 }
1900
1901 #[tokio::test]
1902 async fn maybe_finalize_parent_keeps_sharded_until_all_terminal() {
1903 let dir = tempfile::tempdir().unwrap();
1904 let state = sqlite_state(dir.path()).await;
1905 seed_sharded_with_shards(&state, "r", 2).await;
1906 state
1908 .history()
1909 .finalize_shard("r", "0", true)
1910 .await
1911 .unwrap();
1912 maybe_finalize_parent(&state, "r").await;
1913 assert_eq!(
1914 state.history().get("r").await.unwrap().unwrap().status,
1915 RunStatus::Sharded
1916 );
1917 }
1918
1919 #[tokio::test]
1920 async fn finalize_sharded_parent_is_status_fenced_and_idempotent() {
1921 let dir = tempfile::tempdir().unwrap();
1925 let state = sqlite_state(dir.path()).await;
1926 seed_sharded_with_shards(&state, "r", 2).await;
1927
1928 let first = state
1929 .history()
1930 .finalize_sharded_parent("r", RunStatus::Completed, Utc::now(), None)
1931 .await
1932 .unwrap();
1933 assert!(first, "first finalize wins");
1934 assert_eq!(
1935 state.history().get("r").await.unwrap().unwrap().status,
1936 RunStatus::Completed
1937 );
1938
1939 let second = state
1941 .history()
1942 .finalize_sharded_parent("r", RunStatus::Failed, Utc::now(), Some("late".into()))
1943 .await
1944 .unwrap();
1945 assert!(!second, "second finalize is a no-op");
1946 let r = state.history().get("r").await.unwrap().unwrap();
1947 assert_eq!(r.status, RunStatus::Completed, "status not overwritten");
1948 assert!(r.error.is_none(), "late error must not be stamped");
1949
1950 let missing = state
1952 .history()
1953 .finalize_sharded_parent("does-not-exist", RunStatus::Completed, Utc::now(), None)
1954 .await
1955 .unwrap();
1956 assert!(!missing);
1957 }
1958
1959 #[tokio::test]
1960 async fn execute_shard_runs_a_csv_to_jsonl_shard() {
1961 let dir = tempfile::tempdir().unwrap();
1962 let state = sqlite_state(dir.path()).await;
1963 let input = dir.path().join("in.csv");
1964 std::fs::write(&input, "id,name\n1,alice\n2,bob\n").unwrap();
1965 let output = dir.path().join("out.jsonl");
1966 let yaml = format!(
1967 "version: 1\npipeline:\n \
1968 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n \
1969 sink: {{ type: jsonl, config: {{ path: \"{}\" }} }}\n",
1970 input.display(),
1971 output.display()
1972 );
1973 let l = loaded(&yaml).await;
1974 let ok = execute_shard(
1977 &state,
1978 l,
1979 "r",
1980 "0",
1981 ShardSpec::whole(),
1982 CancellationToken::new(),
1983 None,
1984 None,
1985 Utc::now(),
1986 )
1987 .await;
1988 assert!(ok, "csv→jsonl shard should complete");
1989 let written = std::fs::read_to_string(&output).unwrap();
1990 assert_eq!(written.lines().count(), 2, "both rows written");
1991 assert!(written.contains("alice") && written.contains("bob"));
1992 }
1993
1994 #[tokio::test]
1995 async fn resume_claimed_shard_executes_and_finalizes_parent() {
1996 let dir = tempfile::tempdir().unwrap();
2000 let state = sqlite_state(dir.path()).await;
2001 let input = dir.path().join("in.csv");
2002 std::fs::write(&input, "id,name\n1,alice\n").unwrap();
2003 let output = dir.path().join("out.jsonl");
2004 let yaml = format!(
2005 "version: 1\npipeline:\n \
2006 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n \
2007 sink: {{ type: jsonl, config: {{ path: \"{}\" }} }}\n",
2008 input.display(),
2009 output.display()
2010 );
2011 let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2013 rec.status = RunStatus::Sharded;
2014 rec.config_body = Some(yaml);
2015 state.history().upsert(&rec).await.unwrap();
2016 use crate::serve::history::ShardInsert;
2017 state
2018 .history()
2019 .insert_shards(
2020 "r",
2021 &[ShardInsert {
2022 shard_id: "0".into(),
2023 descriptor: serde_json::Value::Null,
2024 size_estimate: None,
2025 }],
2026 )
2027 .await
2028 .unwrap();
2029 let claimed = state.history().claim_shards(1).await.unwrap();
2030 assert_eq!(claimed.len(), 1);
2031
2032 resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2033
2034 let mut status = RunStatus::Sharded;
2037 for _ in 0..100 {
2038 tokio::time::sleep(Duration::from_millis(50)).await;
2039 status = state.history().get("r").await.unwrap().unwrap().status;
2040 if status.is_terminal() {
2041 break;
2042 }
2043 }
2044 assert_eq!(status, RunStatus::Completed, "shard ran → parent completed");
2045 assert!(output.exists(), "shard wrote its output");
2046 }
2047
2048 #[tokio::test]
2049 async fn resume_claimed_shard_with_no_config_fails_the_shard() {
2050 let dir = tempfile::tempdir().unwrap();
2053 let state = sqlite_state(dir.path()).await;
2054 let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2055 rec.status = RunStatus::Sharded; state.history().upsert(&rec).await.unwrap();
2057 use crate::serve::history::ShardInsert;
2058 state
2059 .history()
2060 .insert_shards(
2061 "r",
2062 &[ShardInsert {
2063 shard_id: "0".into(),
2064 descriptor: serde_json::Value::Null,
2065 size_estimate: None,
2066 }],
2067 )
2068 .await
2069 .unwrap();
2070 let claimed = state.history().claim_shards(1).await.unwrap();
2071 resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2072
2073 let mut status = RunStatus::Sharded;
2074 for _ in 0..100 {
2075 tokio::time::sleep(Duration::from_millis(50)).await;
2076 status = state.history().get("r").await.unwrap().unwrap().status;
2077 if status.is_terminal() {
2078 break;
2079 }
2080 }
2081 assert_eq!(status, RunStatus::Failed, "no-config shard → parent failed");
2082 }
2083
2084 #[tokio::test]
2085 async fn coordinate_returns_err_when_source_build_fails() {
2086 let dir = tempfile::tempdir().unwrap();
2089 let state = sqlite_state(dir.path()).await;
2090 let l = loaded(
2092 "version: 1\npipeline:\n \
2093 source: { type: s3, config: { bucket: b } }\n \
2094 sink: { type: stdout, config: {} }\n",
2095 )
2096 .await;
2097 assert!(coordinate_sharded_run(&state, "r", &l, 4).await.is_err());
2098 }
2099
2100 #[tokio::test]
2101 async fn execute_shard_returns_false_on_malformed_resilience() {
2102 let dir = tempfile::tempdir().unwrap();
2105 let state = sqlite_state(dir.path()).await;
2106 let input = dir.path().join("in.csv");
2107 std::fs::write(&input, "id\n1\n").unwrap();
2108 let yaml = format!(
2109 "version: 1\nresilience:\n retry:\n max_attempts: 0\npipeline:\n \
2110 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n \
2111 sink: {{ type: stdout, config: {{}} }}\n",
2112 input.display()
2113 );
2114 let l = loaded(&yaml).await;
2115 let ok = execute_shard(
2116 &state,
2117 l,
2118 "r",
2119 "0",
2120 ShardSpec::whole(),
2121 CancellationToken::new(),
2122 None,
2123 None,
2124 Utc::now(),
2125 )
2126 .await;
2127 assert!(!ok, "malformed resilience → shard fails fast");
2128 }
2129
2130 #[tokio::test]
2131 async fn resume_claimed_shard_with_unloadable_config_fails_the_shard() {
2132 let dir = tempfile::tempdir().unwrap();
2135 let state = sqlite_state(dir.path()).await;
2136 let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2137 rec.status = RunStatus::Sharded;
2138 rec.config_body = Some("this: is: not: valid: yaml: [".into());
2139 state.history().upsert(&rec).await.unwrap();
2140 use crate::serve::history::ShardInsert;
2141 state
2142 .history()
2143 .insert_shards(
2144 "r",
2145 &[ShardInsert {
2146 shard_id: "0".into(),
2147 descriptor: serde_json::Value::Null,
2148 size_estimate: None,
2149 }],
2150 )
2151 .await
2152 .unwrap();
2153 let claimed = state.history().claim_shards(1).await.unwrap();
2154 resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2155
2156 let mut status = RunStatus::Sharded;
2157 for _ in 0..100 {
2158 tokio::time::sleep(Duration::from_millis(50)).await;
2159 status = state.history().get("r").await.unwrap().unwrap().status;
2160 if status.is_terminal() {
2161 break;
2162 }
2163 }
2164 assert_eq!(
2165 status,
2166 RunStatus::Failed,
2167 "unloadable config → parent failed"
2168 );
2169 }
2170
2171 #[tokio::test]
2174 async fn request_cancel_flags_a_sharded_parent() {
2175 let dir = tempfile::tempdir().unwrap();
2178 let state = sqlite_state(dir.path()).await;
2179 seed_run(&state, "r", RunStatus::Sharded).await;
2180 state.history().request_cancel("r").await.unwrap();
2183 use crate::serve::history::ShardInsert;
2185 state
2186 .history()
2187 .insert_shards(
2188 "r",
2189 &[ShardInsert {
2190 shard_id: "0".into(),
2191 descriptor: serde_json::Value::Null,
2192 size_estimate: None,
2193 }],
2194 )
2195 .await
2196 .unwrap();
2197 let claimed = state.history().claim_shards(1).await.unwrap();
2198 assert_eq!(claimed.len(), 1, "shard claimed (running, owned)");
2199
2200 let flagged = state.history().pending_shard_cancellations().await.unwrap();
2201 assert_eq!(
2202 flagged,
2203 vec!["r".to_string()],
2204 "the flagged sharded parent's run id is returned for its running shard"
2205 );
2206 }
2207
2208 #[tokio::test]
2209 async fn pending_shard_cancellations_filters_unflagged_and_pending_shards() {
2210 let dir = tempfile::tempdir().unwrap();
2211 let state = sqlite_state(dir.path()).await;
2212 use crate::serve::history::ShardInsert;
2213 let one = |id: &str| {
2214 vec![ShardInsert {
2215 shard_id: id.into(),
2216 descriptor: serde_json::Value::Null,
2217 size_estimate: None,
2218 }]
2219 };
2220
2221 seed_run(&state, "A", RunStatus::Sharded).await;
2224 state.history().request_cancel("A").await.unwrap();
2225 state.history().insert_shards("A", &one("0")).await.unwrap();
2226 seed_run(&state, "C", RunStatus::Sharded).await;
2227 state.history().insert_shards("C", &one("0")).await.unwrap();
2228 let claimed = state.history().claim_shards(8).await.unwrap();
2229 assert_eq!(claimed.len(), 2, "A and C shards claimed (running)");
2230
2231 seed_run(&state, "B", RunStatus::Sharded).await;
2235 state.history().request_cancel("B").await.unwrap();
2236 state.history().insert_shards("B", &one("0")).await.unwrap();
2237
2238 let flagged = state.history().pending_shard_cancellations().await.unwrap();
2239 assert_eq!(
2240 flagged,
2241 vec!["A".to_string()],
2242 "only A (flagged + a running owned shard); B pending-shard, C unflagged"
2243 );
2244 }
2245
2246 #[tokio::test]
2249 async fn finalize_sweep_completes_an_all_success_sharded_parent() {
2250 let dir = tempfile::tempdir().unwrap();
2251 let state = sqlite_state(dir.path()).await;
2252 seed_sharded_with_shards(&state, "r", 3).await;
2253 for i in 0..3 {
2254 state
2255 .history()
2256 .finalize_shard("r", &i.to_string(), true)
2257 .await
2258 .unwrap();
2259 }
2260 let n = state
2262 .history()
2263 .finalize_completed_sharded_parents()
2264 .await
2265 .unwrap();
2266 assert_eq!(n, 1, "one sharded parent finalized");
2267 let rec = state.history().get("r").await.unwrap().unwrap();
2268 assert_eq!(rec.status, RunStatus::Completed);
2269 assert!(rec.finished_at.is_some());
2270 assert!(rec.error.is_none());
2271
2272 assert_eq!(
2274 state
2275 .history()
2276 .finalize_completed_sharded_parents()
2277 .await
2278 .unwrap(),
2279 0,
2280 "already-terminal parent is not re-finalized"
2281 );
2282 }
2283
2284 #[tokio::test]
2285 async fn finalize_sweep_fails_a_parent_with_a_failed_shard() {
2286 let dir = tempfile::tempdir().unwrap();
2287 let state = sqlite_state(dir.path()).await;
2288 seed_sharded_with_shards(&state, "r", 3).await;
2289 state
2290 .history()
2291 .finalize_shard("r", "0", true)
2292 .await
2293 .unwrap();
2294 state
2295 .history()
2296 .finalize_shard("r", "1", false)
2297 .await
2298 .unwrap();
2299 state
2300 .history()
2301 .finalize_shard("r", "2", true)
2302 .await
2303 .unwrap();
2304 let n = state
2305 .history()
2306 .finalize_completed_sharded_parents()
2307 .await
2308 .unwrap();
2309 assert_eq!(n, 1);
2310 let rec = state.history().get("r").await.unwrap().unwrap();
2311 assert_eq!(rec.status, RunStatus::Failed);
2312 assert!(rec.finished_at.is_some());
2313 assert_eq!(rec.error.as_deref(), Some("1/3 shard(s) failed"));
2314 }
2315
2316 #[tokio::test]
2317 async fn finalize_sweep_leaves_a_not_all_terminal_parent_sharded() {
2318 let dir = tempfile::tempdir().unwrap();
2319 let state = sqlite_state(dir.path()).await;
2320 seed_sharded_with_shards(&state, "r", 2).await;
2321 state
2323 .history()
2324 .finalize_shard("r", "0", true)
2325 .await
2326 .unwrap();
2327 let n = state
2328 .history()
2329 .finalize_completed_sharded_parents()
2330 .await
2331 .unwrap();
2332 assert_eq!(n, 0, "parent with a still-running shard is not finalized");
2333 assert_eq!(
2334 state.history().get("r").await.unwrap().unwrap().status,
2335 RunStatus::Sharded
2336 );
2337 }
2338 }
2339}