Upload and download with Axum, streaming

[frederikhors]

I'm trying to write a small axum router that uploads and downloads files.

I would like to stream those files directly to the cloud, but I don't know what to use for streaming.

In Go (my primarily language) I'm using http.ResponseWriter with io.Copy.

But in Rust I'm having real difficulties to understand what to use, this is the code:

• Cargo.toml:

[package]
name = "upload"
version = "0.1.0"
edition = "2021"

[dependencies]
axum = { version = "0.5.17", features = ["multipart"] }
futures = "0.3.25"
tokio = { version = "1.21.2", features = ["full"] }
tokio-util = { version = "0.7.4", features = ["io"] }
rust-s3 = "0.32.3"

• main.rs:

use axum::{
    body::{boxed, StreamBody},
    extract::{Multipart, Query},
    http::StatusCode,
    response::Response,
    routing::{on, MethodFilter},
    Router,
};
use futures::TryStreamExt;
use s3::{creds::Credentials, Bucket};
use std::{io, net::SocketAddr, pin::Pin};
use tokio::io::{AsyncRead, AsyncWrite, BufWriter};
use tokio_util::io::StreamReader;

#[tokio::main]
async fn main() {
    let router = Router::new()
        .route("/upload", on(MethodFilter::POST, upload))
        .route("/download/*key", on(MethodFilter::GET, download));

    let addr = SocketAddr::from(([127, 0, 0, 1], 8080));

    axum::Server::bind(&addr)
        .serve(router.into_make_service())
        .await
        .unwrap();
}

pub async fn upload(mut multipart: Multipart) -> Result<Response, StatusCode> {
    while let Some(field) = multipart.next_field().await.unwrap() {
        let filename = if let Some(filename) = field.file_name() {
            filename.to_string()
        } else {
            continue;
        };

        let bucket = Bucket::new(
            "test",
            "us-east-1".parse().unwrap(),
            Credentials::default().unwrap(),
        )
        .unwrap();

        let body_with_io_error = field.map_err(|err| io::Error::new(io::ErrorKind::Other, err));

        let body_reader = StreamReader::new(body_with_io_error);

        futures::pin_mut!(body_reader);

        put_file(bucket, &filename, body_reader);

        return Ok(Response::builder()
            .status(StatusCode::CREATED)
            .body(boxed("OK".to_string()))
            .unwrap());
    }

    Err(StatusCode::INTERNAL_SERVER_ERROR)
}

async fn put_file(
    bucket: Bucket,
    filename: &str,
    mut reader: Pin<&mut (dyn AsyncRead + Send)>,
) -> Result<(), ()> {
    bucket
        .put_object_stream(&mut reader, filename)
        .await
        .unwrap();

    Ok(())
}

pub async fn download(Query(params): Query<Vec<(String, String)>>) -> Result<Response, StatusCode> {
    let filename = params[0].1.to_string();

    let bucket = Bucket::new(
        "test",
        "us-east-1".parse().unwrap(),
        Credentials::default().unwrap(),
    )
    .unwrap();

    // I DON'T KNOW HOW TO START HERE! HELP!!! :smile:
    // What should I use here?

    let writer = BufWriter::new();
    // let writer = ReaderStream::new(reader);
    // futures::pin_mut!(writer);

    get_file(bucket, &filename, writer).await.unwrap();

    let response = Response::builder()
        .body(boxed(StreamBody::new(writer)))
        .unwrap();

    Ok(response)
}

async fn get_file(
    bucket: Bucket,
    filename: &str,
    mut writer: Pin<&mut (dyn AsyncWrite + Send)>,
) -> Result<(), ()> {
    bucket.get_object_stream(filename, &mut writer).await?;

    Ok(())
}

QUESTIONS

1. The upload function works but I don't know what these lines mean:

   let body_with_io_error = field.map_err(|err| io::Error::new(io::ErrorKind::Other, err));
   let body_reader = StreamReader::new(body_with_io_error);
   futures::pin_mut!(body_reader);

2. The download function doesn't work because I don't know what to use, how to create the writer that get_object_stream() needs.

---

I can PR examples and docs, as soon as I figure out how to do it, with your help.

Thanks!

[davidpdrsn]

// `field` is of type `axum::extract::multipart::Field` which implements `Stream` as seen here
// https://docs.rs/axum/latest/axum/extract/multipart/struct.Field.html#impl-Stream-for-Field%3C%27a%3E
// `map_err` comes from https://docs.rs/futures/latest/futures/future/trait.TryFutureExt.html#method.map_err which
// produces a new stream with a different error type, similarly to `Result::map_err`
let body_with_io_error = field.map_err(|err| io::Error::new(io::ErrorKind::Other, err));

// `StreamReader` comes from
// https://docs.rs/tokio-util/latest/tokio_util/io/struct.StreamReader.html and like the docs says
// it "Convert a Stream of byte chunks into an AsyncRead." so `body_reader` is something that
// implements `AsyncRead`
let body_reader = StreamReader::new(body_with_io_error);

// This pins value on the stack but it doesn't appear to be necessary
futures::pin_mut!(body_reader);

This works as well

let mut body_reader = StreamReader::new(body_with_io_error);

bucket
    .put_object_stream(&mut body_reader, filename)
    .await
    .unwrap();

get_object_stream takes something that implements AsyncWrite. If you click on that in the docs you can see what implements it https://docs.rs/tokio/1.20.1/tokio/io/trait.AsyncWrite.html#foreign-impls. I can't tell you which you need. Do you wanna write to a file then use tokio::fs::File. Do you wanna write to an in-memory buffer then use Vec<u8>.

[frederikhors]

Your suggestion works!

Questions

1. Why do you suggest to use spawn?:

   tokio::spawn(async move {
       let _ = bucket.get_object_stream(filename, &mut tx).await;
   });

2. I tried to use the same code with a simple file system adapter (instead of rust-s3) and even if it compiles when I use the download function it remains stucked on the line below, why? What do you think can be the issue?

   async fn get_file(
       &self,
       filename: &str,
       mut writer: Pin<&mut (dyn AsyncWrite + Send)>,
   ) -> Result<()> {
       let path = Path::new(&self.files_location).join(filename);
   
       let mut file = BufReader::new(File::open(path).await?);
   
       dbg!("It logs this");
   
       let bytes_copied = tokio::io::copy(&mut file, &mut writer).await?;
   
       dbg!("It cannot log this, it's stucked here, why?");
   
       Ok(())
   }

My updated code that works with rust-s3 thanks to you:

use axum::{
    body::{boxed},
    extract::{Multipart, Query},
    http::StatusCode,
    response::Response,
    routing::{on, MethodFilter},
    Router,
};
use axum_extra::body::AsyncReadBody;
use futures::TryStreamExt;
use s3::{creds::Credentials, Bucket};
use std::{io, net::SocketAddr, pin::Pin};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_util::io::StreamReader;

#[tokio::main]
async fn main() {
    let router = Router::new()
        .route("/upload", on(MethodFilter::POST, upload))
        .route("/download/*key", on(MethodFilter::GET, download));

    let addr = SocketAddr::from(([127, 0, 0, 1], 8080));

    axum::Server::bind(&addr)
        .serve(router.into_make_service())
        .await
        .unwrap();
}

pub async fn upload(mut multipart: Multipart) -> Result<Response, StatusCode> {
    while let Some(field) = multipart.next_field().await.unwrap() {
        let filename = if let Some(filename) = field.file_name() {
            filename.to_string()
        } else {
            continue;
        };

        let bucket = Bucket::new(
            "test",
            "us-east-1".parse().unwrap(),
            Credentials::default().unwrap(),
        )
        .unwrap();

        let body_with_io_error = field.map_err(|err| io::Error::new(io::ErrorKind::Other, err));

        let body_reader = StreamReader::new(body_with_io_error);

        futures::pin_mut!(body_reader);

        put_file(bucket, &filename, body_reader).await.unwrap();

        return Ok(Response::builder()
            .status(StatusCode::CREATED)
            .body(boxed("OK".to_string()))
            .unwrap());
    }

    Err(StatusCode::INTERNAL_SERVER_ERROR)
}

async fn put_file(
    bucket: Bucket,
    filename: &str,
    mut reader: Pin<&mut (dyn AsyncRead + Send)>,
) -> Result<(), ()> {
    bucket
        .put_object_stream(&mut reader, filename)
        .await
        .unwrap();

    Ok(())
}

pub async fn download(Query(params): Query<Vec<(String, String)>>) -> Result<Response, StatusCode> {
    let filename = params[0].1.to_string();

    let bucket = Bucket::new(
        "test",
        "us-east-1".parse().unwrap(),
        Credentials::default().unwrap(),
    )
    .unwrap();

    let (tx, rx) = tokio::io::duplex(65_536);

    futures::pin_mut!(tx);

    let body = AsyncReadBody::new(rx);

    get_file(bucket, &filename, tx).await.unwrap();

    let response = Response::builder().body(boxed(body)).unwrap();

    Ok(response)
}

async fn get_file(
    bucket: Bucket,
    filename: &str,
    mut writer: Pin<&mut (dyn AsyncWrite + Send)>,
) -> Result<(), ()> {
    bucket.get_object_stream(filename, &mut writer).await.unwrap();

    Ok(())
}

[frederikhors]

I spoke too soon. It's stucked on download with bucket adapter too.

[frederikhors]

It compiles but it's stucked... :(

[frederikhors]

I think I figured out what makes it stand still without proceeding.

If I increase the value in tokio::io::duplex(65_536); it works!

What value do you think I should put there?

Is that the size of chunks?

[davidpdrsn]

What value do you think I should put there?

Dunno 🤷

Is that the size of chunks?

From the docs:

The max_buf_size argument is the maximum amount of bytes that can be written to a side before the write returns Poll::Pending.

So its the max number of bytes that'll be held in memory waiting to be read.