feat: Implement Async/Streaming for uploads

nuclearcat · nuclearcat · commit ff6777277261 · 2025-10-23T10:36:39.000+03:00
We had very ugly upload, with buffering whole file in memory,
then writing to temp file.
Now we implement proper async multipart stream to Azure, with
max buffering 10MB (might be configurable in near future).
More details in CHANGELOG.md

Signed-off-by: Denys Fedoryshchenko &lt;denys.f@collabora.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,72 @@
 # kernelci-storage ChangeLog
 
-## 0.1.0 (2025-01-06)
+All notable changes to this project will be documented in this file.
 
-- Initial release
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [0.2.0] - 2025-10-23
+
+### Added
+- Streaming multipart upload support for large files (200MB+)
+- New `write_file_streaming()` async method in Driver trait
+- `FieldStream` wrapper to convert multipart fields to AsyncRead
+- Streaming implementations for both Azure and Local storage backends
+- `async-trait` dependency (0.1) for trait async methods
+- `bytes` dependency (1.9) for efficient byte buffer handling
+
+### Changed
+- Refactored `ax_post_file` handler to stream uploads directly without loading entire file into memory
+- Azure backend now streams uploads in 10MB chunks directly to blob storage
+- Local backend now streams uploads in 10MB chunks directly to filesystem
+- Upload handler processes multipart fields sequentially and starts streaming immediately
+
+### Performance Improvements
+- Reduced memory usage: Only 10MB buffer in memory at any time instead of entire file
+- Improved upload performance for large files by eliminating full memory buffering
+- Removed temporary file creation during Azure uploads (streaming directly from multipart)
+- Better scalability: Can handle files larger than available RAM
+
+### Technical Details
+- No breaking changes to existing API - fully backward compatible
+- Existing upload clients continue to work without modifications
+- All existing functionality preserved (permissions, content-type, owner tags, file locking)
+- Proper lifetime management for multipart field streaming
+- Uses async/await with tokio AsyncRead trait
+
+## [0.1.0] - 2025-01-06
+
+### Added
+- Initial release of KernelCI Storage Server
+- JWT token-based authentication with HMAC-SHA256
+- File upload/download with local caching
+- Azure Blob Storage backend support with chunked uploads (10MB blocks)
+- Local filesystem backend support
+- Range request support for partial content downloads
+- Prometheus metrics endpoint (`/metrics`)
+- File locking mechanism to prevent concurrent uploads to same path
+- Automatic cache cleanup when disk space < 12%
+- LRU-style cache deletion (removes files older than 60 minutes)
+- User-specific upload path restrictions via configuration
+- SSL/TLS support with certificate configuration
+- Docker support with Dockerfile
+
+### Features
+- RESTful API with endpoints:
+  - `GET /` - Server status
+  - `POST /upload` or `POST /v1/file` - File upload (requires JWT)
+  - `GET /*filepath` - File download (public)
+  - `GET /v1/checkauth` - Validate JWT token
+  - `GET /v1/list` - List all files (public)
+  - `GET /metrics` - Prometheus metrics
+- Header preservation for HTTP caching (ETag, Last-Modified)
+- Content-type detection with heuristics
+- Configurable via TOML configuration file
+- Environment variable override for config path (`KCI_STORAGE_CONFIG`)
+- Verbose logging support via `--verbose` flag
+- Command-line utilities:
+  - `--generate-jwt-secret` - Generate JWT secret
+  - `--generate-jwt-token` - Generate JWT token for user
+- SHA-512 based cache filenames to avoid path conflicts
+- Client IP detection with X-Forwarded-For support
+- Disk space monitoring with hysteresis
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,14 +1,16 @@
 [package]
 name = "kernelci-storage"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2021"
 
 [dependencies]
+async-trait = "0.1"
 axum = { version = "0.7.9", features = ["tracing", "multipart", "macros"] }
 axum-server = { version = "0.7.1", features = ["rustls", "rustls-pemfile", "tls-rustls"] }
 azure_blob_uploader = "0.1.4"
 azure_storage = "0.21.0"
 azure_storage_blobs = "0.21.0"
+bytes = "1.9"
 chksum-hash-sha2-512 = "0.0.1"
 chrono = "0.4.39"
 clap = { version = "4.5.23", features = ["derive"] }
diff --git a/src/azure.rs b/src/azure.rs
@@ -11,6 +11,7 @@ impl AzureDriver {
 }
 
 use crate::{debug_log, get_config_content, ReceivedFile};
+use async_trait::async_trait;
 use axum::http::{HeaderName, HeaderValue};
 use azure_storage::StorageCredentials;
 use azure_storage_blobs::container::operations::BlobItem;
@@ -26,6 +27,7 @@ use std::io::Read;
 use std::io::Write;
 use std::sync::Arc;
 use tempfile::Builder;
+use tokio::io::AsyncReadExt;
 use toml::Table;
 
 #[derive(Deserialize)]
@@ -75,7 +77,110 @@ fn calculate_checksum(filename: &String, data: &[u8]) {
     debug_log!("File: {} Checksum: {}", filename, digest.to_hex_lowercase());
 }
 
-/// Write file to Azure blob storage
+/// Write file to Azure blob storage using streaming (new version)
+async fn write_file_to_blob_streaming(
+    filename: String,
+    data: &mut (dyn tokio::io::AsyncRead + Unpin + Send),
+    cont_type: String,
+    owner_email: Option<String>,
+) -> &'static str {
+    let azure_cfg = Arc::new(get_azure_credentials("azure"));
+
+    let storage_account = azure_cfg.account.as_str();
+    let storage_key = azure_cfg.key.clone();
+    let storage_container = azure_cfg.container.as_str();
+    let storage_blob = filename.as_str();
+    let storage_credential = StorageCredentials::access_key(storage_account, storage_key);
+    let blob_client = ClientBuilder::new(storage_account, storage_credential)
+        .blob_client(storage_container, storage_blob);
+
+    let mut total_bytes_uploaded: usize = 0;
+    let chunk_size = 10 * 1024 * 1024; // 10MB chunks
+    let mut blocks = BlockList::default();
+
+    loop {
+        let mut buffer = vec![0u8; chunk_size];
+        let bytes_read = data.read(&mut buffer).await;
+        match bytes_read {
+            Ok(bytes_read) => {
+                if bytes_read == 0 {
+                    break;
+                }
+                buffer.truncate(bytes_read);
+                let block_id = BlockId::new(hex::encode(total_bytes_uploaded.to_le_bytes()));
+                blocks
+                    .blocks
+                    .push(BlobBlockType::Uncommitted(block_id.clone()));
+                match blob_client.put_block(block_id, buffer).await {
+                    Ok(_) => {
+                        total_bytes_uploaded += bytes_read;
+                        debug_log!("Uploaded {} bytes", total_bytes_uploaded);
+                    }
+                    Err(e) => {
+                        eprintln!("Error uploading block: {:?}", e);
+                        break;
+                    }
+                }
+            }
+            Err(e) => {
+                eprintln!("Error reading stream: {:?}", e);
+                break;
+            }
+        }
+    }
+    match blob_client
+        .put_block_list(blocks)
+        .content_type(cont_type)
+        .await
+    {
+        Ok(_) => {
+            debug_log!("Block list uploaded");
+            let blob_url_res = blob_client.url();
+            match blob_url_res {
+                Ok(blob_url) => {
+                    debug_log!("Blob URL: {}", blob_url);
+                }
+                Err(e) => {
+                    eprintln!("Error getting blob URL: {:?}", e);
+                }
+            }
+
+            // Set owner tag if email is provided
+            if let Some(email) = owner_email {
+                let mut tags = Tags::new();
+                let sanitized = sanitize_tag_component(&email);
+                if sanitized != email {
+                    debug_log!(
+                        "Sanitized owner tag value from '{}' to '{}'",
+                        email,
+                        sanitized
+                    );
+                }
+                // Ensure non-empty value
+                let final_value = if sanitized.is_empty() {
+                    "_".to_string()
+                } else {
+                    sanitized
+                };
+                tags.insert("owner".to_string(), final_value);
+                match blob_client.set_tags(tags).await {
+                    Ok(_) => {
+                        debug_log!("Owner tag set successfully");
+                    }
+                    Err(e) => {
+                        eprintln!("Error setting owner tag: {:?}", e);
+                    }
+                }
+            }
+        }
+        Err(e) => {
+            eprintln!("Error uploading block list: {:?}", e);
+        }
+    }
+    "OK"
+}
+
+/// Write file to Azure blob storage (legacy version using Vec<u8>)
 /// TBD: Rework, do not keep whole file as Vec<u8> in memory!!!
 async fn write_file_to_blob(
     filename: String,
@@ -429,6 +534,7 @@ async fn azure_list_files(directory: String) -> Vec<String> {
 }
 
 /// Implement Driver trait for AzureDriver
+#[async_trait]
 impl super::Driver for AzureDriver {
     fn write_file(
         &self,
@@ -445,6 +551,16 @@ impl super::Driver for AzureDriver {
         });
         filenameret
     }
+    async fn write_file_streaming(
+        &self,
+        filename: String,
+        data: &mut (dyn tokio::io::AsyncRead + Unpin + Send),
+        cont_type: String,
+        owner_email: Option<String>,
+    ) -> String {
+        write_file_to_blob_streaming(filename.clone(), data, cont_type, owner_email).await;
+        filename
+    }
     fn tag_file(
         &self,
         filename: String,
diff --git a/src/local.rs b/src/local.rs
@@ -3,13 +3,15 @@
 // Author: Denys Fedoryshchenko <denys.f@collabora.com>
 
 use crate::{debug_log, get_config_content, ReceivedFile};
+use async_trait::async_trait;
 use axum::http::{HeaderName, HeaderValue};
 use chksum_hash_sha2_512 as sha2_512;
 use headers::HeaderMap;
 use serde::Deserialize;
 use std::fs::{self, File, OpenOptions};
 use std::io::Write;
 use std::path::{Path, PathBuf};
+use tokio::io::AsyncReadExt;
 use toml::Table;
 
 pub struct LocalDriver;
@@ -91,7 +93,57 @@ fn calculate_checksum(filename: &str, data: &[u8]) {
     debug_log!("File: {} Checksum: {}", filename, digest.to_hex_lowercase());
 }
 
-/// Write file to local storage
+/// Write file to local storage using streaming (new version)
+async fn write_file_to_local_streaming(
+    filename: String,
+    data: &mut (dyn tokio::io::AsyncRead + Unpin + Send),
+    cont_type: String,
+    owner_email: Option<String>,
+) -> Result<String, String> {
+    let file_path = get_storage_file_path(&filename);
+
+    // Ensure directory structure exists
+    if let Err(e) = ensure_directory_exists(&file_path) {
+        return Err(format!("Failed to create directory structure: {}", e));
+    }
+
+    // Write the file with streaming
+    let mut file = tokio::fs::File::create(&file_path)
+        .await
+        .map_err(|e| format!("Failed to create file: {}", e))?;
+
+    let mut buffer = vec![0u8; 10 * 1024 * 1024]; // 10MB buffer
+    let mut total_bytes = 0;
+
+    loop {
+        match data.read(&mut buffer).await {
+            Ok(0) => break, // EOF
+            Ok(n) => {
+                tokio::io::AsyncWriteExt::write_all(&mut file, &buffer[..n])
+                    .await
+                    .map_err(|e| format!("Failed to write file: {}", e))?;
+                total_bytes += n;
+                debug_log!("Written {} bytes to local storage", total_bytes);
+            }
+            Err(e) => return Err(format!("Failed to read stream: {}", e)),
+        }
+    }
+
+    // Create and write metadata (headers and owner tag)
+    let metadata_path = get_metadata_file_path(&filename);
+    if let Ok(mut metadata_file) = File::create(&metadata_path) {
+        let mut metadata_content = format!("content-type:{}\n", cont_type);
+        if let Some(email) = owner_email {
+            metadata_content.push_str(&format!("tag-owner:{}\n", email));
+        }
+        let _ = metadata_file.write_all(metadata_content.as_bytes());
+    }
+
+    debug_log!("File written to local storage: {}", file_path.display());
+    Ok(filename)
+}
+
+/// Write file to local storage (legacy version using Vec<u8>)
 fn write_file_to_local(
     filename: String,
     data: Vec<u8>,
@@ -272,6 +324,7 @@ fn set_tags_for_local_file(
 }
 
 /// Implement Driver trait for LocalDriver
+#[async_trait]
 impl super::Driver for LocalDriver {
     fn write_file(
         &self,
@@ -289,6 +342,22 @@ impl super::Driver for LocalDriver {
         }
     }
 
+    async fn write_file_streaming(
+        &self,
+        filename: String,
+        data: &mut (dyn tokio::io::AsyncRead + Unpin + Send),
+        cont_type: String,
+        owner_email: Option<String>,
+    ) -> String {
+        match write_file_to_local_streaming(filename.clone(), data, cont_type, owner_email).await {
+            Ok(_) => filename,
+            Err(e) => {
+                eprintln!("Local storage streaming write error: {}", e);
+                String::new()
+            }
+        }
+    }
+
     fn get_file(&self, filename: String) -> ReceivedFile {
         get_file_from_local(filename)
     }
diff --git a/src/main.rs b/src/main.rs
