use bytes::Bytes; use futures::TryStreamExt; use md5::Digest; use md5::Md5; use std::pin::Pin; use std::task::Context; use std::task::Poll; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tokio::io::ReadBuf; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio_util::io::ReaderStream; use tokio_util::io::StreamReader; use tracing::error; use tracing::warn; pub type FileReader = Box; pub type FileWriter = Box; pub const READ_BUFFER_SIZE: usize = 1024 * 1024; #[derive(Debug)] pub struct HttpFileWriter { wd: tokio::io::DuplexStream, err_rx: oneshot::Receiver, } impl HttpFileWriter { pub fn new(url: &str, disk: &str, volume: &str, path: &str, size: usize, append: bool) -> std::io::Result { let (rd, wd) = tokio::io::duplex(READ_BUFFER_SIZE); let (err_tx, err_rx) = oneshot::channel::(); let body = reqwest::Body::wrap_stream(ReaderStream::with_capacity(rd, READ_BUFFER_SIZE)); let url = url.to_owned(); let disk = disk.to_owned(); let volume = volume.to_owned(); let path = path.to_owned(); tokio::spawn(async move { let client = reqwest::Client::new(); if let Err(err) = client .put(format!( "{}/rustfs/rpc/put_file_stream?disk={}&volume={}&path={}&append={}&size={}", url, urlencoding::encode(&disk), urlencoding::encode(&volume), urlencoding::encode(&path), append, size )) .body(body) .send() .await .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)) { error!("HttpFileWriter put file err: {:?}", err); if let Err(er) = err_tx.send(err) { error!("HttpFileWriter tx.send err: {:?}", er); } } }); Ok(Self { wd, err_rx }) } } impl AsyncWrite for HttpFileWriter { #[tracing::instrument(level = "debug", skip(self, buf))] fn poll_write( mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, buf: &[u8], ) -> Poll> { if let Ok(err) = self.as_mut().err_rx.try_recv() { return Poll::Ready(Err(err)); } Pin::new(&mut self.wd).poll_write(cx, buf) } #[tracing::instrument(level = "debug", skip(self))] fn poll_flush(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll> { Pin::new(&mut self.wd).poll_flush(cx) } #[tracing::instrument(level = "debug", skip(self))] fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll> { Pin::new(&mut self.wd).poll_shutdown(cx) } } pub struct HttpFileReader { inner: FileReader, } impl HttpFileReader { pub async fn new(url: &str, disk: &str, volume: &str, path: &str, offset: usize, length: usize) -> std::io::Result { let resp = reqwest::Client::new() .get(format!( "{}/rustfs/rpc/read_file_stream?disk={}&volume={}&path={}&offset={}&length={}", url, urlencoding::encode(disk), urlencoding::encode(volume), urlencoding::encode(path), offset, length )) .send() .await .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; let inner = Box::new(StreamReader::new(resp.bytes_stream().map_err(std::io::Error::other))); Ok(Self { inner }) } } impl AsyncRead for HttpFileReader { fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { Pin::new(&mut self.inner).poll_read(cx, buf) } } pub struct EtagReader { inner: R, bytes_tx: mpsc::Sender, md5_rx: oneshot::Receiver, } impl EtagReader { pub fn new(inner: R) -> Self { let (bytes_tx, mut bytes_rx) = mpsc::channel::(8); let (md5_tx, md5_rx) = oneshot::channel::(); tokio::task::spawn_blocking(move || { let mut md5 = Md5::new(); while let Some(bytes) = bytes_rx.blocking_recv() { md5.update(&bytes); } let digest = md5.finalize(); let etag = hex_simd::encode_to_string(digest, hex_simd::AsciiCase::Lower); let _ = md5_tx.send(etag); }); EtagReader { inner, bytes_tx, md5_rx } } pub async fn etag(self) -> String { drop(self.inner); drop(self.bytes_tx); self.md5_rx.await.unwrap() } } impl AsyncRead for EtagReader { #[tracing::instrument(level = "debug", skip_all)] fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { let poll = Pin::new(&mut self.inner).poll_read(cx, buf); if let Poll::Ready(Ok(())) = &poll { if buf.remaining() == 0 { let bytes = buf.filled(); let bytes = Bytes::copy_from_slice(bytes); let tx = self.bytes_tx.clone(); tokio::spawn(async move { if let Err(e) = tx.send(bytes).await { warn!("EtagReader send error: {:?}", e); } }); } } poll } }