Fix and simplify tar downloads

This commit is contained in:
Joscha 2023-08-12 16:06:20 +02:00
parent 81328fcf04
commit 4f63b02509
2 changed files with 35 additions and 43 deletions

View file

@ -1,3 +1,5 @@
mod tree;
use tracing::error; use tracing::error;
use crate::config::Config; use crate::config::Config;

View file

@ -5,9 +5,9 @@ use std::{io, path::PathBuf};
use axum::body::Bytes; use axum::body::Bytes;
use flate2::read::GzDecoder; use flate2::read::GzDecoder;
use futures::{Stream, StreamExt}; use futures::{Stream, StreamExt};
use reqwest::Response;
use tempfile::TempDir; use tempfile::TempDir;
use tokio::{select, sync::mpsc}; use tokio::sync::mpsc;
use tracing::debug;
use crate::somehow; use crate::somehow;
@ -41,47 +41,37 @@ impl io::Read for ReceiverReader {
} }
} }
pub struct UnpackedTree { async fn receive_bytes(
pub hash: String,
pub dir: TempDir,
}
impl UnpackedTree {
async fn stream(
mut stream: impl Stream<Item = reqwest::Result<Bytes>> + Unpin, mut stream: impl Stream<Item = reqwest::Result<Bytes>> + Unpin,
tx: mpsc::Sender<Bytes>, tx: mpsc::Sender<Bytes>,
) -> somehow::Result<()> { ) -> somehow::Result<()> {
while let Some(bytes) = stream.next().await { while let Some(bytes) = stream.next().await {
tx.send(bytes?).await?; tx.send(bytes?).await?;
} }
Ok(()) Ok(())
} }
fn unpack(rx: mpsc::Receiver<Bytes>, path: PathBuf) -> somehow::Result<()> { fn unpack_archive(rx: mpsc::Receiver<Bytes>, path: PathBuf) -> somehow::Result<()> {
let reader = ReceiverReader::new(rx); let reader = ReceiverReader::new(rx);
let reader = GzDecoder::new(reader); let reader = GzDecoder::new(reader);
let mut reader = tar::Archive::new(reader); let mut reader = tar::Archive::new(reader);
reader.unpack(path)?; reader.unpack(path)?;
Ok(()) Ok(())
} }
pub async fn download(url: &str, hash: String) -> somehow::Result<Self> { pub async fn download(response: Response) -> somehow::Result<TempDir> {
let dir = TempDir::new()?; let stream = response.error_for_status()?.bytes_stream();
debug!(
"Downloading and unpacking {url} to {}", let dir = TempDir::new()?;
dir.path().display() let path = dir.path().to_path_buf();
); let (tx, rx) = mpsc::channel(1);
let (tx, rx) = mpsc::channel(1);
let stream = reqwest::get(url).await?.bytes_stream(); let (received, unpacked) = tokio::join!(
receive_bytes(stream, tx),
let path = dir.path().to_path_buf(); tokio::task::spawn_blocking(move || unpack_archive(rx, path)),
let unpack_task = tokio::task::spawn_blocking(move || Self::unpack(rx, path)); );
received?;
select! { unpacked??;
r = Self::stream(stream, tx) => r?,
r = unpack_task => r??, Ok(dir)
}
Ok(Self { hash, dir })
}
} }