Something went wrong. Try again.
A Forgejo Runner plugin that spawn jobs as Kubernetes pods
Something went wrong. Try again.
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070use std::{collections::VecDeque, future::Future};
use futures::{ FutureExt as _, Stream, StreamExt as _, future::{BoxFuture, pending},};use k8s_openapi::apimachinery::pkg::apis::meta::v1::Status as RemoteStatus;use kube::api::AttachedProcess;use tokio::{ io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _}, sync::mpsc,};use tonic::Status;
use crate::plugin::{CopyInChunk, CopyOutChunk, DataChunk, ExecOutput, data_chunk, exec_output};
const STDERR_LIMIT: usize = 16 * 1024;type Reader = Box<dyn AsyncRead + Unpin + Send>;type Writer = Box<dyn AsyncWrite + Unpin + Send>;type StatusFuture = BoxFuture<'static, Option<RemoteStatus>>;
/// Owns the WebSocket driver as well as its extracted streams. Dropping an/// operation requests local driver abortion, not termination of the remote command.pub struct Attachment { _process: Option<AttachedProcess>, stdin: Option<Writer>, stdout: Option<Reader>, stderr: Option<Reader>, status: Option<StatusFuture>, #[cfg(test)] _drop_guard: Option<tests::DropGuard>,}
impl From<AttachedProcess> for Attachment { fn from(mut process: AttachedProcess) -> Self { Self { stdin: process.stdin().map(|writer| -> Writer { Box::new(writer) }), stdout: process .stdout() .map(|reader| -> Reader { Box::new(reader) }), stderr: process .stderr() .map(|reader| -> Reader { Box::new(reader) }), status: process .take_status() .map(|status| -> StatusFuture { Box::pin(status) }), _process: Some(process), #[cfg(test)] _drop_guard: None, } }}
#[derive(Debug)]pub enum Failure { Remote(String), Stream(Status),}
/// Forwards data only; environment policy and terminal messages belong to the caller.pub async fn command( process: Attachment, output: &mpsc::Sender<Result<ExecOutput, Status>>,) -> Result<i32, Failure> { supervise(process, async { Ok(0) }, Output::Command(output)) .await .map(|completion| completion.exit_code)}
pub async fn upload_archive<Chunks>( mut process: Attachment, first_data: Vec<u8>, remaining: Chunks,) -> Result<usize, Failure>where Chunks: Stream<Item = Result<CopyInChunk, Status>> + Unpin,{ let stdin = process .stdin .take() .ok_or_else(|| internal("Kubernetes exec did not provide stdin"))?; supervise( process, write_archive(stdin, first_data, remaining), Output::Upload, ) .await .map(|completion| completion.bytes)}
pub async fn download_archive( process: Attachment, output: &mpsc::Sender<Result<CopyOutChunk, Status>>,) -> Result<(), Failure> { supervise(process, async { Ok(0) }, Output::Download(output)) .await .map(|_| ())}
#[derive(Clone, Copy)]enum Output<'sender> { Command(&'sender mpsc::Sender<Result<ExecOutput, Status>>), Upload, Download(&'sender mpsc::Sender<Result<CopyOutChunk, Status>>),}
impl Output<'_> { async fn closed(self) { match self { Self::Command(sender) => sender.closed().await, Self::Download(sender) => sender.closed().await, Self::Upload => pending().await, } }
fn cancelled(self) -> Failure { Failure::Stream(Status::cancelled(match self { Self::Command(_) => "Exec stream closed", Self::Upload => "CopyIn stream closed", Self::Download(_) => "CopyOut stream closed", })) }}
struct Completion { exit_code: i32, bytes: usize,}
enum InputFailure { Source(Status), Stdin(Status),}
async fn supervise( process: Attachment, input: impl Future<Output = Result<usize, InputFailure>>, output: Output<'_>,) -> Result<Completion, Failure> { let mut stderr = StderrTail::default(); let result = supervise_inner(process, input, output, &mut stderr).await; result.map_err(|failure| stderr.annotate(failure))}
async fn supervise_inner( mut process: Attachment, input: impl Future<Output = Result<usize, InputFailure>>, output: Output<'_>, stderr_tail: &mut StderrTail,) -> Result<Completion, Failure> { let stdout = process .stdout .take() .ok_or_else(|| internal("Kubernetes exec did not provide stdout"))?; let stderr = process .stderr .take() .ok_or_else(|| internal("Kubernetes exec did not provide stderr"))?; let mut status = process .status .take() .ok_or_else(|| internal("Kubernetes exec did not provide a status channel"))?; let capture = (!matches!(output, Output::Command(_))).then_some(stderr_tail); let drains = async { tokio::try_join!( pump(stdout, output, data_chunk::Stream::Stdout, None), pump(stderr, output, data_chunk::Stream::Stderr, capture), )?; Ok::<(), Failure>(()) }; tokio::pin!(drains); let mut drained = false; let mut bytes = None; let remote = { // Tar can consume its end markers and exit zero before the source's EOF. // Keep that upload alive; other statuses leave this scope and drop even // an upload blocked reading a chunk or writing to the bounded pipe. let mut succeeded = false; tokio::pin!(input); loop { if succeeded && bytes.is_some() { break Ok(0); } tokio::select! { biased; () = output.closed() => return Err(output.cancelled()), result = &mut input, if bytes.is_none() => match result { Ok(count) => bytes = Some(count), Err(error) => { // A ready remote failure can explain a broken stdin pipe. // Collect only ready diagnostics/status: an input error must // not wait for remote completion or output EOF. if !drained { let _ = drains.as_mut().now_or_never(); } if matches!(&error, InputFailure::Stdin(_)) && !succeeded && let Some(Some(remote)) = status.as_mut().now_or_never() { let exit_code = parse_exit_code(Some(remote))?; if exit_code != 0 { return Err(internal(format!("tar exited with status {exit_code}"))); } } return Err(match error { InputFailure::Source(error_status) | InputFailure::Stdin(error_status) => Failure::Stream(error_status), }); } }, remote = &mut status, if !succeeded => match parse_exit_code(remote) { Ok(0) => succeeded = true, result => break result, }, result = &mut drains, if !drained => { result?; drained = true; } } } }; if !drained { tokio::select! { biased; () = output.closed() => return Err(output.cancelled()), result = &mut drains => result?, } } // Do not join the driver: join consumes the owner and drops untaken // endpoints; cancelling that owned join would detach its driver task. let exit_code = remote?; if !matches!(output, Output::Command(_)) && exit_code != 0 { return Err(internal(format!("tar exited with status {exit_code}"))); } let bytes = bytes.ok_or_else(|| { internal("CopyIn tar completed before all input was consumed and stdin reached EOF") })?; Ok(Completion { exit_code, bytes })}
async fn write_archive( mut stdin: Writer, first_data: Vec<u8>, mut remaining: impl Stream<Item = Result<CopyInChunk, Status>> + Unpin,) -> Result<usize, InputFailure> { stdin.write_all(&first_data).await.map_err(|error| { InputFailure::Stdin(Status::internal(format!( "write first CopyIn chunk: {error}" ))) })?; let mut bytes = first_data.len(); while let Some(chunk) = remaining.next().await { let chunk = chunk.map_err(InputFailure::Source)?; if chunk.environment_id.is_some() || chunk.dest_path.is_some() { return Err(InputFailure::Source(Status::invalid_argument( "only the first CopyIn chunk may contain environment_id or dest_path", ))); } stdin.write_all(&chunk.data).await.map_err(|error| { InputFailure::Stdin(Status::internal(format!("write CopyIn chunk: {error}"))) })?; bytes = bytes .checked_add(chunk.data.len()) .ok_or_else(|| InputFailure::Source(Status::internal("CopyIn byte count overflow")))?; } stdin.shutdown().await.map_err(|error| { InputFailure::Stdin(Status::internal(format!("close CopyIn stdin: {error}"))) })?; drop(stdin); Ok(bytes)}
async fn pump( mut reader: Reader, output: Output<'_>, stream: data_chunk::Stream, mut tail: Option<&mut StderrTail>,) -> Result<(), Failure> { let mut buffer = vec![0; 8192]; loop { let count = reader .read(&mut buffer) .await .map_err(|error| internal(format!("read Kubernetes exec {stream:?}: {error}")))?; if count == 0 { return Ok(()); } let data = &buffer[..count]; if let Some(tail) = tail.as_mut() { tail.push(data); } match output { Output::Command(sender) => sender .send(Ok(ExecOutput { output: Some(exec_output::Output::Data(DataChunk { stream: stream as i32, data: data.to_vec(), })), })) .await .map_err(|_| output.cancelled())?, Output::Download(sender) if stream == data_chunk::Stream::Stdout => sender .send(Ok(CopyOutChunk { data: data.to_vec(), })) .await .map_err(|_| output.cancelled())?, Output::Upload | Output::Download(_) => {} } }}
#[derive(Default)]struct StderrTail { bytes: VecDeque<u8>, truncated: bool,}
impl StderrTail { fn push(&mut self, data: &[u8]) { for byte in data { if self.bytes.len() == STDERR_LIMIT { self.bytes.pop_front(); self.truncated = true; } self.bytes.push_back(*byte); } }
fn annotate(self, failure: Failure) -> Failure { if self.bytes.is_empty() { return failure; } let bytes: Vec<_> = self.bytes.into_iter().collect(); let marker = if self.truncated { "[stderr truncated; showing last 16384 bytes]\n" } else { "" }; let diagnostic = format!("\ntar stderr:\n{marker}{}", String::from_utf8_lossy(&bytes)); match failure { Failure::Remote(message) => Failure::Remote(format!("{message}{diagnostic}")), Failure::Stream(status) => Failure::Stream(Status::with_details_and_metadata( status.code(), format!("{}{diagnostic}", status.message()), status.details().to_vec().into(), status.metadata().clone(), )), } }}
fn internal(message: impl Into<String>) -> Failure { Failure::Stream(Status::internal(message))}
fn parse_exit_code(status: Option<RemoteStatus>) -> Result<i32, Failure> { let status = status .ok_or_else(|| Failure::Remote("Kubernetes exec ended without a status".to_owned()))?; if status.status.as_deref() == Some("Success") { return Ok(0); } if status.reason.as_deref() == Some("NonZeroExitCode") && let Some(cause) = status .details .as_ref() .and_then(|details| details.causes.as_ref()) .and_then(|causes| { causes .iter() .find(|cause| cause.reason.as_deref() == Some("ExitCode")) }) { return cause .message .as_deref() .ok_or_else(|| Failure::Remote("Kubernetes exit-code cause had no message".to_owned()))? .parse::<i32>() .map_err(|error| Failure::Remote(format!("invalid Kubernetes exit code: {error}"))); } Err(Failure::Remote(format!( "Kubernetes exec failed without a parseable exit code: reason={:?}, message={:?}", status.reason, status.message )))}
#[cfg(test)]mod tests { use std::{ sync::{ Arc, atomic::{AtomicBool, Ordering}, }, time::Duration, };
use futures::{ future::{Either, select}, stream, }; use k8s_openapi::apimachinery::pkg::apis::meta::v1::{StatusCause, StatusDetails}; use tokio::{ io::{DuplexStream, duplex}, sync::oneshot, time::timeout, };
use tokio_stream::wrappers::ReceiverStream;
use super::*;
pub(super) struct DropGuard(Arc<AtomicBool>);
impl Drop for DropGuard { fn drop(&mut self) { self.0.store(true, Ordering::SeqCst); } }
struct Script { stdin: DuplexStream, stdout: DuplexStream, stderr: DuplexStream, status: oneshot::Sender<RemoteStatus>, }
fn scripted() -> (Attachment, Script, Arc<AtomicBool>) { let (stdin, remote_stdin) = duplex(8); let (stdout, remote_stdout) = duplex(8); let (stderr, remote_stderr) = duplex(8); let (status_tx, status_rx) = oneshot::channel(); let dropped = Arc::new(AtomicBool::new(false)); ( Attachment { _process: None, stdin: Some(Box::new(stdin)), stdout: Some(Box::new(stdout)), stderr: Some(Box::new(stderr)), status: Some(Box::pin(async { status_rx.await.ok() })), _drop_guard: Some(DropGuard(Arc::clone(&dropped))), }, Script { stdin: remote_stdin, stdout: remote_stdout, stderr: remote_stderr, status: status_tx, }, dropped, ) }
fn exit_status(code: i32) -> RemoteStatus { if code == 0 { return RemoteStatus { status: Some("Success".to_owned()), ..RemoteStatus::default() }; } RemoteStatus { status: Some("Failure".to_owned()), reason: Some("NonZeroExitCode".to_owned()), details: Some(StatusDetails { causes: Some(vec![StatusCause { reason: Some("ExitCode".to_owned()), message: Some(code.to_string()), ..StatusCause::default() }]), ..StatusDetails::default() }), ..RemoteStatus::default() } }
async fn bounded<T>(future: impl Future<Output = T>) -> T { timeout(Duration::from_secs(2), future) .await .expect("scripted execution should not hang") }
fn stream_failure(failure: Failure, code: tonic::Code) -> Status { let Failure::Stream(status) = failure else { panic!("expected a stream failure, got {failure:?}"); }; assert_eq!(status.code(), code); status }
#[tokio::test] async fn command_forwards_both_streams_and_returns_zero_or_nonzero() { for code in [0, 23] { let (attachment, mut script, dropped) = scripted(); let (sender, mut receiver) = mpsc::channel(1); let operation = async { let result = command(attachment, &sender).await; drop(sender); result }; let driver = async { script.stdout.write_all(b"out").await.unwrap(); script.stderr.write_all(b"err").await.unwrap(); script.status.send(exit_status(code)).unwrap(); drop(script.stdout); drop(script.stderr); }; let consume = async { let mut stdout = Vec::new(); let mut stderr = Vec::new(); while let Some(message) = receiver.recv().await { let Some(exec_output::Output::Data(chunk)) = message.unwrap().output else { panic!("command must not send terminal messages"); }; if chunk.stream == data_chunk::Stream::Stdout as i32 { stdout.extend(chunk.data); } else { stderr.extend(chunk.data); } } (stdout, stderr) }; let (result, (), (stdout, stderr)) = bounded(async { tokio::join!(operation, driver, consume) }).await; assert_eq!(result.unwrap(), code); assert_eq!(stdout, b"out"); assert_eq!(stderr, b"err"); assert!(dropped.load(Ordering::SeqCst)); } }
#[tokio::test] async fn download_requires_zero_and_preserves_a_bounded_lossy_stderr_tail() { for code in [0, 2] { let (attachment, mut script, dropped) = scripted(); let (sender, mut receiver) = mpsc::channel(1); let mut diagnostic = vec![b'x'; 100]; diagnostic.extend_from_slice("\u{20ac}".as_bytes()); diagnostic.extend(vec![b'z'; STDERR_LIMIT - 2]); let operation = async { let result = download_archive(attachment, &sender).await; drop(sender); result }; let driver = async { script.stdout.write_all(b"archive bytes").await.unwrap(); script.stderr.write_all(&diagnostic).await.unwrap(); script.status.send(exit_status(code)).unwrap(); drop(script.stdout); drop(script.stderr); }; let consume = async { let mut archive = Vec::new(); while let Some(chunk) = receiver.recv().await { archive.extend(chunk.unwrap().data); } archive }; let (result, (), archive) = bounded(async { tokio::join!(operation, driver, consume) }).await; assert_eq!(archive, b"archive bytes"); if code == 0 { result.unwrap(); } else { let status = stream_failure(result.unwrap_err(), tonic::Code::Internal); assert!(status.message().contains("tar exited with status 2")); assert!( status .message() .contains("[stderr truncated; showing last 16384 bytes]") ); assert!(status.message().contains('\u{fffd}')); assert!(!status.message().contains(&"x".repeat(100))); assert!(status.message().ends_with(&"z".repeat(STDERR_LIMIT - 2))); } assert!(dropped.load(Ordering::SeqCst)); } }
#[tokio::test] async fn upload_drains_both_outputs_concurrently_and_requires_zero() { for code in [0, 3] { let (attachment, mut script, dropped) = scripted(); let first = vec![b'a'; 4096]; let remaining = stream::iter([Ok(CopyInChunk { data: vec![b'b'; 4096], ..CopyInChunk::default() })]); let operation = upload_archive(attachment, first, remaining); let driver = async { let mut uploaded = Vec::new(); let stdout = vec![b'o'; 32 * 1024]; let stderr = vec![b'e'; 32 * 1024]; let (input, out, err) = tokio::join!( script.stdin.read_to_end(&mut uploaded), script.stdout.write_all(&stdout), script.stderr.write_all(&stderr), ); assert_eq!(input.unwrap(), 8192); out.unwrap(); err.unwrap(); script.status.send(exit_status(code)).unwrap(); drop(script.stdout); drop(script.stderr); uploaded }; let (result, uploaded) = bounded(async { tokio::join!(operation, driver) }).await; assert_eq!(&uploaded[..4096], vec![b'a'; 4096]); assert_eq!(&uploaded[4096..], vec![b'b'; 4096]); if code == 0 { assert_eq!(result.unwrap(), 8192); } else { let status = stream_failure(result.unwrap_err(), tonic::Code::Internal); assert!(status.message().contains("tar exited with status 3")); assert!( status .message() .contains("[stderr truncated; showing last 16384 bytes]") ); assert!(status.message().ends_with(&"e".repeat(STDERR_LIMIT))); } assert!(dropped.load(Ordering::SeqCst)); } }
#[tokio::test] async fn missing_status_is_remote_for_all_operations() { let (attachment, script, dropped) = scripted(); drop(script); let (sender, _receiver) = mpsc::channel(1); let failure = bounded(command(attachment, &sender)).await.unwrap_err(); assert!( matches!(failure, Failure::Remote(message) if message.contains("without a status")) ); assert!(dropped.load(Ordering::SeqCst));
let (attachment, script, dropped) = scripted(); drop(script); let (sender, _receiver) = mpsc::channel(1); let failure = bounded(download_archive(attachment, &sender)) .await .unwrap_err(); assert!( matches!(failure, Failure::Remote(message) if message.contains("without a status")) ); assert!(dropped.load(Ordering::SeqCst));
let (attachment, script, dropped) = scripted(); drop(script.stdout); drop(script.stderr); drop(script.status); let failure = bounded(upload_archive(attachment, Vec::new(), stream::empty())) .await .unwrap_err(); assert!( matches!(failure, Failure::Remote(message) if message.contains("without a status")) ); assert!(dropped.load(Ordering::SeqCst)); }
#[tokio::test] async fn malformed_remote_status_preserves_stopped_container_detail() { let malformed = RemoteStatus { reason: Some("NonZeroExitCode".to_owned()), details: Some(StatusDetails { causes: Some(vec![StatusCause { reason: Some("ExitCode".to_owned()), message: Some("not a number".to_owned()), ..StatusCause::default() }]), ..StatusDetails::default() }), ..RemoteStatus::default() }; let stopped = RemoteStatus { message: Some("cannot exec in a stopped state".to_owned()), ..RemoteStatus::default() }; for (remote, expected) in [ (malformed, "invalid Kubernetes exit code"), (stopped, "cannot exec in a stopped state"), ] { let (attachment, script, dropped) = scripted(); script.status.send(remote).unwrap(); drop(script.stdout); drop(script.stderr); let (sender, _receiver) = mpsc::channel(1); let Failure::Remote(message) = bounded(command(attachment, &sender)).await.unwrap_err() else { panic!("remote status failure must remain distinguishable"); }; assert!(message.contains(expected)); assert!(dropped.load(Ordering::SeqCst)); } }
#[tokio::test] async fn remote_failure_interrupts_pending_upload_source_or_write() { for blocked_write in [false, true] { for (remote, remote_failure) in [ (Some(exit_status(9)), false), (None, true), (Some(RemoteStatus::default()), true), ] { let (attachment, script, dropped) = scripted(); let first = if blocked_write { vec![b'a'; 128] } else { Vec::new() }; let mut operation = Box::pin(upload_archive(attachment, first, stream::pending())); assert!(futures::poll!(operation.as_mut()).is_pending()); assert!(!dropped.load(Ordering::SeqCst)); if let Some(status) = remote { script.status.send(status).unwrap(); } else { drop(script.status); } drop(script.stdout); drop(script.stderr); let failure = bounded(operation).await.unwrap_err(); if remote_failure { assert!(matches!(failure, Failure::Remote(_))); } else { let status = stream_failure(failure, tonic::Code::Internal); assert!(status.message().contains("tar exited with status 9")); } assert!(dropped.load(Ordering::SeqCst)); let mut stdin = script.stdin; let mut bytes = Vec::new(); bounded(stdin.read_to_end(&mut bytes)).await.unwrap(); assert!(bytes.len() <= 8); } } }
#[tokio::test] async fn upload_zero_exit_before_source_eof_waits_then_succeeds() { let (attachment, mut script, dropped) = scripted(); let (chunks, receiver) = mpsc::channel(1); // Two zero-filled blocks are a complete empty tar archive. let archive = vec![0; 1024]; let mut uploaded = vec![0; archive.len()]; let mut operation = Box::pin(upload_archive( attachment, archive.clone(), ReceiverStream::new(receiver), )); let read = Box::pin(script.stdin.read_exact(&mut uploaded)); assert!(matches!( bounded(select(operation.as_mut(), read)).await, Either::Right((Ok(_), _)) )); assert_eq!(uploaded, archive); assert!(futures::poll!(operation.as_mut()).is_pending()); script.status.send(exit_status(0)).unwrap(); drop(script.stdin); drop(script.stdout); drop(script.stderr); assert!(futures::poll!(operation.as_mut()).is_pending()); assert!(!dropped.load(Ordering::SeqCst));
drop(chunks); assert_eq!(bounded(operation).await.unwrap(), archive.len()); assert!(dropped.load(Ordering::SeqCst)); }
#[tokio::test] async fn upload_zero_exit_keeps_pending_input_until_future_drop() { for blocked_write in [false, true] { let (attachment, script, dropped) = scripted(); let first = if blocked_write { vec![b'a'; 128] } else { Vec::new() }; let mut operation = Box::pin(upload_archive(attachment, first, stream::pending())); assert!(futures::poll!(operation.as_mut()).is_pending()); script.status.send(exit_status(0)).unwrap(); drop(script.stdout); drop(script.stderr); assert!(futures::poll!(operation.as_mut()).is_pending()); assert!(!dropped.load(Ordering::SeqCst)); drop(operation); assert!(dropped.load(Ordering::SeqCst)); let mut stdin = script.stdin; let mut bytes = Vec::new(); bounded(stdin.read_to_end(&mut bytes)).await.unwrap(); assert!(bytes.len() <= 8); } }
#[tokio::test] async fn upload_zero_exit_does_not_discard_later_payload_for_closed_stdin() { let (attachment, script, dropped) = scripted(); let (chunks, receiver) = mpsc::channel(1); let mut operation = Box::pin(upload_archive( attachment, Vec::new(), ReceiverStream::new(receiver), )); assert!(futures::poll!(operation.as_mut()).is_pending()); script.status.send(exit_status(0)).unwrap(); drop(script.stdin); drop(script.stdout); drop(script.stderr); assert!(futures::poll!(operation.as_mut()).is_pending()); chunks .send(Ok(CopyInChunk { data: vec![b'a'], ..CopyInChunk::default() })) .await .unwrap(); drop(chunks); let status = stream_failure(bounded(operation).await.unwrap_err(), tonic::Code::Internal); assert!(status.message().contains("write CopyIn chunk")); assert!(dropped.load(Ordering::SeqCst)); }
#[tokio::test] async fn upload_ready_nonzero_does_not_replace_forbidden_metadata_error() { for (environment_id, dest_path) in [(Some(String::new()), None), (None, Some(String::new()))] { let (attachment, script, dropped) = scripted(); script.status.send(exit_status(2)).unwrap(); let remaining = stream::iter([Ok(CopyInChunk { environment_id, dest_path, data: Vec::new(), })]); let status = stream_failure( bounded(upload_archive(attachment, Vec::new(), remaining)) .await .unwrap_err(), tonic::Code::InvalidArgument, ); assert!(status.message().contains("only the first CopyIn chunk")); assert!(dropped.load(Ordering::SeqCst)); } }
#[tokio::test] async fn upload_ready_nonzero_does_not_replace_source_error() { for code in [tonic::Code::Cancelled, tonic::Code::DataLoss] { let (attachment, script, dropped) = scripted(); script.status.send(exit_status(2)).unwrap(); let remaining = stream::iter([Err(Status::new(code, "archive input failed"))]); let status = stream_failure( bounded(upload_archive(attachment, Vec::new(), remaining)) .await .unwrap_err(), code, ); assert!(status.message().contains("archive input failed")); assert!(dropped.load(Ordering::SeqCst)); } }
#[tokio::test] async fn upload_broken_pipe_preserves_ready_tar_failure_and_stderr() { let (attachment, mut script, dropped) = scripted(); script.stderr.write_all(b"bad tar").await.unwrap(); script.status.send(exit_status(2)).unwrap(); drop(script.stdin); drop(script.stdout); drop(script.stderr);
let status = stream_failure( bounded(upload_archive(attachment, vec![b'a'], stream::pending())) .await .unwrap_err(), tonic::Code::Internal, ); assert!(status.message().contains("tar exited with status 2")); assert!(status.message().contains("bad tar")); assert!(dropped.load(Ordering::SeqCst)); }
#[tokio::test] async fn upload_errors_do_not_wait_for_status_or_output_eof() { for remote_code in [None, Some(0)] { for broken_pipe in [false, true] { let (attachment, mut script, dropped) = scripted(); script.stderr.write_all(b"context").await.unwrap(); if let Some(code) = remote_code { script.status.send(exit_status(code)).unwrap(); } let first = if broken_pipe { drop(script.stdin); vec![b'a'] } else { Vec::new() }; let remaining = stream::iter([Err(Status::data_loss("archive input failed"))]); let expected = if broken_pipe { tonic::Code::Internal } else { tonic::Code::DataLoss }; let status = stream_failure( bounded(upload_archive(attachment, first, remaining)) .await .unwrap_err(), expected, ); assert!(status.message().contains("context")); assert!(dropped.load(Ordering::SeqCst)); } } }
#[tokio::test] async fn upload_rejects_metadata_on_later_chunks() { for (environment_id, dest_path) in [(Some(String::new()), None), (None, Some(String::new()))] { let (attachment, _script, dropped) = scripted(); let remaining = stream::iter([Ok(CopyInChunk { environment_id, dest_path, data: Vec::new(), })]); let status = stream_failure( bounded(upload_archive(attachment, Vec::new(), remaining)) .await .unwrap_err(), tonic::Code::InvalidArgument, ); assert!(status.message().contains("only the first CopyIn chunk")); assert!(dropped.load(Ordering::SeqCst)); } }
#[tokio::test] async fn upload_preserves_input_error_status() { let (attachment, _script, dropped) = scripted(); let remaining = stream::iter([Err(Status::data_loss("archive input failed"))]); let status = stream_failure( bounded(upload_archive(attachment, Vec::new(), remaining)) .await .unwrap_err(), tonic::Code::DataLoss, ); assert!(status.message().contains("archive input failed")); assert!(dropped.load(Ordering::SeqCst)); }
#[tokio::test] async fn command_cancels_during_silence_and_when_output_is_backpressured() { for blocked in [false, true] { let (attachment, mut script, dropped) = scripted(); let (sender, receiver) = mpsc::channel(1); if blocked { sender.send(Ok(ExecOutput::default())).await.unwrap(); script.stdout.write_all(b"out").await.unwrap(); script.status.send(exit_status(0)).unwrap(); drop(script.stdout); drop(script.stderr); } let mut operation = Box::pin(command(attachment, &sender)); assert!(futures::poll!(operation.as_mut()).is_pending()); drop(receiver); stream_failure( bounded(operation).await.unwrap_err(), tonic::Code::Cancelled, ); assert!(dropped.load(Ordering::SeqCst)); } }
#[tokio::test] async fn download_cancels_during_silence_and_when_output_is_backpressured() { for blocked in [false, true] { let (attachment, mut script, dropped) = scripted(); let (sender, receiver) = mpsc::channel(1); if blocked { sender.send(Ok(CopyOutChunk::default())).await.unwrap(); script.stdout.write_all(b"tar").await.unwrap(); script.status.send(exit_status(0)).unwrap(); drop(script.stdout); drop(script.stderr); } let mut operation = Box::pin(download_archive(attachment, &sender)); assert!(futures::poll!(operation.as_mut()).is_pending()); drop(receiver); stream_failure( bounded(operation).await.unwrap_err(), tonic::Code::Cancelled, ); assert!(dropped.load(Ordering::SeqCst)); } }
#[tokio::test] async fn abrupt_future_drop_releases_every_operation_owner_and_endpoint() { let (attachment, script, dropped) = scripted(); let (sender, _receiver) = mpsc::channel(1); let mut operation = Box::pin(command(attachment, &sender)); assert!(futures::poll!(operation.as_mut()).is_pending()); drop(operation); assert!(dropped.load(Ordering::SeqCst)); assert!(script.status.send(exit_status(0)).is_err());
let (attachment, script, dropped) = scripted(); let (sender, _receiver) = mpsc::channel(1); let mut operation = Box::pin(download_archive(attachment, &sender)); assert!(futures::poll!(operation.as_mut()).is_pending()); drop(operation); assert!(dropped.load(Ordering::SeqCst)); assert!(script.status.send(exit_status(0)).is_err());
let (attachment, mut script, dropped) = scripted(); let mut operation = Box::pin(upload_archive( attachment, vec![b'a'; 128], stream::pending(), )); assert!(futures::poll!(operation.as_mut()).is_pending()); drop(operation); assert!(dropped.load(Ordering::SeqCst)); assert!(script.status.send(exit_status(0)).is_err()); let mut bytes = Vec::new(); bounded(script.stdin.read_to_end(&mut bytes)).await.unwrap(); assert_eq!(bytes.len(), 8); assert!(script.stdout.write_all(b"out").await.is_err()); assert!(script.stderr.write_all(b"err").await.is_err()); }
#[tokio::test] async fn completion_waits_for_output_drain_after_status() { let (attachment, mut script, dropped) = scripted(); let (sender, mut receiver) = mpsc::channel(1); script.status.send(exit_status(0)).unwrap(); let mut operation = Box::pin(command(attachment, &sender)); assert!(futures::poll!(operation.as_mut()).is_pending()); assert!(!dropped.load(Ordering::SeqCst)); script.stdout.write_all(b"late").await.unwrap(); drop(script.stdout); drop(script.stderr); assert_eq!(bounded(operation).await.unwrap(), 0); let Some(exec_output::Output::Data(chunk)) = receiver.recv().await.unwrap().unwrap().output else { panic!("expected the final stdout bytes"); }; assert_eq!(chunk.data, b"late"); assert!(dropped.load(Ordering::SeqCst)); }
#[tokio::test] async fn missing_required_endpoints_fail_without_waiting_for_status() { let (mut attachment, _script, dropped) = scripted(); attachment.stdout = None; let (sender, _receiver) = mpsc::channel(1); let status = stream_failure( bounded(command(attachment, &sender)).await.unwrap_err(), tonic::Code::Internal, ); assert!(status.message().contains("did not provide stdout")); assert!(dropped.load(Ordering::SeqCst));
let (mut attachment, _script, dropped) = scripted(); attachment.stdin = None; let status = stream_failure( bounded(upload_archive(attachment, Vec::new(), stream::empty())) .await .unwrap_err(), tonic::Code::Internal, ); assert!(status.message().contains("did not provide stdin")); assert!(dropped.load(Ordering::SeqCst)); }}