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