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 metrics: None,
1412 }],
1413 };
1414 let (status, reason, records, _, error) = classify_run(Ok(summary)).into_parts();
1415 assert_eq!(status, RunStatus::Completed);
1416 assert_eq!(reason, "ok");
1417 assert_eq!(records, 3);
1418 assert!(error.is_none());
1419 }
1420
1421 #[test]
1422 fn classify_ok_with_failures_is_failed() {
1423 let summary = RunSummary {
1424 invocations: vec![crate::executor::InvocationOutcome {
1425 row_id: "r".into(),
1426 parent_record_key: None,
1427 records_written: 0,
1428 error: Some("boom".into()),
1429 metrics: None,
1430 }],
1431 };
1432 let (status, reason, _, _, error) = classify_run(Ok(summary)).into_parts();
1433 assert_eq!(status, RunStatus::Failed);
1434 assert_eq!(reason, "error");
1435 assert!(error.unwrap().contains("invocation(s) failed"));
1436 }
1437
1438 #[test]
1439 fn timeout_maps_to_failed_with_timeout_reason() {
1440 let (status, reason, _, _, error) = Terminal::Timeout { secs: 30 }.into_parts();
1441 assert_eq!(status, RunStatus::Failed);
1442 assert_eq!(reason, "timeout");
1443 assert!(error.unwrap().contains("30s"));
1444 }
1445
1446 #[test]
1447 fn resolve_clock_defaults_and_parses() {
1448 let default = Utc::now();
1449 assert_eq!(
1450 resolve_clock(None, default).unwrap(),
1451 default.fixed_offset()
1452 );
1453 assert!(resolve_clock(Some("2026-01-31T00:00:00Z"), default).is_ok());
1454 assert!(resolve_clock(Some("not-a-time"), default).is_err());
1455 }
1456
1457 #[tokio::test]
1458 async fn conflict_releases_reservation() {
1459 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1460 use crate::serve::history::RunHistory;
1461 use crate::serve::history::memory::MemoryHistory;
1462 use crate::serve::state::ServerState;
1463 use std::sync::Arc;
1464 use tokio_util::sync::CancellationToken;
1465
1466 let cfg = ServeConfig {
1467 listen: "127.0.0.1:0".parse().unwrap(),
1468 auth: AuthMode::None,
1469 max_concurrent_runs: 4,
1470 max_queued_runs: 4,
1471 default_config_path: None,
1472 history: HistoryBackendSpec::Memory,
1473 cors_origins: vec![],
1474 body_limit_bytes: 1_048_576,
1475 shutdown_grace: Duration::from_secs(60),
1476 retain_terminal_runs: Duration::from_secs(60),
1477 idempotency_retention: Duration::from_secs(60),
1478 lease_ttl: Duration::from_secs(30),
1479 probe_timeout: Duration::from_secs(10),
1480 env_file: None,
1481 no_env_file: false,
1482 log_level: "info".into(),
1483 ui_enabled: true,
1484 cluster: crate::serve::cluster::ClusterConfig::disabled(),
1485 triggers_path: None,
1486 };
1487 let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1488 let state = ServerState::new(
1489 &cfg,
1490 None,
1491 CancellationToken::new(),
1492 history,
1493 crate::serve::logs::LogHub::new(),
1494 None,
1495 #[cfg(feature = "triggers")]
1496 crate::serve::triggers::health::TriggersHandle::empty(),
1497 );
1498
1499 state
1501 .history()
1502 .claim_idempotency("k", "different-fp", "prior", Duration::from_secs(60))
1503 .await
1504 .unwrap();
1505
1506 let req = SubmitRequest {
1507 config: "version: 1\npipeline:\n source: { type: csv, config: { path: x.csv } }\n sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1508 config_format: ConfigFormatWire::Yaml,
1509 name: None,
1510 labels: BTreeMap::new(),
1511 timeout_secs: None,
1512 doctor_first: false,
1513 idempotency_key: Some("k".into()),
1514 clock: None,
1515 };
1516
1517 let err = submit(state.clone(), req, admin_actor()).await.unwrap_err();
1518 assert!(
1519 matches!(err, ServeError::Conflict(_)),
1520 "expected Conflict, got {err:?}"
1521 );
1522 assert_eq!(state.registry().queued(), 0);
1524 }
1525
1526 #[tokio::test]
1527 async fn cluster_submit_writes_pending_with_config_and_does_not_spawn() {
1528 use crate::serve::cluster::ClusterConfig;
1529 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1530 use crate::serve::history::RunHistory;
1531 use crate::serve::history::memory::MemoryHistory;
1532 use crate::serve::state::ServerState;
1533 use std::sync::Arc;
1534 use tokio_util::sync::CancellationToken;
1535
1536 let mut cluster = ClusterConfig::disabled();
1537 cluster.enabled = true;
1538 let cfg = ServeConfig {
1539 listen: "127.0.0.1:0".parse().unwrap(),
1540 auth: AuthMode::None,
1541 max_concurrent_runs: 4,
1542 max_queued_runs: 4,
1543 default_config_path: None,
1544 history: HistoryBackendSpec::Memory,
1545 cors_origins: vec![],
1546 body_limit_bytes: 1_048_576,
1547 shutdown_grace: Duration::from_secs(60),
1548 retain_terminal_runs: Duration::from_secs(60),
1549 idempotency_retention: Duration::from_secs(60),
1550 lease_ttl: Duration::from_secs(30),
1551 probe_timeout: Duration::from_secs(10),
1552 env_file: None,
1553 no_env_file: false,
1554 log_level: "info".into(),
1555 ui_enabled: true,
1556 cluster,
1557 triggers_path: None,
1558 };
1559 let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1560 let state = ServerState::new(
1561 &cfg,
1562 None,
1563 CancellationToken::new(),
1564 history,
1565 crate::serve::logs::LogHub::new(),
1566 None,
1567 #[cfg(feature = "triggers")]
1568 crate::serve::triggers::health::TriggersHandle::empty(),
1569 );
1570
1571 let req = SubmitRequest {
1572 config: "version: 1\npipeline:\n source: { type: csv, config: { path: x.csv } }\n sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1573 config_format: ConfigFormatWire::Yaml,
1574 name: Some("n".into()),
1575 labels: BTreeMap::new(),
1576 timeout_secs: Some(99),
1577 doctor_first: false,
1578 idempotency_key: None,
1579 clock: None,
1580 };
1581 let resp = submit(state.clone(), req, admin_actor()).await.unwrap();
1582 assert_eq!(resp.status, RunStatus::Pending);
1583 assert_eq!(state.registry().queued(), 0);
1585 let rec = state.history().get(&resp.run_id).await.unwrap().unwrap();
1586 assert_eq!(rec.status, RunStatus::Pending);
1587 assert!(rec.config_body.as_deref().unwrap().contains("version: 1"));
1588 assert_eq!(rec.timeout_secs, Some(99));
1589 }
1590
1591 fn memory_state() -> crate::serve::state::ServerState {
1593 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1594 use crate::serve::history::RunHistory;
1595 use crate::serve::history::memory::MemoryHistory;
1596 use crate::serve::state::ServerState;
1597 use std::sync::Arc;
1598 use tokio_util::sync::CancellationToken;
1599
1600 let cfg = ServeConfig {
1601 listen: "127.0.0.1:0".parse().unwrap(),
1602 auth: AuthMode::None,
1603 max_concurrent_runs: 4,
1604 max_queued_runs: 4,
1605 default_config_path: None,
1606 history: HistoryBackendSpec::Memory,
1607 cors_origins: vec![],
1608 body_limit_bytes: 1_048_576,
1609 shutdown_grace: Duration::from_secs(60),
1610 retain_terminal_runs: Duration::from_secs(60),
1611 idempotency_retention: Duration::from_secs(60),
1612 lease_ttl: Duration::from_secs(30),
1613 probe_timeout: Duration::from_secs(10),
1614 env_file: None,
1615 no_env_file: false,
1616 log_level: "info".into(),
1617 ui_enabled: true,
1618 cluster: crate::serve::cluster::ClusterConfig::disabled(),
1619 triggers_path: None,
1620 };
1621 let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1622 ServerState::new(
1623 &cfg,
1624 None,
1625 CancellationToken::new(),
1626 history,
1627 crate::serve::logs::LogHub::new(),
1628 None,
1629 #[cfg(feature = "triggers")]
1630 crate::serve::triggers::health::TriggersHandle::empty(),
1631 )
1632 }
1633
1634 #[tokio::test]
1635 async fn finalize_writes_terminal_record_when_record_is_missing() {
1636 let state = memory_state();
1642 let started = Utc::now();
1643 finalize(
1644 &state,
1645 "ghost",
1646 started,
1647 Terminal::Failed {
1648 reason: "boom".into(),
1649 records: 0,
1650 invs: Vec::new(),
1651 },
1652 )
1653 .await;
1654 let rec = state
1655 .history()
1656 .get("ghost")
1657 .await
1658 .unwrap()
1659 .expect("finalize must create a terminal record even when none existed");
1660 assert_eq!(rec.status, RunStatus::Failed);
1661 assert!(rec.finished_at.is_some());
1662 assert!(rec.started_at.is_some());
1663 assert_eq!(rec.error.as_deref(), Some("boom"));
1664 }
1665
1666 #[tokio::test]
1667 async fn finalize_preserves_metadata_of_existing_record() {
1668 let state = memory_state();
1670 let started = Utc::now();
1671 let mut rec = RunRecord::queued(
1672 "r1".into(),
1673 Some("nightly".into()),
1674 BTreeMap::new(),
1675 Some("idem-k".into()),
1676 started,
1677 );
1678 rec.status = RunStatus::Running;
1679 rec.started_at = Some(started);
1680 state.history().upsert(&rec).await.unwrap();
1681
1682 finalize(
1683 &state,
1684 "r1",
1685 started,
1686 Terminal::Completed {
1687 records: 5,
1688 invs: Vec::new(),
1689 },
1690 )
1691 .await;
1692 let got = state.history().get("r1").await.unwrap().unwrap();
1693 assert_eq!(got.status, RunStatus::Completed);
1694 assert_eq!(got.records_written, 5);
1695 assert_eq!(got.name.as_deref(), Some("nightly"));
1696 assert_eq!(got.idempotency_key.as_deref(), Some("idem-k"));
1697 }
1698
1699 #[cfg(any(feature = "serve-history-sqlite", feature = "serve-history-postgres"))]
1700 #[tokio::test]
1701 async fn cluster_submit_503s_when_history_degraded() {
1702 use crate::serve::cluster::ClusterConfig;
1705 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1706 use crate::serve::history::RunHistory;
1707 use crate::serve::history::fallback::FallbackHistory;
1708 use crate::serve::state::ServerState;
1709 use std::sync::Arc;
1710 use tokio_util::sync::CancellationToken;
1711
1712 let mut cluster = ClusterConfig::disabled();
1713 cluster.enabled = true;
1714 let cfg = ServeConfig {
1715 listen: "127.0.0.1:0".parse().unwrap(),
1716 auth: AuthMode::None,
1717 max_concurrent_runs: 4,
1718 max_queued_runs: 4,
1719 default_config_path: None,
1720 history: HistoryBackendSpec::Memory,
1721 cors_origins: vec![],
1722 body_limit_bytes: 1_048_576,
1723 shutdown_grace: Duration::from_secs(60),
1724 retain_terminal_runs: Duration::from_secs(60),
1725 idempotency_retention: Duration::from_secs(60),
1726 lease_ttl: Duration::from_secs(30),
1727 probe_timeout: Duration::from_secs(10),
1728 env_file: None,
1729 no_env_file: false,
1730 log_level: "info".into(),
1731 ui_enabled: true,
1732 cluster,
1733 triggers_path: None,
1734 };
1735 let history = Arc::new(FallbackHistory::degraded_at_startup(
1737 Duration::from_secs(60),
1738 "test",
1739 )) as Arc<dyn RunHistory>;
1740 assert!(history.degraded());
1741 let state = ServerState::new(
1742 &cfg,
1743 None,
1744 CancellationToken::new(),
1745 history,
1746 crate::serve::logs::LogHub::new(),
1747 None,
1748 #[cfg(feature = "triggers")]
1749 crate::serve::triggers::health::TriggersHandle::empty(),
1750 );
1751 let req = SubmitRequest {
1752 config: "version: 1\npipeline:\n source: { type: csv, config: { path: x.csv } }\n sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1753 config_format: ConfigFormatWire::Yaml,
1754 name: None,
1755 labels: BTreeMap::new(),
1756 timeout_secs: None,
1757 doctor_first: false,
1758 idempotency_key: None,
1759 clock: None,
1760 };
1761 let err = submit(state.clone(), req, admin_actor()).await.unwrap_err();
1762 assert!(
1763 matches!(err, ServeError::Unavailable(_)),
1764 "expected 503 Unavailable, got {err:?}"
1765 );
1766 assert_eq!(state.registry().queued(), 0);
1768 }
1769
1770 #[cfg(feature = "serve-history-sqlite")]
1776 mod shards {
1777 use super::*;
1778 use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1779 use crate::serve::history::RunHistory;
1780 use crate::serve::history::sqlite::SqliteHistory;
1781 use crate::serve::load::{ConfigFormat, load_submission};
1782 use crate::serve::state::ServerState;
1783 use faucet_core::ShardSpec;
1784 use std::collections::BTreeMap;
1785 use std::sync::Arc;
1786 use tokio_util::sync::CancellationToken;
1787
1788 async fn sqlite_state(dir: &std::path::Path) -> ServerState {
1789 let url = format!("sqlite://{}/h.db", dir.display());
1790 let history = Arc::new(
1791 SqliteHistory::connect(
1792 &url,
1793 Duration::from_secs(300),
1794 Duration::from_secs(300),
1795 "inst-test".into(),
1796 )
1797 .await
1798 .expect("sqlite history"),
1799 ) as Arc<dyn RunHistory>;
1800 let cfg = ServeConfig {
1801 listen: "127.0.0.1:0".parse().unwrap(),
1802 auth: AuthMode::None,
1803 max_concurrent_runs: 4,
1804 max_queued_runs: 4,
1805 default_config_path: None,
1806 history: HistoryBackendSpec::Memory,
1807 cors_origins: vec![],
1808 body_limit_bytes: 1_048_576,
1809 shutdown_grace: Duration::from_secs(60),
1810 retain_terminal_runs: Duration::from_secs(60),
1811 idempotency_retention: Duration::from_secs(60),
1812 lease_ttl: Duration::from_secs(30),
1813 probe_timeout: Duration::from_secs(10),
1814 env_file: None,
1815 no_env_file: false,
1816 log_level: "info".into(),
1817 ui_enabled: true,
1818 cluster: crate::serve::cluster::ClusterConfig::disabled(),
1819 triggers_path: None,
1820 };
1821 ServerState::new(
1822 &cfg,
1823 None,
1824 CancellationToken::new(),
1825 history,
1826 crate::serve::logs::LogHub::new(),
1827 None,
1828 #[cfg(feature = "triggers")]
1829 crate::serve::triggers::health::TriggersHandle::empty(),
1830 )
1831 }
1832
1833 async fn loaded(yaml: &str) -> LoadedSubmission {
1834 load_submission(yaml, ConfigFormat::Yaml, None)
1835 .await
1836 .expect("load submission")
1837 }
1838
1839 async fn seed_run(state: &ServerState, run_id: &str, status: RunStatus) {
1840 let mut rec = RunRecord::queued(run_id.into(), None, BTreeMap::new(), None, Utc::now());
1841 rec.status = status;
1842 rec.config_body = Some("version: 1".into());
1843 state.history().upsert(&rec).await.expect("seed run");
1844 }
1845
1846 #[tokio::test]
1847 async fn coordinate_matrix_run_is_not_shardable() {
1848 let dir = tempfile::tempdir().unwrap();
1850 let state = sqlite_state(dir.path()).await;
1851 let l = loaded(
1852 "version: 1\nname: m\nmatrix:\n - id: a\n - id: b\npipeline:\n \
1853 source: { type: rest, config: { url: \"http://localhost/x\" } }\n \
1854 sink: { type: stdout, config: {} }\n",
1855 )
1856 .await;
1857 assert!(!coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1858 }
1859
1860 #[tokio::test]
1861 async fn coordinate_non_shardable_source_runs_whole() {
1862 let dir = tempfile::tempdir().unwrap();
1864 let state = sqlite_state(dir.path()).await;
1865 let input = dir.path().join("in.csv");
1866 std::fs::write(&input, "id\n1\n").unwrap();
1867 let l = loaded(&format!(
1868 "version: 1\npipeline:\n \
1869 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n \
1870 sink: {{ type: stdout, config: {{}} }}\n",
1871 input.display()
1872 ))
1873 .await;
1874 assert!(!coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1875 }
1876
1877 #[tokio::test]
1878 async fn coordinate_s3_source_inserts_shards_and_marks_sharded() {
1879 let dir = tempfile::tempdir().unwrap();
1880 let state = sqlite_state(dir.path()).await;
1881 seed_run(&state, "r", RunStatus::Running).await;
1882 let l = loaded(
1883 "version: 1\npipeline:\n \
1884 source: { type: s3, config: { bucket: my-bucket, prefix: null, \
1885 region: null, endpoint_url: null, file_format: json_lines, \
1886 max_objects: null, concurrency: 10 } }\n \
1887 sink: { type: stdout, config: {} }\n",
1888 )
1889 .await;
1890 assert!(coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1891 let prog = state.history().shard_progress("r").await.unwrap();
1893 assert_eq!(prog.total, 4);
1894 assert_eq!(prog.pending, 4);
1895 assert_eq!(
1896 state.history().get("r").await.unwrap().unwrap().status,
1897 RunStatus::Sharded
1898 );
1899 }
1900
1901 #[tokio::test]
1902 async fn at_least_once_risky_sinks_flags_append_cluster_and_shard() {
1903 let append = loaded(
1906 "version: 1\npipeline:\n \
1907 source: { type: rest, config: { url: \"http://localhost/x\" } }\n \
1908 sink: { type: stdout, config: {} }\n",
1909 )
1910 .await;
1911 assert!(at_least_once_risky_sinks(&append, false, false).is_empty());
1913 assert_eq!(
1915 at_least_once_risky_sinks(&append, true, false),
1916 vec!["stdout"]
1917 );
1918 assert_eq!(
1920 at_least_once_risky_sinks(&append, false, true),
1921 vec!["stdout"]
1922 );
1923 }
1924
1925 #[tokio::test]
1926 async fn at_least_once_risky_sinks_safe_for_upsert_and_exactly_once() {
1927 let upsert = loaded(
1929 "version: 1\npipeline:\n \
1930 source: { type: postgres, config: { connection_url: \"postgres://x\", \
1931 query: \"select 1\" } }\n \
1932 sink: { type: postgres, config: { connection_url: \"postgres://y\", \
1933 table_name: t, column_mapping: auto_map, write_mode: upsert, key: [id] } }\n",
1934 )
1935 .await;
1936 assert!(at_least_once_risky_sinks(&upsert, true, true).is_empty());
1937
1938 let eo = loaded(
1940 "version: 1\ndelivery: exactly_once\npipeline:\n \
1941 source: { type: postgres-cdc, config: {} }\n \
1942 sink: { type: sqlite, config: {} }\n \
1943 state: { type: file, config: { path: \"/tmp/x.json\" } }\n",
1944 )
1945 .await;
1946 assert!(at_least_once_risky_sinks(&eo, true, false).is_empty());
1947 }
1948
1949 async fn seed_sharded_with_shards(state: &ServerState, run_id: &str, n: usize) {
1950 use crate::serve::history::ShardInsert;
1951 seed_run(state, run_id, RunStatus::Sharded).await;
1952 let shards: Vec<ShardInsert> = (0..n)
1953 .map(|i| ShardInsert {
1954 shard_id: i.to_string(),
1955 descriptor: serde_json::json!({ "i": i }),
1956 size_estimate: None,
1957 })
1958 .collect();
1959 state
1960 .history()
1961 .insert_shards(run_id, &shards)
1962 .await
1963 .unwrap();
1964 let claimed = state.history().claim_shards(n).await.unwrap();
1966 assert_eq!(claimed.len(), n);
1967 }
1968
1969 #[tokio::test]
1970 async fn maybe_finalize_parent_completes_when_all_shards_succeed() {
1971 let dir = tempfile::tempdir().unwrap();
1972 let state = sqlite_state(dir.path()).await;
1973 seed_sharded_with_shards(&state, "r", 3).await;
1974 for i in 0..3 {
1975 state
1976 .history()
1977 .finalize_shard("r", &i.to_string(), true)
1978 .await
1979 .unwrap();
1980 }
1981 maybe_finalize_parent(&state, "r").await;
1982 assert_eq!(
1983 state.history().get("r").await.unwrap().unwrap().status,
1984 RunStatus::Completed
1985 );
1986 }
1987
1988 #[tokio::test]
1989 async fn maybe_finalize_parent_fails_when_a_shard_fails() {
1990 let dir = tempfile::tempdir().unwrap();
1991 let state = sqlite_state(dir.path()).await;
1992 seed_sharded_with_shards(&state, "r", 2).await;
1993 state
1994 .history()
1995 .finalize_shard("r", "0", true)
1996 .await
1997 .unwrap();
1998 state
1999 .history()
2000 .finalize_shard("r", "1", false)
2001 .await
2002 .unwrap();
2003 maybe_finalize_parent(&state, "r").await;
2004 assert_eq!(
2005 state.history().get("r").await.unwrap().unwrap().status,
2006 RunStatus::Failed
2007 );
2008 }
2009
2010 #[tokio::test]
2011 async fn maybe_finalize_parent_keeps_sharded_until_all_terminal() {
2012 let dir = tempfile::tempdir().unwrap();
2013 let state = sqlite_state(dir.path()).await;
2014 seed_sharded_with_shards(&state, "r", 2).await;
2015 state
2017 .history()
2018 .finalize_shard("r", "0", true)
2019 .await
2020 .unwrap();
2021 maybe_finalize_parent(&state, "r").await;
2022 assert_eq!(
2023 state.history().get("r").await.unwrap().unwrap().status,
2024 RunStatus::Sharded
2025 );
2026 }
2027
2028 #[tokio::test]
2029 async fn finalize_sharded_parent_is_status_fenced_and_idempotent() {
2030 let dir = tempfile::tempdir().unwrap();
2034 let state = sqlite_state(dir.path()).await;
2035 seed_sharded_with_shards(&state, "r", 2).await;
2036
2037 let first = state
2038 .history()
2039 .finalize_sharded_parent("r", RunStatus::Completed, Utc::now(), None)
2040 .await
2041 .unwrap();
2042 assert!(first, "first finalize wins");
2043 assert_eq!(
2044 state.history().get("r").await.unwrap().unwrap().status,
2045 RunStatus::Completed
2046 );
2047
2048 let second = state
2050 .history()
2051 .finalize_sharded_parent("r", RunStatus::Failed, Utc::now(), Some("late".into()))
2052 .await
2053 .unwrap();
2054 assert!(!second, "second finalize is a no-op");
2055 let r = state.history().get("r").await.unwrap().unwrap();
2056 assert_eq!(r.status, RunStatus::Completed, "status not overwritten");
2057 assert!(r.error.is_none(), "late error must not be stamped");
2058
2059 let missing = state
2061 .history()
2062 .finalize_sharded_parent("does-not-exist", RunStatus::Completed, Utc::now(), None)
2063 .await
2064 .unwrap();
2065 assert!(!missing);
2066 }
2067
2068 #[tokio::test]
2069 async fn execute_shard_runs_a_csv_to_jsonl_shard() {
2070 let dir = tempfile::tempdir().unwrap();
2071 let state = sqlite_state(dir.path()).await;
2072 let input = dir.path().join("in.csv");
2073 std::fs::write(&input, "id,name\n1,alice\n2,bob\n").unwrap();
2074 let output = dir.path().join("out.jsonl");
2075 let yaml = format!(
2076 "version: 1\npipeline:\n \
2077 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n \
2078 sink: {{ type: jsonl, config: {{ path: \"{}\" }} }}\n",
2079 input.display(),
2080 output.display()
2081 );
2082 let l = loaded(&yaml).await;
2083 let ok = execute_shard(
2086 &state,
2087 l,
2088 "r",
2089 "0",
2090 ShardSpec::whole(),
2091 CancellationToken::new(),
2092 None,
2093 None,
2094 Utc::now(),
2095 )
2096 .await;
2097 assert!(ok, "csv→jsonl shard should complete");
2098 let written = std::fs::read_to_string(&output).unwrap();
2099 assert_eq!(written.lines().count(), 2, "both rows written");
2100 assert!(written.contains("alice") && written.contains("bob"));
2101 }
2102
2103 #[tokio::test]
2104 async fn resume_claimed_shard_executes_and_finalizes_parent() {
2105 let dir = tempfile::tempdir().unwrap();
2109 let state = sqlite_state(dir.path()).await;
2110 let input = dir.path().join("in.csv");
2111 std::fs::write(&input, "id,name\n1,alice\n").unwrap();
2112 let output = dir.path().join("out.jsonl");
2113 let yaml = format!(
2114 "version: 1\npipeline:\n \
2115 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n \
2116 sink: {{ type: jsonl, config: {{ path: \"{}\" }} }}\n",
2117 input.display(),
2118 output.display()
2119 );
2120 let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2122 rec.status = RunStatus::Sharded;
2123 rec.config_body = Some(yaml);
2124 state.history().upsert(&rec).await.unwrap();
2125 use crate::serve::history::ShardInsert;
2126 state
2127 .history()
2128 .insert_shards(
2129 "r",
2130 &[ShardInsert {
2131 shard_id: "0".into(),
2132 descriptor: serde_json::Value::Null,
2133 size_estimate: None,
2134 }],
2135 )
2136 .await
2137 .unwrap();
2138 let claimed = state.history().claim_shards(1).await.unwrap();
2139 assert_eq!(claimed.len(), 1);
2140
2141 resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2142
2143 let mut status = RunStatus::Sharded;
2146 for _ in 0..100 {
2147 tokio::time::sleep(Duration::from_millis(50)).await;
2148 status = state.history().get("r").await.unwrap().unwrap().status;
2149 if status.is_terminal() {
2150 break;
2151 }
2152 }
2153 assert_eq!(status, RunStatus::Completed, "shard ran → parent completed");
2154 assert!(output.exists(), "shard wrote its output");
2155 }
2156
2157 #[tokio::test]
2158 async fn resume_claimed_shard_with_no_config_fails_the_shard() {
2159 let dir = tempfile::tempdir().unwrap();
2162 let state = sqlite_state(dir.path()).await;
2163 let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2164 rec.status = RunStatus::Sharded; state.history().upsert(&rec).await.unwrap();
2166 use crate::serve::history::ShardInsert;
2167 state
2168 .history()
2169 .insert_shards(
2170 "r",
2171 &[ShardInsert {
2172 shard_id: "0".into(),
2173 descriptor: serde_json::Value::Null,
2174 size_estimate: None,
2175 }],
2176 )
2177 .await
2178 .unwrap();
2179 let claimed = state.history().claim_shards(1).await.unwrap();
2180 resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2181
2182 let mut status = RunStatus::Sharded;
2183 for _ in 0..100 {
2184 tokio::time::sleep(Duration::from_millis(50)).await;
2185 status = state.history().get("r").await.unwrap().unwrap().status;
2186 if status.is_terminal() {
2187 break;
2188 }
2189 }
2190 assert_eq!(status, RunStatus::Failed, "no-config shard → parent failed");
2191 }
2192
2193 #[tokio::test]
2194 async fn coordinate_returns_err_when_source_build_fails() {
2195 let dir = tempfile::tempdir().unwrap();
2198 let state = sqlite_state(dir.path()).await;
2199 let l = loaded(
2201 "version: 1\npipeline:\n \
2202 source: { type: s3, config: { bucket: b } }\n \
2203 sink: { type: stdout, config: {} }\n",
2204 )
2205 .await;
2206 assert!(coordinate_sharded_run(&state, "r", &l, 4).await.is_err());
2207 }
2208
2209 #[tokio::test]
2210 async fn execute_shard_returns_false_on_malformed_resilience() {
2211 let dir = tempfile::tempdir().unwrap();
2214 let state = sqlite_state(dir.path()).await;
2215 let input = dir.path().join("in.csv");
2216 std::fs::write(&input, "id\n1\n").unwrap();
2217 let yaml = format!(
2218 "version: 1\nresilience:\n retry:\n max_attempts: 0\npipeline:\n \
2219 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n \
2220 sink: {{ type: stdout, config: {{}} }}\n",
2221 input.display()
2222 );
2223 let l = loaded(&yaml).await;
2224 let ok = execute_shard(
2225 &state,
2226 l,
2227 "r",
2228 "0",
2229 ShardSpec::whole(),
2230 CancellationToken::new(),
2231 None,
2232 None,
2233 Utc::now(),
2234 )
2235 .await;
2236 assert!(!ok, "malformed resilience → shard fails fast");
2237 }
2238
2239 #[tokio::test]
2240 async fn resume_claimed_shard_with_unloadable_config_fails_the_shard() {
2241 let dir = tempfile::tempdir().unwrap();
2244 let state = sqlite_state(dir.path()).await;
2245 let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2246 rec.status = RunStatus::Sharded;
2247 rec.config_body = Some("this: is: not: valid: yaml: [".into());
2248 state.history().upsert(&rec).await.unwrap();
2249 use crate::serve::history::ShardInsert;
2250 state
2251 .history()
2252 .insert_shards(
2253 "r",
2254 &[ShardInsert {
2255 shard_id: "0".into(),
2256 descriptor: serde_json::Value::Null,
2257 size_estimate: None,
2258 }],
2259 )
2260 .await
2261 .unwrap();
2262 let claimed = state.history().claim_shards(1).await.unwrap();
2263 resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2264
2265 let mut status = RunStatus::Sharded;
2266 for _ in 0..100 {
2267 tokio::time::sleep(Duration::from_millis(50)).await;
2268 status = state.history().get("r").await.unwrap().unwrap().status;
2269 if status.is_terminal() {
2270 break;
2271 }
2272 }
2273 assert_eq!(
2274 status,
2275 RunStatus::Failed,
2276 "unloadable config → parent failed"
2277 );
2278 }
2279
2280 #[tokio::test]
2283 async fn request_cancel_flags_a_sharded_parent() {
2284 let dir = tempfile::tempdir().unwrap();
2287 let state = sqlite_state(dir.path()).await;
2288 seed_run(&state, "r", RunStatus::Sharded).await;
2289 state.history().request_cancel("r").await.unwrap();
2292 use crate::serve::history::ShardInsert;
2294 state
2295 .history()
2296 .insert_shards(
2297 "r",
2298 &[ShardInsert {
2299 shard_id: "0".into(),
2300 descriptor: serde_json::Value::Null,
2301 size_estimate: None,
2302 }],
2303 )
2304 .await
2305 .unwrap();
2306 let claimed = state.history().claim_shards(1).await.unwrap();
2307 assert_eq!(claimed.len(), 1, "shard claimed (running, owned)");
2308
2309 let flagged = state.history().pending_shard_cancellations().await.unwrap();
2310 assert_eq!(
2311 flagged,
2312 vec!["r".to_string()],
2313 "the flagged sharded parent's run id is returned for its running shard"
2314 );
2315 }
2316
2317 #[tokio::test]
2318 async fn pending_shard_cancellations_filters_unflagged_and_pending_shards() {
2319 let dir = tempfile::tempdir().unwrap();
2320 let state = sqlite_state(dir.path()).await;
2321 use crate::serve::history::ShardInsert;
2322 let one = |id: &str| {
2323 vec![ShardInsert {
2324 shard_id: id.into(),
2325 descriptor: serde_json::Value::Null,
2326 size_estimate: None,
2327 }]
2328 };
2329
2330 seed_run(&state, "A", RunStatus::Sharded).await;
2333 state.history().request_cancel("A").await.unwrap();
2334 state.history().insert_shards("A", &one("0")).await.unwrap();
2335 seed_run(&state, "C", RunStatus::Sharded).await;
2336 state.history().insert_shards("C", &one("0")).await.unwrap();
2337 let claimed = state.history().claim_shards(8).await.unwrap();
2338 assert_eq!(claimed.len(), 2, "A and C shards claimed (running)");
2339
2340 seed_run(&state, "B", RunStatus::Sharded).await;
2344 state.history().request_cancel("B").await.unwrap();
2345 state.history().insert_shards("B", &one("0")).await.unwrap();
2346
2347 let flagged = state.history().pending_shard_cancellations().await.unwrap();
2348 assert_eq!(
2349 flagged,
2350 vec!["A".to_string()],
2351 "only A (flagged + a running owned shard); B pending-shard, C unflagged"
2352 );
2353 }
2354
2355 #[tokio::test]
2358 async fn finalize_sweep_completes_an_all_success_sharded_parent() {
2359 let dir = tempfile::tempdir().unwrap();
2360 let state = sqlite_state(dir.path()).await;
2361 seed_sharded_with_shards(&state, "r", 3).await;
2362 for i in 0..3 {
2363 state
2364 .history()
2365 .finalize_shard("r", &i.to_string(), true)
2366 .await
2367 .unwrap();
2368 }
2369 let n = state
2371 .history()
2372 .finalize_completed_sharded_parents()
2373 .await
2374 .unwrap();
2375 assert_eq!(n, 1, "one sharded parent finalized");
2376 let rec = state.history().get("r").await.unwrap().unwrap();
2377 assert_eq!(rec.status, RunStatus::Completed);
2378 assert!(rec.finished_at.is_some());
2379 assert!(rec.error.is_none());
2380
2381 assert_eq!(
2383 state
2384 .history()
2385 .finalize_completed_sharded_parents()
2386 .await
2387 .unwrap(),
2388 0,
2389 "already-terminal parent is not re-finalized"
2390 );
2391 }
2392
2393 #[tokio::test]
2394 async fn finalize_sweep_fails_a_parent_with_a_failed_shard() {
2395 let dir = tempfile::tempdir().unwrap();
2396 let state = sqlite_state(dir.path()).await;
2397 seed_sharded_with_shards(&state, "r", 3).await;
2398 state
2399 .history()
2400 .finalize_shard("r", "0", true)
2401 .await
2402 .unwrap();
2403 state
2404 .history()
2405 .finalize_shard("r", "1", false)
2406 .await
2407 .unwrap();
2408 state
2409 .history()
2410 .finalize_shard("r", "2", true)
2411 .await
2412 .unwrap();
2413 let n = state
2414 .history()
2415 .finalize_completed_sharded_parents()
2416 .await
2417 .unwrap();
2418 assert_eq!(n, 1);
2419 let rec = state.history().get("r").await.unwrap().unwrap();
2420 assert_eq!(rec.status, RunStatus::Failed);
2421 assert!(rec.finished_at.is_some());
2422 assert_eq!(rec.error.as_deref(), Some("1/3 shard(s) failed"));
2423 }
2424
2425 #[tokio::test]
2426 async fn finalize_sweep_leaves_a_not_all_terminal_parent_sharded() {
2427 let dir = tempfile::tempdir().unwrap();
2428 let state = sqlite_state(dir.path()).await;
2429 seed_sharded_with_shards(&state, "r", 2).await;
2430 state
2432 .history()
2433 .finalize_shard("r", "0", true)
2434 .await
2435 .unwrap();
2436 let n = state
2437 .history()
2438 .finalize_completed_sharded_parents()
2439 .await
2440 .unwrap();
2441 assert_eq!(n, 0, "parent with a still-running shard is not finalized");
2442 assert_eq!(
2443 state.history().get("r").await.unwrap().unwrap().status,
2444 RunStatus::Sharded
2445 );
2446 }
2447 }
2448}