eba4fd1
use std::path::PathBuf;
use tokio::io::{AsyncBufRead, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
use tokio::process::Child;
use crate::server::error::{ServerError, ServerResult};
use crate::server::limits::DEFAULT_OUTPUT_CAP_BYTES;
pub struct ChildPackReader<R: AsyncBufRead + Unpin + Send> {
reader: R,
child: Option<Child>,
finished: bool,
}
impl<R: AsyncBufRead + Unpin + Send> ChildPackReader<R> {
pub fn new(reader: R, child: Option<Child>) -> Self {
Self {
reader,
child,
finished: false,
}
}
pub fn finished(&self) -> bool {
self.finished
}
pub fn child_handle(&mut self) -> Option<&mut Child> {
self.child.as_mut()
}
pub async fn next_chunk(&mut self, buf: &mut [u8]) -> ServerResult<usize> {
if self.finished {
return Ok(0);
}
let n = self
.reader
.read(buf)
.await
.map_err(|source| ServerError::Io {
path: PathBuf::new(),
source,
})?;
if n == 0 {
self.finished = true;
}
Ok(n)
}
pub async fn read_to_end_capped(&mut self, cap: Option<u64>) -> ServerResult<StreamOutcome> {
let cap = cap.unwrap_or(DEFAULT_OUTPUT_CAP_BYTES as u64);
let mut buf = vec![0u8; 32 * 1024];
let mut total: u64 = 0;
loop {
let n = self.next_chunk(&mut buf).await?;
if n == 0 {
return Ok(StreamOutcome::Eof);
}
total = total.saturating_add(n as u64);
if total > cap {
return Err(ServerError::PayloadTooLarge {
cap_bytes: cap,
actual_bytes: total,
});
}
}
}
pub async fn collect_pack_stats(&mut self, cap: Option<u64>) -> ServerResult<PackStats> {
let mut stats = PackStats::default();
let mut buf = vec![0u8; 32 * 1024];
loop {
let n = self.next_chunk(&mut buf).await?;
if n == 0 {
break;
}
stats.bytes_streamed = stats.bytes_streamed.saturating_add(n as u64);
if let Some(limit) = cap
&& stats.bytes_streamed > limit
{
return Err(ServerError::PayloadTooLarge {
cap_bytes: limit,
actual_bytes: stats.bytes_streamed,
});
}
}
Ok(stats)
}
}
impl<R: AsyncBufRead + Unpin + Send> Drop for ChildPackReader<R> {
fn drop(&mut self) {
if let Some(mut child) = self.child.take() {
// If the stream was cancelled mid-flight (not cleanly finished),
// kill the child so it cannot linger as a runaway process.
if !self.finished {
let _ = child.start_kill();
}
reap_child(child);
}
}
}
pub struct ChildPackWriter<W: AsyncWrite + Unpin + Send> {
writer: W,
child: Option<Child>,
finished: bool,
}
impl<W: AsyncWrite + Unpin + Send> ChildPackWriter<W> {
pub fn new(writer: W, child: Option<Child>) -> Self {
Self {
writer,
child,
finished: false,
}
}
pub fn finished(&self) -> bool {
self.finished
}
pub fn child_handle(&mut self) -> Option<&mut Child> {
self.child.as_mut()
}
pub async fn write_chunk(&mut self, buf: &[u8]) -> ServerResult<usize> {
if self.finished {
return Err(ServerError::Pipe(
"cannot write to a finished child writer".to_string(),
));
}
self.writer
.write_all(buf)
.await
.map_err(|source| ServerError::Io {
path: PathBuf::new(),
source,
})?;
Ok(buf.len())
}
pub async fn copy_from<R>(&mut self, mut reader: R) -> ServerResult<StreamOutcome>
where
R: AsyncRead + Unpin + Send,
{
let mut buf = vec![0u8; 32 * 1024];
let mut total: u64 = 0;
loop {
let n = reader
.read(&mut buf)
.await
.map_err(|source| ServerError::Io {
path: PathBuf::new(),
source,
})?;
if n == 0 {
break;
}
self.write_chunk(&buf[..n]).await?;
total = total.saturating_add(n as u64);
}
self.writer
.flush()
.await
.map_err(|source| ServerError::Io {
path: PathBuf::new(),
source,
})?;
Ok(StreamOutcome::Bytes(total))
}
pub async fn finish(mut self) -> ServerResult<()> {
self.writer
.shutdown()
.await
.map_err(|source| ServerError::Io {
path: PathBuf::new(),
source,
})?;
self.finished = true;
Ok(())
}
}
impl<W: AsyncWrite + Unpin + Send> Drop for ChildPackWriter<W> {
fn drop(&mut self) {
if let Some(mut child) = self.child.take() {
if !self.finished {
let _ = child.start_kill();
}
reap_child(child);
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct PackStats {
pub bytes_streamed: u64,
pub objects_reported: u64,
}
impl PackStats {
pub fn is_empty(&self) -> bool {
self.bytes_streamed == 0 && self.objects_reported == 0
}
}
#[derive(Debug, Clone, Default)]
pub struct IngestReport {
pub stats: PackStats,
pub refs_updated: Vec<String>,
}
impl IngestReport {
pub fn new(stats: PackStats, refs_updated: Vec<String>) -> Self {
Self {
stats,
refs_updated,
}
}
pub fn empty() -> Self {
Self::default()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StreamOutcome {
Eof,
Bytes(u64),
}
impl StreamOutcome {
pub fn bytes(&self) -> Option<u64> {
match self {
Self::Eof => None,
Self::Bytes(n) => Some(*n),
}
}
}
pub type BoxedAsyncRead = Box<dyn AsyncRead + Unpin + Send + Sync>;
pub type BoxedAsyncWrite = Box<dyn AsyncWrite + Unpin + Send + Sync>;
pub type BoxedAsyncBufRead = Box<dyn AsyncBufRead + Unpin + Send + Sync>;
pub fn empty_reader() -> BoxedAsyncRead {
Box::new(tokio::io::empty())
}
pub fn sink_writer() -> BoxedAsyncWrite {
Box::new(tokio::io::sink())
}
pub fn buf_empty_reader() -> BoxedAsyncBufRead {
Box::new(BufReader::new(tokio::io::empty()))
}
/// Reap a child process without blocking the current (potentially sync) drop
/// path. Spawns a reaper task if a tokio runtime is available; otherwise the
/// child is abandoned to the OS (best-effort — cannot `wait` without a runtime).
fn reap_child(mut child: Child) {
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
let _ = child.wait().await;
});
}
}