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::cost::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 = tokio::time::Instant::now() + 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 = tokio::time::Instant::now() + 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 {
332 exec_id: "e1".into(),
333 command: vec!["bash".into(), "-c".into(), "echo hi".into()],
334 shell: false.into(),
335 cwd: None,
336 env: Default::default(),
337 user: None,
338 group: None,
339 timeout_sec,
340 stdin: false,
341 reap_group_on_exit: false,
342 inherit_image_env: false,
343 }
344 }
345
346 type Seen = Arc<Mutex<Vec<(u64, Vec<u8>)>>>;
347
348 fn recording_sink(flow: std::ops::ControlFlow<()>) -> (OutputSink, Seen) {
350 let seen: Seen = Arc::new(Mutex::new(Vec::new()));
351 let into = Arc::clone(&seen);
352 let sink: OutputSink = Box::new(move |event| {
353 if let ExecEvent::Output { offset, data, .. } = event {
354 into.lock().expect("unpoisoned").push((offset, data));
355 }
356 Box::pin(std::future::ready(flow))
357 });
358 (sink, seen)
359 }
360
361 fn route(recorder: &Recorder) -> Vec<String> {
363 recorder
364 .requests()
365 .iter()
366 .map(|request| {
367 let path = request.path.split('?').next().unwrap_or_default();
368 format!("{} {path}", request.method)
369 })
370 .collect()
371 }
372
373 fn options(grace_secs: u64) -> CompletionOptions {
374 CompletionOptions {
375 client_grace: Duration::from_secs(grace_secs),
376 stream: StreamOptions {
377 reconnect: false,
378 ..StreamOptions::default()
379 },
380 }
381 }
382
383 #[tokio::test(start_paused = true)]
390 async fn a_streamed_exit_event_is_acked_directly_with_every_chunk_delivered() {
391 let recorder = Recorder::with([
392 started(),
393 Reply::Chunks(
394 200,
395 vec![
396 output(0, b"AB"),
397 output(2, b"CD"),
398 exit_frame(Some(0), None, false, 4),
399 ],
400 ),
401 finished("acked", Some(0), None, "ABCD"),
402 ]);
403 let (session, _, _) = session_with(Arc::clone(&recorder));
404 let (sink, seen) = recording_sink(std::ops::ControlFlow::Continue(()));
405
406 let result = session
407 .run_to_completion(request(None), options(60), Some(sink))
408 .await
409 .expect("completes");
410
411 assert_eq!(result.stdout(), "ABCD");
412 assert_eq!(result.posix_exit_code(), Some(0));
413 assert!(result.client_deadline.is_none());
414 assert_eq!(
415 *seen.lock().expect("unpoisoned"),
416 [(0, b"AB".to_vec()), (2, b"CD".to_vec())],
417 "the callback missed or reordered a chunk"
418 );
419 assert_eq!(
420 route(&recorder),
421 [
422 "POST /v1/exec/start",
423 "GET /v1/exec/e1/stream",
424 "POST /v1/exec/e1/ack"
425 ],
426 "an exit event needs one ack and nothing else"
427 );
428 }
429
430 #[tokio::test(start_paused = true)]
435 async fn a_cut_stream_falls_back_to_wait_and_ack() {
436 let recorder = Recorder::with([
437 started(),
438 Reply::Chunks(200, vec![output(0, b"AB")]),
439 finished("exited", Some(0), None, "ABCD"),
440 finished("acked", Some(0), None, "ABCD"),
441 ]);
442 let (session, _, _) = session_with(Arc::clone(&recorder));
443 let (sink, seen) = recording_sink(std::ops::ControlFlow::Continue(()));
444
445 let result = session
446 .run_to_completion(request(None), options(60), Some(sink))
447 .await
448 .expect("the fallback returns the result");
449
450 assert_eq!(
451 result.stdout(),
452 "ABCD",
453 "the ack's output, not the stream's"
454 );
455 assert_eq!(seen.lock().expect("unpoisoned").len(), 1);
456 assert_eq!(
457 route(&recorder),
458 [
459 "POST /v1/exec/start",
460 "GET /v1/exec/e1/stream",
461 "GET /v1/exec/e1",
462 "POST /v1/exec/e1/ack"
463 ],
464 "a cut stream must fall back to wait and ack"
465 );
466 }
467
468 #[tokio::test(start_paused = true)]
470 async fn a_failed_ack_after_the_exit_event_falls_back_to_wait_and_ack() {
471 let recorder = Recorder::with([
472 started(),
473 Reply::Chunks(
474 200,
475 vec![output(0, b"x"), exit_frame(Some(3), None, false, 1)],
476 ),
477 Reply::Cut("connection reset"),
478 finished("exited", Some(3), None, "x"),
479 finished("acked", Some(3), None, "x"),
480 ]);
481 let (session, _, _) = session_with(Arc::clone(&recorder));
482 let (sink, _) = recording_sink(std::ops::ControlFlow::Continue(()));
483
484 let result = session
485 .run_to_completion(request(None), options(60), Some(sink))
486 .await
487 .expect("the fallback returns the result");
488 assert_eq!(result.posix_exit_code(), Some(3));
489 assert_eq!(route(&recorder).len(), 5, "{:?}", route(&recorder));
490 }
491
492 #[tokio::test(start_paused = true)]
497 async fn a_callback_that_stops_still_gets_one_acked_result() {
498 let recorder = Recorder::with([
499 started(),
500 Reply::Chunks(
501 200,
502 vec![
503 output(0, b"A"),
504 output(1, b"B"),
505 exit_frame(Some(0), None, false, 2),
506 ],
507 ),
508 finished("exited", Some(0), None, "AB"),
509 finished("acked", Some(0), None, "AB"),
510 ]);
511 let (session, _, _) = session_with(Arc::clone(&recorder));
512 let (sink, seen) = recording_sink(std::ops::ControlFlow::Break(()));
513
514 let result = session
515 .run_to_completion(request(None), options(60), Some(sink))
516 .await
517 .expect("completes");
518 assert_eq!(result.stdout(), "AB");
519 assert_eq!(
520 seen.lock().expect("unpoisoned").len(),
521 1,
522 "delivery stopped"
523 );
524 assert_eq!(
525 route(&recorder).last().map(String::as_str),
526 Some("POST /v1/exec/e1/ack")
527 );
528 }
529
530 #[tokio::test(start_paused = true)]
532 async fn without_a_callback_the_call_waits_and_acks() {
533 let recorder = Recorder::with([
534 started(),
535 running(),
536 finished("exited", None, Some(11), ""),
537 finished("acked", None, Some(11), ""),
538 ]);
539 let (session, _, _) = session_with(Arc::clone(&recorder));
540
541 let result = session
542 .run_to_completion(request(Some(30.0)), options(60), None)
543 .await
544 .expect("completes");
545 assert_eq!(result.posix_exit_code(), Some(128 + 11));
546 assert!(!route(&recorder).iter().any(|r| r.contains("stream")));
547 }
548
549 #[tokio::test(start_paused = true)]
558 async fn a_dropped_ack_in_the_wait_is_retried() {
559 let recorder = Recorder::with([
560 started(),
561 finished("exited", Some(0), None, "out"),
562 Reply::Cut("connection reset"),
563 finished("exited", Some(0), None, "out"),
564 finished("acked", Some(0), None, "out"),
565 ]);
566 let (session, _, _) = session_with(Arc::clone(&recorder));
567
568 let result = session
569 .run_to_completion(request(Some(30.0)), options(60), None)
570 .await
571 .expect("the retried ack returns the result");
572 assert_eq!(result.stdout(), "out");
573 assert!(result.client_deadline.is_none());
574 assert_eq!(route(&recorder).len(), 5, "{:?}", route(&recorder));
575 }
576
577 #[tokio::test(start_paused = true)]
588 async fn a_client_deadline_kills_then_acks_within_the_grace() {
589 let recorder = Recorder::with([
590 started(),
591 running(),
592 running(),
593 running(),
594 killed(true),
595 finished("exited", None, Some(15), "partial"),
596 finished("acked", None, Some(15), "partial"),
597 ]);
598 let (session, _, _) = session_with(Arc::clone(&recorder));
599
600 let result = session
601 .run_to_completion(request(Some(1.0)), options(1), None)
602 .await
603 .expect("the killed exec's result");
604
605 let seen = route(&recorder);
606 let kill = seen
607 .iter()
608 .position(|r| r.ends_with("/kill"))
609 .expect("a kill");
610 let ack = seen
611 .iter()
612 .rposition(|r| r.ends_with("/ack"))
613 .expect("an ack");
614 assert!(kill < ack, "the ack went out before the kill: {seen:?}");
615 assert_eq!(result.stdout(), "partial");
616 assert_eq!(
617 result.client_deadline,
618 Some(ClientDeadline {
619 after: Duration::from_secs(2),
620 kill: KillAnswer::Signalled,
621 ack_error: None,
622 })
623 );
624 assert!(!result.synthesized());
625 assert_eq!(
626 result.posix_exit_code(),
627 Some(124),
628 "the client deadline ended it"
629 );
630 assert!(
631 result
632 .notes()
633 .iter()
634 .any(|note| note.contains("client deadline")),
635 "{:?}",
636 result.notes()
637 );
638 }
639
640 #[tokio::test(start_paused = true)]
643 async fn a_deadline_while_streaming_kills_then_acks() {
644 let recorder = Recorder::with([
645 started(),
646 Reply::Stalled(vec![output(0, b"tick\n")]),
647 killed(true),
648 finished("exited", None, Some(15), "tick\n"),
649 finished("acked", None, Some(15), "tick\n"),
650 ]);
651 let (session, _, _) = session_with(Arc::clone(&recorder));
652 let (sink, seen) = recording_sink(std::ops::ControlFlow::Continue(()));
653
654 let result = session
655 .run_to_completion(request(Some(5.0)), options(5), Some(sink))
656 .await
657 .expect("the killed exec's result");
658 assert_eq!(seen.lock().expect("unpoisoned").len(), 1);
659 assert_eq!(result.posix_exit_code(), Some(124));
660 assert_eq!(
661 route(&recorder)[2..],
662 [
663 "POST /v1/exec/e1/kill",
664 "GET /v1/exec/e1",
665 "POST /v1/exec/e1/ack"
666 ]
667 );
668 }
669
670 #[tokio::test(start_paused = true)]
673 async fn a_failed_kill_still_acks_and_keeps_the_commands_own_status() {
674 let recorder = Recorder::with([
675 started(),
676 running(),
677 running(),
678 running(),
679 Reply::Cut("connection reset"),
680 finished("exited", Some(0), None, "done"),
681 finished("acked", Some(0), None, "done"),
682 ]);
683 let (session, _, _) = session_with(Arc::clone(&recorder));
684
685 let result = session
686 .run_to_completion(request(Some(1.0)), options(1), None)
687 .await
688 .expect("the result arrives after the failed kill");
689 assert!(matches!(
690 result.client_deadline.as_ref().map(|d| &d.kill),
691 Some(KillAnswer::Failed(_))
692 ));
693 assert_eq!(
694 result.posix_exit_code(),
695 Some(0),
696 "nothing the client did ended it"
697 );
698 assert!(!result.synthesized());
699 }
700
701 #[tokio::test(start_paused = true)]
705 async fn a_failed_post_kill_ack_synthesizes_124() {
706 let recorder = Recorder::with([
707 started(),
708 running(),
709 running(),
710 running(),
711 killed(true),
712 running(),
713 running(),
714 ]);
715 let (session, _, _) = session_with(Arc::clone(&recorder));
716
717 let result = session
718 .run_to_completion(request(Some(1.0)), options(1), None)
719 .await
720 .expect("a synthesized result, not an error");
721 assert!(result.synthesized());
722 assert_eq!(result.posix_exit_code(), Some(124));
723 assert_eq!(
724 result.phase,
725 protocol::exec::Phase::Running,
726 "nothing was observed"
727 );
728 assert!(result.outcome.is_none());
729 assert!(
730 result
731 .notes()
732 .iter()
733 .any(|note| note.contains("synthesized")),
734 "{:?}",
735 result.notes()
736 );
737 assert_eq!(
738 route(&recorder).last().map(String::as_str),
739 Some("GET /v1/exec/e1")
740 );
741 }
742
743 #[tokio::test(start_paused = true)]
745 async fn a_failed_kill_and_a_failed_ack_synthesize_124_naming_both() {
746 let recorder = Recorder::with([
747 started(),
748 running(),
749 running(),
750 running(),
751 Reply::Cut("kill refused by the test"),
752 Reply::Cut("poll refused by the test"),
753 Reply::Cut("poll refused by the test"),
754 ]);
755 let (session, _, _) = session_with(Arc::clone(&recorder));
756
757 let result = session
758 .run_to_completion(request(Some(1.0)), options(1), None)
759 .await
760 .expect("a synthesized result");
761 assert!(result.synthesized());
762 assert_eq!(result.posix_exit_code(), Some(124));
763 let notes = result.notes().join("\n");
764 assert!(notes.contains("kill refused by the test"), "{notes}");
765 }
766
767 #[tokio::test]
772 async fn the_plan_computes_the_deadline_and_refuses_before_starting() {
773 let plan = CompletionPlan::new(&request(Some(30.0)), options(60)).expect("plans");
774 assert_eq!(plan.deadline(), Duration::from_secs(90));
775 let plan = CompletionPlan::new(&request(None), options(60)).expect("plans");
776 assert_eq!(plan.deadline(), NO_TIMEOUT_CEILING);
777
778 let recorder = Recorder::with([]);
779 let (session, _, _) = session_with(Arc::clone(&recorder));
780 for bad in [f64::NAN, -1.0, f64::INFINITY] {
781 let error = session
782 .run_to_completion(request(Some(bad)), options(60), None)
783 .await
784 .expect_err("refused");
785 assert_eq!(error.kind(), crate::ErrorKind::InvalidArg, "{bad}: {error}");
786 }
787 assert!(
788 recorder.requests().is_empty(),
789 "a refused plan started an exec"
790 );
791 }
792}