122 lines
3.5 KiB
Rust
122 lines
3.5 KiB
Rust
//! Stream gzipped tar-ed repository worktrees.
|
|
|
|
use std::{
|
|
io::{self, BufWriter, Read},
|
|
sync::Arc,
|
|
};
|
|
|
|
use axum::{body::Bytes, BoxError};
|
|
use flate2::{write::GzEncoder, Compression};
|
|
use futures::TryStream;
|
|
use gix::{
|
|
bstr::ByteSlice, objs::tree::EntryMode, prelude::ObjectIdExt, worktree::stream::Entry,
|
|
ObjectId, ThreadSafeRepository,
|
|
};
|
|
use tokio::sync::mpsc;
|
|
use tokio_stream::wrappers::ReceiverStream;
|
|
use tracing::{debug, warn};
|
|
|
|
const BLOCK_SIZE: usize = 1024 * 1024;
|
|
const COMPRESSION_LEVEL: Compression = Compression::fast();
|
|
|
|
struct SenderWriter(mpsc::Sender<Result<Bytes, BoxError>>);
|
|
|
|
impl SenderWriter {
|
|
fn new(tx: mpsc::Sender<Result<Bytes, BoxError>>) -> Self {
|
|
Self(tx)
|
|
}
|
|
|
|
fn finish(self) {}
|
|
}
|
|
|
|
impl io::Write for SenderWriter {
|
|
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
|
match self.0.blocking_send(Ok(Bytes::copy_from_slice(buf))) {
|
|
Ok(()) => Ok(buf.len()),
|
|
Err(_) => Err(io::ErrorKind::ConnectionAborted.into()),
|
|
}
|
|
}
|
|
|
|
fn flush(&mut self) -> io::Result<()> {
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
fn write_entry(
|
|
mut entry: Entry<'_>,
|
|
writer: &mut tar::Builder<impl io::Write>,
|
|
) -> Result<(), BoxError> {
|
|
let mut header = tar::Header::new_gnu();
|
|
header.set_entry_type(match entry.mode {
|
|
EntryMode::Tree | EntryMode::Commit => tar::EntryType::Directory,
|
|
EntryMode::Blob | EntryMode::BlobExecutable => tar::EntryType::Regular,
|
|
EntryMode::Link => tar::EntryType::Symlink,
|
|
});
|
|
header.set_mode(match entry.mode {
|
|
EntryMode::BlobExecutable => 0o755, // rwxr-xr-x
|
|
_ => 0o644, // rw-r--r--
|
|
});
|
|
|
|
if entry.mode == EntryMode::Link {
|
|
let mut buf = vec![];
|
|
entry.read_to_end(&mut buf)?;
|
|
header.set_size(0);
|
|
let path = gix::path::from_bstr(entry.relative_path());
|
|
let target = gix::path::from_bstr(buf.as_bstr());
|
|
writer.append_link(&mut header, path, target)?;
|
|
} else {
|
|
header.set_size(entry.bytes_remaining().unwrap_or(0) as u64);
|
|
let path = gix::path::from_bstr(entry.relative_path()).to_path_buf();
|
|
writer.append_data(&mut header, path, entry)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn write_worktree(
|
|
repo: Arc<ThreadSafeRepository>,
|
|
commit_id: ObjectId,
|
|
tx: mpsc::Sender<Result<Bytes, BoxError>>,
|
|
) -> Result<(), BoxError> {
|
|
let repo = repo.to_thread_local();
|
|
let tree_id = commit_id
|
|
.attach(&repo)
|
|
.object()?
|
|
.try_into_commit()?
|
|
.tree_id()?;
|
|
let (mut stream, _) = repo.worktree_stream(tree_id)?;
|
|
|
|
let writer = SenderWriter::new(tx);
|
|
let writer = BufWriter::with_capacity(BLOCK_SIZE, writer);
|
|
let writer = GzEncoder::new(writer, COMPRESSION_LEVEL);
|
|
let mut writer = tar::Builder::new(writer);
|
|
|
|
while let Some(entry) = stream.next_entry()? {
|
|
write_entry(entry, &mut writer)?;
|
|
}
|
|
|
|
writer
|
|
.into_inner()?
|
|
.finish()?
|
|
.into_inner()
|
|
.map_err(|e| e.into_error())?
|
|
.finish();
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub fn tar_and_gzip(
|
|
repo: Arc<ThreadSafeRepository>,
|
|
id: ObjectId,
|
|
) -> impl TryStream<Ok = Bytes, Error = BoxError> {
|
|
let (tx, rx) = mpsc::channel(1);
|
|
tokio::task::spawn_blocking(move || {
|
|
if let Err(e) = write_worktree(repo, id, tx.clone()) {
|
|
warn!("Error while streaming tar:\n{e:?}");
|
|
let _ = tx.blocking_send(Err(e));
|
|
} else {
|
|
debug!("Tar streamed successfully");
|
|
}
|
|
});
|
|
ReceiverStream::new(rx)
|
|
}
|