Re-use try_from_stream where possible

This commit is contained in:
asonix 2024-02-22 16:10:34 -06:00
parent 0ebee2a07c
commit 3a7d5b7bfb
3 changed files with 12 additions and 37 deletions

View File

@ -26,6 +26,7 @@ impl BytesStream {
}
}
#[tracing::instrument(skip(stream))]
pub(crate) async fn try_from_stream<S, E>(stream: S) -> Result<Self, E>
where
S: Stream<Item = Result<Bytes, E>>,
@ -35,6 +36,7 @@ impl BytesStream {
let mut bs = Self::new();
while let Some(bytes) = stream.try_next().await? {
tracing::trace!("try_from_stream: looping");
bs.add_bytes(bytes);
}

View File

@ -14,7 +14,6 @@ use actix_web::web::Bytes;
use futures_core::Stream;
use reqwest::Body;
use streem::IntoStreamer;
use tracing::{Instrument, Span};
mod hasher;
@ -29,25 +28,6 @@ pub(crate) struct Session {
identifier: Option<Arc<str>>,
}
#[tracing::instrument(skip(stream))]
async fn aggregate<S>(stream: S) -> Result<BytesStream, Error>
where
S: Stream<Item = Result<Bytes, Error>>,
{
let mut buf = BytesStream::new();
let stream = std::pin::pin!(stream);
let mut stream = stream.into_streamer();
while let Some(res) = stream.next().await {
tracing::trace!("aggregate: looping");
buf.add_bytes(res?);
}
Ok(buf)
}
async fn process_ingest<S>(
state: &State<S>,
stream: impl Stream<Item = Result<Bytes, Error>> + 'static,
@ -63,9 +43,12 @@ async fn process_ingest<S>(
where
S: Store,
{
let bytes = tokio::time::timeout(Duration::from_secs(60), aggregate(stream))
.await
.map_err(|_| UploadError::AggregateTimeout)??;
let bytes = tokio::time::timeout(
Duration::from_secs(60),
BytesStream::try_from_stream(stream),
)
.await
.map_err(|_| UploadError::AggregateTimeout)??;
let permit = crate::process_semaphore().acquire().await?;

View File

@ -1,7 +1,6 @@
use actix_web::web::Bytes;
use futures_core::Stream;
use std::{fmt::Debug, sync::Arc};
use streem::IntoStreamer;
use tokio::io::{AsyncRead, AsyncWrite};
use crate::{bytes_stream::BytesStream, error_code::ErrorCode, stream::LocalBoxStream};
@ -123,20 +122,11 @@ pub(crate) trait Store: Clone + Debug {
from_start: Option<u64>,
len: Option<u64>,
) -> Result<BytesStream, StoreError> {
let mut buf = BytesStream::new();
let stream = self.to_stream(identifier, from_start, len).await?;
let mut streamer = self
.to_stream(identifier, from_start, len)
.await?
.into_streamer();
while let Some(bytes) = streamer.try_next().await.map_err(StoreError::ReadStream)? {
tracing::trace!("to_bytes: looping");
buf.add_bytes(bytes);
}
Ok(buf)
BytesStream::try_from_stream(stream)
.await
.map_err(StoreError::ReadStream)
}
async fn read_into<Writer>(