rust 66 lines · 8 steps

Throttling file uploads in Axum with a Semaphore

An Axum handler streams multipart uploads to disk while a shared semaphore caps how many uploads run at once.

Explained by highlit
1use std::sync::Arc;
2use axum::{
3 extract::{Multipart, State},
4 http::StatusCode,
5 response::{IntoResponse, Response},
6 Json,
7};
8use serde_json::json;
9use tokio::sync::Semaphore;
10use tokio::io::AsyncWriteExt;
11 
12#[derive(Clone)]
13pub struct AppState {
14 pub upload_slots: Arc<Semaphore>,
15 pub storage_dir: Arc<std::path::PathBuf>,
16}
17 
18pub async fn upload(
19 State(state): State<AppState>,
20 mut multipart: Multipart,
21) -> Response {
22 let permit = match state.upload_slots.clone().try_acquire_owned() {
23 Ok(permit) => permit,
24 Err(_) => {
25 return (
26 StatusCode::SERVICE_UNAVAILABLE,
27 [("Retry-After", "5")],
28 Json(json!({ "error": "upload capacity reached, retry shortly" })),
29 )
30 .into_response();
31 }
32 };
33 
34 let mut written = 0u64;
35 while let Some(mut field) = match multipart.next_field().await {
36 Ok(field) => field,
37 Err(e) => return (StatusCode::BAD_REQUEST, e.to_string()).into_response(),
38 } {
39 let name = field.file_name().unwrap_or("blob").to_owned();
40 let path = state.storage_dir.join(sanitize(&name));
41 let mut file = match tokio::fs::File::create(&path).await {
42 Ok(f) => f,
43 Err(e) => return (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(),
44 };
45 
46 while let Some(chunk) = match field.chunk().await {
47 Ok(chunk) => chunk,
48 Err(e) => return (StatusCode::BAD_REQUEST, e.to_string()).into_response(),
49 } {
50 if let Err(e) = file.write_all(&chunk).await {
51 return (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response();
52 }
53 written += chunk.len() as u64;
54 }
55 file.flush().await.ok();
56 }
57 
58 drop(permit);
59 (StatusCode::CREATED, Json(json!({ "bytes": written }))).into_response()
60}
61 
62fn sanitize(name: &str) -> String {
63 name.chars()
64 .filter(|c| c.is_alphanumeric() || matches!(c, '.' | '-' | '_'))
65 .collect()
66}
01 / 01
STEP 01

Walkthrough

Space play step click any line
Three takeaways
  1. 1A shared semaphore turns unbounded concurrency into a fixed number of admission slots, protecting scarce resources like disk or memory.
  2. 2Streaming multipart fields chunk-by-chunk keeps memory flat regardless of upload size.
  3. 3Holding a permit for the whole handler and dropping it at the end ties resource lifetime to the request's actual work.

Related explainers

Share this explainer

Here's the card — post it anywhere.

Throttling file uploads in Axum with a Semaphore — share card
Made with highlit — turn any snippet into a walkthrough like this in about a minute.
Explain your code