1use std::time::Duration;
12
13use futures_util::future::BoxFuture;
14
15use super::Session;
16use super::exec::{EndReason, ExecHandle, ExecResult, StreamEnd, StreamOptions};
17use super::sse::ExecEvent;
18use crate::error::{Error, WireKind};
19
20pub const DEFAULT_CLIENT_GRACE: Duration = Duration::from_secs(60);
26
27pub const NO_TIMEOUT_CEILING: Duration =
30 Duration::from_secs(crate::constants::MAX_DURATION_SEC as u64);
31
32pub type OutputSink = Box<dyn FnMut(ExecEvent) -> OutputFlow + Send>;
36
37pub type OutputFlow = BoxFuture<'static, std::ops::ControlFlow<()>>;
40
41#[derive(Clone, Debug)]
43pub struct CompletionOptions {
44 pub client_grace: Duration,
47 pub stream: StreamOptions,
49}
50
51impl Default for CompletionOptions {
52 fn default() -> Self {
53 Self {
54 client_grace: DEFAULT_CLIENT_GRACE,
55 stream: StreamOptions::default(),
56 }
57 }
58}
59
60#[derive(Clone, Debug, Eq, PartialEq)]
62pub enum KillAnswer {
63 Signalled,
65 AlreadyGone,
67 Failed(String),
69}
70
71#[derive(Clone, Debug, Eq, PartialEq)]
73pub struct ClientDeadline {
74 pub after: Duration,
76 pub kill: KillAnswer,
78 pub ack_error: Option<String>,
81}
82
83#[derive(Clone, Debug)]
89pub struct CompletionPlan {
90 deadline: Duration,
91 options: CompletionOptions,
92}
93
94impl CompletionPlan {
95 pub fn new(
98 request: &protocol::exec::StartRequest,
99 options: CompletionOptions,
100 ) -> Result<Self, Error> {
101 let deadline = match request.timeout_sec {
102 Some(seconds) => {
103 crate::duration::of_secs_f64(seconds)?.saturating_add(options.client_grace)
104 }
105 None => NO_TIMEOUT_CEILING,
106 };
107 Ok(Self { deadline, options })
108 }
109
110 pub fn deadline(&self) -> Duration {
112 self.deadline
113 }
114
115 pub async fn drive(
139 &self,
140 handle: &ExecHandle,
141 on_output: Option<OutputSink>,
142 ) -> Result<ExecResult, Error> {
143 let deadline = super::deadline_after(self.deadline);
144 if let Some(sink) = on_output {
145 match tokio::time::timeout_at(deadline, self.stream(handle, sink)).await {
146 Err(_elapsed) => return self.after_deadline(handle).await,
147 Ok(true) => match handle.ack().await {
148 Ok(result) if result.done() && result.outcome.is_some() => return Ok(result),
149 Ok(_) | Err(_) => {}
152 },
153 Ok(false) => {}
154 }
155 }
156 match wait_and_ack_until(handle, deadline).await {
157 Ok(result) => Ok(result),
158 Err(error) if error.wire_kind() == Some(WireKind::ExecTimeout) => {
159 self.after_deadline(handle).await
160 }
161 Err(error) => Err(error),
162 }
163 }
164
165 async fn stream(&self, handle: &ExecHandle, mut sink: OutputSink) -> bool {
172 let end = handle
173 .for_each_event_async(self.options.stream.clone(), |event| match event {
174 ExecEvent::Output { .. } => sink(event),
175 ExecEvent::Gap { .. } | ExecEvent::Exit(_) => {
176 Box::pin(std::future::ready(std::ops::ControlFlow::Continue(())))
177 }
178 })
179 .await;
180 matches!(
181 end,
182 Ok(StreamEnd {
183 reason: EndReason::Exited,
184 ..
185 })
186 )
187 }
188
189 async fn after_deadline(&self, handle: &ExecHandle) -> Result<ExecResult, Error> {
191 let kill = match handle.kill().await {
192 Ok(true) => KillAnswer::Signalled,
193 Ok(false) => KillAnswer::AlreadyGone,
194 Err(error) => KillAnswer::Failed(error.to_string()),
195 };
196 let grace = super::deadline_after(self.options.client_grace);
197 match wait_and_ack_until(handle, grace).await {
198 Ok(mut result) => {
199 result.client_deadline = Some(ClientDeadline {
200 after: self.deadline,
201 kill,
202 ack_error: None,
203 });
204 Ok(result)
205 }
206 Err(error) => Ok(ExecResult {
210 exec_id: handle.exec_id().to_string(),
211 phase: protocol::exec::Phase::Running,
212 outcome: None,
213 client_deadline: Some(ClientDeadline {
214 after: self.deadline,
215 kill,
216 ack_error: Some(error.to_string()),
217 }),
218 }),
219 }
220 }
221}
222
223const ACK_RETRY_INTERVAL: Duration = Duration::from_secs(1);
225
226async fn wait_and_ack_until(
236 handle: &ExecHandle,
237 deadline: tokio::time::Instant,
238) -> Result<ExecResult, Error> {
239 loop {
240 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
241 match handle.wait_and_ack(remaining).await {
242 Err(error) if error.retryable() && error.wire_kind() != Some(WireKind::ExecTimeout) => {
243 let left = deadline.saturating_duration_since(tokio::time::Instant::now());
244 if left.is_zero() {
245 return Err(Error::wire(
246 WireKind::ExecTimeout,
247 format!(
248 "exec {} was not acknowledged before the deadline; the last \
249 attempt failed: {error}",
250 handle.exec_id()
251 ),
252 ));
253 }
254 tokio::time::sleep(ACK_RETRY_INTERVAL.min(left)).await;
255 }
256 other => return other,
257 }
258 }
259}
260
261impl Session {
262 pub async fn run_to_completion(
264 &self,
265 request: protocol::exec::StartRequest,
266 options: CompletionOptions,
267 on_output: Option<OutputSink>,
268 ) -> Result<ExecResult, Error> {
269 let plan = CompletionPlan::new(&request, options)?;
270 let handle = self.run(request).await?;
271 plan.drive(&handle, on_output).await
272 }
273}
274
275#[cfg(test)]
276mod tests {
277 use std::sync::{Arc, Mutex};
278
279 use base64::Engine as _;
280
281 use super::*;
282 use crate::session::testing::{Recorder, Reply, session_with};
283
284 fn output(offset: u64, bytes: &[u8]) -> Vec<u8> {
285 let encoded = base64::engine::general_purpose::STANDARD.encode(bytes);
286 format!(
287 "event: output\ndata: {{\"offset\":{offset},\"stream\":\"stdout\",\
288 \"output\":\"{encoded}\"}}\n\n"
289 )
290 .into_bytes()
291 }
292
293 fn exit_frame(code: Option<i32>, signal: Option<i32>, timed_out: bool, total: u64) -> Vec<u8> {
294 let event = protocol::exec::ExitEvent {
295 exit_code: code,
296 signal,
297 timed_out,
298 truncated: false,
299 writers_may_be_alive: false,
300 offset: total,
301 };
302 format!(
303 "event: exit\ndata: {}\n\n",
304 serde_json::to_string(&event).expect("serializes")
305 )
306 .into_bytes()
307 }
308
309 fn started() -> Reply {
310 Reply::ok(serde_json::json!({"exec_id": "e1", "phase": "running"}))
311 }
312
313 fn running() -> Reply {
314 Reply::ok(serde_json::json!({"exec_id": "e1", "phase": "running"}))
315 }
316
317 fn killed(killed: bool) -> Reply {
318 Reply::ok(serde_json::json!({"exec_id": "e1", "killed": killed}))
319 }
320
321 fn finished(phase: &str, code: Option<i32>, signal: Option<i32>, stdout: &str) -> Reply {
323 Reply::ok(serde_json::json!({
324 "exec_id": "e1", "phase": phase, "exit_code": code, "signal": signal,
325 "timed_out": false, "stdout": stdout, "stderr": "", "truncated": false,
326 "writers_may_be_alive": false,
327 }))
328 }
329
330 fn request(timeout_sec: Option<f64>) -> protocol::exec::StartRequest {
331 protocol::exec::StartRequest::new("e1", vec!["bash".into(), "-c".into(), "echo hi".into()])
332 .with_timeout_sec(timeout_sec)
333 }
334
335 type Seen = Arc<Mutex<Vec<(u64, Vec<u8>)>>>;
336
337 fn recording_sink(flow: std::ops::ControlFlow<()>) -> (OutputSink, Seen) {
339 let seen: Seen = Arc::new(Mutex::new(Vec::new()));
340 let into = Arc::clone(&seen);
341 let sink: OutputSink = Box::new(move |event| {
342 if let ExecEvent::Output { offset, data, .. } = event {
343 into.lock().expect("unpoisoned").push((offset, data));
344 }
345 Box::pin(std::future::ready(flow))
346 });
347 (sink, seen)
348 }
349
350 fn route(recorder: &Recorder) -> Vec<String> {
352 recorder
353 .requests()
354 .iter()
355 .map(|request| {
356 let path = request.path.split('?').next().unwrap_or_default();
357 format!("{} {path}", request.method)
358 })
359 .collect()
360 }
361
362 fn options(grace_secs: u64) -> CompletionOptions {
363 CompletionOptions {
364 client_grace: Duration::from_secs(grace_secs),
365 stream: StreamOptions {
366 reconnect: false,
367 ..StreamOptions::default()
368 },
369 }
370 }
371
372 #[tokio::test(start_paused = true)]
379 async fn a_streamed_exit_event_is_acked_directly_with_every_chunk_delivered() {
380 let recorder = Recorder::with([
381 started(),
382 Reply::Chunks(
383 200,
384 vec![
385 output(0, b"AB"),
386 output(2, b"CD"),
387 exit_frame(Some(0), None, false, 4),
388 ],
389 ),
390 finished("acked", Some(0), None, "ABCD"),
391 ]);
392 let (session, _, _) = session_with(Arc::clone(&recorder));
393 let (sink, seen) = recording_sink(std::ops::ControlFlow::Continue(()));
394
395 let result = session
396 .run_to_completion(request(None), options(60), Some(sink))
397 .await
398 .expect("completes");
399
400 assert_eq!(result.stdout(), "ABCD");
401 assert_eq!(result.posix_exit_code(), Some(0));
402 assert!(result.client_deadline.is_none());
403 assert_eq!(
404 *seen.lock().expect("unpoisoned"),
405 [(0, b"AB".to_vec()), (2, b"CD".to_vec())],
406 "the callback missed or reordered a chunk"
407 );
408 assert_eq!(
409 route(&recorder),
410 [
411 "POST /v1/exec/start",
412 "GET /v1/exec/e1/stream",
413 "POST /v1/exec/e1/ack"
414 ],
415 "an exit event needs one ack and nothing else"
416 );
417 }
418
419 #[tokio::test(start_paused = true)]
424 async fn a_cut_stream_falls_back_to_wait_and_ack() {
425 let recorder = Recorder::with([
426 started(),
427 Reply::Chunks(200, vec![output(0, b"AB")]),
428 finished("exited", Some(0), None, "ABCD"),
429 finished("acked", Some(0), None, "ABCD"),
430 ]);
431 let (session, _, _) = session_with(Arc::clone(&recorder));
432 let (sink, seen) = recording_sink(std::ops::ControlFlow::Continue(()));
433
434 let result = session
435 .run_to_completion(request(None), options(60), Some(sink))
436 .await
437 .expect("the fallback returns the result");
438
439 assert_eq!(
440 result.stdout(),
441 "ABCD",
442 "the ack's output, not the stream's"
443 );
444 assert_eq!(seen.lock().expect("unpoisoned").len(), 1);
445 assert_eq!(
446 route(&recorder),
447 [
448 "POST /v1/exec/start",
449 "GET /v1/exec/e1/stream",
450 "GET /v1/exec/e1",
451 "POST /v1/exec/e1/ack"
452 ],
453 "a cut stream must fall back to wait and ack"
454 );
455 }
456
457 #[tokio::test(start_paused = true)]
459 async fn a_failed_ack_after_the_exit_event_falls_back_to_wait_and_ack() {
460 let recorder = Recorder::with([
461 started(),
462 Reply::Chunks(
463 200,
464 vec![output(0, b"x"), exit_frame(Some(3), None, false, 1)],
465 ),
466 Reply::Cut("connection reset"),
467 finished("exited", Some(3), None, "x"),
468 finished("acked", Some(3), None, "x"),
469 ]);
470 let (session, _, _) = session_with(Arc::clone(&recorder));
471 let (sink, _) = recording_sink(std::ops::ControlFlow::Continue(()));
472
473 let result = session
474 .run_to_completion(request(None), options(60), Some(sink))
475 .await
476 .expect("the fallback returns the result");
477 assert_eq!(result.posix_exit_code(), Some(3));
478 assert_eq!(route(&recorder).len(), 5, "{:?}", route(&recorder));
479 }
480
481 #[tokio::test(start_paused = true)]
486 async fn a_callback_that_stops_still_gets_one_acked_result() {
487 let recorder = Recorder::with([
488 started(),
489 Reply::Chunks(
490 200,
491 vec![
492 output(0, b"A"),
493 output(1, b"B"),
494 exit_frame(Some(0), None, false, 2),
495 ],
496 ),
497 finished("exited", Some(0), None, "AB"),
498 finished("acked", Some(0), None, "AB"),
499 ]);
500 let (session, _, _) = session_with(Arc::clone(&recorder));
501 let (sink, seen) = recording_sink(std::ops::ControlFlow::Break(()));
502
503 let result = session
504 .run_to_completion(request(None), options(60), Some(sink))
505 .await
506 .expect("completes");
507 assert_eq!(result.stdout(), "AB");
508 assert_eq!(
509 seen.lock().expect("unpoisoned").len(),
510 1,
511 "delivery stopped"
512 );
513 assert_eq!(
514 route(&recorder).last().map(String::as_str),
515 Some("POST /v1/exec/e1/ack")
516 );
517 }
518
519 #[tokio::test(start_paused = true)]
521 async fn without_a_callback_the_call_waits_and_acks() {
522 let recorder = Recorder::with([
523 started(),
524 running(),
525 finished("exited", None, Some(11), ""),
526 finished("acked", None, Some(11), ""),
527 ]);
528 let (session, _, _) = session_with(Arc::clone(&recorder));
529
530 let result = session
531 .run_to_completion(request(Some(30.0)), options(60), None)
532 .await
533 .expect("completes");
534 assert_eq!(result.posix_exit_code(), Some(128 + 11));
535 assert!(!route(&recorder).iter().any(|r| r.contains("stream")));
536 }
537
538 #[tokio::test(start_paused = true)]
547 async fn a_dropped_ack_in_the_wait_is_retried() {
548 let recorder = Recorder::with([
549 started(),
550 finished("exited", Some(0), None, "out"),
551 Reply::Cut("connection reset"),
552 finished("exited", Some(0), None, "out"),
553 finished("acked", Some(0), None, "out"),
554 ]);
555 let (session, _, _) = session_with(Arc::clone(&recorder));
556
557 let result = session
558 .run_to_completion(request(Some(30.0)), options(60), None)
559 .await
560 .expect("the retried ack returns the result");
561 assert_eq!(result.stdout(), "out");
562 assert!(result.client_deadline.is_none());
563 assert_eq!(route(&recorder).len(), 5, "{:?}", route(&recorder));
564 }
565
566 #[tokio::test(start_paused = true)]
577 async fn a_client_deadline_kills_then_acks_within_the_grace() {
578 let recorder = Recorder::with([
579 started(),
580 running(),
581 running(),
582 running(),
583 killed(true),
584 finished("exited", None, Some(15), "partial"),
585 finished("acked", None, Some(15), "partial"),
586 ]);
587 let (session, _, _) = session_with(Arc::clone(&recorder));
588
589 let result = session
590 .run_to_completion(request(Some(1.0)), options(1), None)
591 .await
592 .expect("the killed exec's result");
593
594 let seen = route(&recorder);
595 let kill = seen
596 .iter()
597 .position(|r| r.ends_with("/kill"))
598 .expect("a kill");
599 let ack = seen
600 .iter()
601 .rposition(|r| r.ends_with("/ack"))
602 .expect("an ack");
603 assert!(kill < ack, "the ack went out before the kill: {seen:?}");
604 assert_eq!(result.stdout(), "partial");
605 assert_eq!(
606 result.client_deadline,
607 Some(ClientDeadline {
608 after: Duration::from_secs(2),
609 kill: KillAnswer::Signalled,
610 ack_error: None,
611 })
612 );
613 assert!(!result.synthesized());
614 assert_eq!(
615 result.posix_exit_code(),
616 Some(124),
617 "the client deadline ended it"
618 );
619 assert!(
620 result
621 .notes()
622 .iter()
623 .any(|note| note.contains("client deadline")),
624 "{:?}",
625 result.notes()
626 );
627 }
628
629 #[tokio::test(start_paused = true)]
632 async fn a_deadline_while_streaming_kills_then_acks() {
633 let recorder = Recorder::with([
634 started(),
635 Reply::Stalled(vec![output(0, b"tick\n")]),
636 killed(true),
637 finished("exited", None, Some(15), "tick\n"),
638 finished("acked", None, Some(15), "tick\n"),
639 ]);
640 let (session, _, _) = session_with(Arc::clone(&recorder));
641 let (sink, seen) = recording_sink(std::ops::ControlFlow::Continue(()));
642
643 let result = session
644 .run_to_completion(request(Some(5.0)), options(5), Some(sink))
645 .await
646 .expect("the killed exec's result");
647 assert_eq!(seen.lock().expect("unpoisoned").len(), 1);
648 assert_eq!(result.posix_exit_code(), Some(124));
649 assert_eq!(
650 route(&recorder)[2..],
651 [
652 "POST /v1/exec/e1/kill",
653 "GET /v1/exec/e1",
654 "POST /v1/exec/e1/ack"
655 ]
656 );
657 }
658
659 #[tokio::test(start_paused = true)]
662 async fn a_failed_kill_still_acks_and_keeps_the_commands_own_status() {
663 let recorder = Recorder::with([
664 started(),
665 running(),
666 running(),
667 running(),
668 Reply::Cut("connection reset"),
669 finished("exited", Some(0), None, "done"),
670 finished("acked", Some(0), None, "done"),
671 ]);
672 let (session, _, _) = session_with(Arc::clone(&recorder));
673
674 let result = session
675 .run_to_completion(request(Some(1.0)), options(1), None)
676 .await
677 .expect("the result arrives after the failed kill");
678 assert!(matches!(
679 result.client_deadline.as_ref().map(|d| &d.kill),
680 Some(KillAnswer::Failed(_))
681 ));
682 assert_eq!(
683 result.posix_exit_code(),
684 Some(0),
685 "nothing the client did ended it"
686 );
687 assert!(!result.synthesized());
688 }
689
690 #[tokio::test(start_paused = true)]
694 async fn a_failed_post_kill_ack_synthesizes_124() {
695 let recorder = Recorder::with([
696 started(),
697 running(),
698 running(),
699 running(),
700 killed(true),
701 running(),
702 running(),
703 ]);
704 let (session, _, _) = session_with(Arc::clone(&recorder));
705
706 let result = session
707 .run_to_completion(request(Some(1.0)), options(1), None)
708 .await
709 .expect("a synthesized result, not an error");
710 assert!(result.synthesized());
711 assert_eq!(result.posix_exit_code(), Some(124));
712 assert_eq!(
713 result.phase,
714 protocol::exec::Phase::Running,
715 "nothing was observed"
716 );
717 assert!(result.outcome.is_none());
718 assert!(
719 result
720 .notes()
721 .iter()
722 .any(|note| note.contains("synthesized")),
723 "{:?}",
724 result.notes()
725 );
726 assert_eq!(
727 route(&recorder).last().map(String::as_str),
728 Some("GET /v1/exec/e1")
729 );
730 }
731
732 #[tokio::test(start_paused = true)]
734 async fn a_failed_kill_and_a_failed_ack_synthesize_124_naming_both() {
735 let recorder = Recorder::with([
736 started(),
737 running(),
738 running(),
739 running(),
740 Reply::Cut("kill refused by the test"),
741 Reply::Cut("poll refused by the test"),
742 Reply::Cut("poll refused by the test"),
743 ]);
744 let (session, _, _) = session_with(Arc::clone(&recorder));
745
746 let result = session
747 .run_to_completion(request(Some(1.0)), options(1), None)
748 .await
749 .expect("a synthesized result");
750 assert!(result.synthesized());
751 assert_eq!(result.posix_exit_code(), Some(124));
752 let notes = result.notes().join("\n");
753 assert!(notes.contains("kill refused by the test"), "{notes}");
754 }
755
756 #[tokio::test]
761 async fn the_plan_computes_the_deadline_and_refuses_before_starting() {
762 let plan = CompletionPlan::new(&request(Some(30.0)), options(60)).expect("plans");
763 assert_eq!(plan.deadline(), Duration::from_secs(90));
764 let plan = CompletionPlan::new(&request(None), options(60)).expect("plans");
765 assert_eq!(plan.deadline(), NO_TIMEOUT_CEILING);
766
767 let recorder = Recorder::with([]);
768 let (session, _, _) = session_with(Arc::clone(&recorder));
769 for bad in [f64::NAN, -1.0, f64::INFINITY] {
770 let error = session
771 .run_to_completion(request(Some(bad)), options(60), None)
772 .await
773 .expect_err("refused");
774 assert_eq!(error.kind(), crate::ErrorKind::InvalidArg, "{bad}: {error}");
775 }
776 assert!(
777 recorder.requests().is_empty(),
778 "a refused plan started an exec"
779 );
780 }
781
782 #[tokio::test(start_paused = true)]
798 async fn a_client_deadline_longer_than_the_clock_can_hold_still_waits() {
799 let recorder = Recorder::with([
800 started(),
801 running(),
802 finished("exited", Some(0), None, "out"),
803 finished("acked", Some(0), None, "out"),
804 ]);
805 let (session, _, _) = session_with(Arc::clone(&recorder));
806
807 let result = session
808 .run_to_completion(request(Some(1e19)), options(60), None)
809 .await
810 .expect("a far deadline waits for the exit");
811 assert_eq!(result.stdout(), "out");
812 assert!(
813 result.client_deadline.is_none(),
814 "{:?}",
815 result.client_deadline
816 );
817 }
818
819 #[tokio::test(start_paused = true)]
836 async fn a_client_grace_longer_than_the_clock_can_hold_still_acks() {
837 let recorder = Recorder::with([
838 started(),
839 running(),
840 killed(true),
841 finished("exited", None, Some(15), "partial"),
842 finished("acked", None, Some(15), "partial"),
843 ]);
844 let (session, _, _) = session_with(Arc::clone(&recorder));
845 let plan = CompletionPlan {
846 deadline: Duration::ZERO,
847 options: CompletionOptions {
848 client_grace: Duration::MAX,
849 ..options(0)
850 },
851 };
852
853 let handle = session.run(request(None)).await.expect("starts");
854 let result = plan
855 .drive(&handle, None)
856 .await
857 .expect("the killed exec's result");
858 assert_eq!(result.stdout(), "partial");
859 assert_eq!(
860 result.client_deadline.map(|deadline| deadline.kill),
861 Some(KillAnswer::Signalled)
862 );
863 assert!(
864 route(&recorder).iter().any(|r| r.ends_with("/kill")),
865 "{:?}",
866 route(&recorder)
867 );
868 }
869}