From 0d8e2dd008fd6a6bbd2deb31d46139a5f9f9a29c Mon Sep 17 00:00:00 2001 From: Alexey Aristov Date: Thu, 28 Aug 2025 13:56:48 +0200 Subject: [PATCH] move upload code to blob.rs Signed-off-by: Alexey Aristov --- src/blob.rs | 143 ++++++++++++++++++++++++++++++++++++++++++++++++ src/handlers.rs | 142 ++--------------------------------------------- src/main.rs | 1 + 3 files changed, 149 insertions(+), 137 deletions(-) create mode 100644 src/blob.rs diff --git a/src/blob.rs b/src/blob.rs new file mode 100644 index 0000000000..573c8861ba --- /dev/null +++ b/src/blob.rs @@ -0,0 +1,143 @@ +use actix_web::web::Payload; +use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart}; +use blake3::Hasher; +use bytes::BytesMut; +use futures_util::StreamExt; +use tracing::*; + +use crate::s3::S3Client; +use crate::{ + config::CONFIG, + postgres::{self, Pool}, +}; + +use crate::handlers::{ApiError, HandlerResult}; + +pub struct Blob { + pub s3_key: String, + pub size: u64, +} + +// upload and deduplicate blob +#[instrument(level = "debug", skip_all, fields(s3_bucket, s3_key, upload))] +pub async fn upload(s3: &S3Client, pool: &Pool, mut payload: Payload) -> Result { + let span = Span::current(); + + let s3_bucket = &CONFIG.s3_bucket; + let s3_key = ksuid::Ksuid::generate().to_base62(); + + span.record("s3_bucket", &s3_bucket); + span.record("s3_key", &s3_key); + + let create_multipart = s3 + .create_multipart_upload() + .bucket(s3_bucket) + .key(&s3_key) + .send() + .await?; + + let upload_id = create_multipart.upload_id().unwrap(); + + span.record("upload", &upload_id[upload_id.len().saturating_sub(16)..]); + + debug!("upload start"); + + let upload_part = async |number, buffer: BytesMut| -> HandlerResult { + let upload = s3 + .upload_part() + .bucket(s3_bucket) + .key(&s3_key) + .upload_id(upload_id) + .body(buffer.freeze().into()) + .part_number(number) + .send() + .await?; + + let part = CompletedPart::builder() + .e_tag(upload.e_tag.unwrap()) + .part_number(number) + .build(); + + Ok(part) + }; + + let mut buffer = BytesMut::with_capacity(1024 * 1024 * 6); + let mut complete = CompletedMultipartUpload::builder(); + let mut part_number = 1; + let mut hash = Hasher::new(); + let mut total_in = 0; + let mut total_uploaded = 0; + + while let Some(part) = payload.next().await { + if let Ok(part) = part { + hash.update(&part); + + total_in += part.len(); + + buffer.extend_from_slice(&part); + + // each part must be at least 5MB + if buffer.len() > 1024 * 1024 * 5 { + trace!(length = buffer.len(), part_number, "upload part"); + + total_uploaded += buffer.len(); + + let uploaded = upload_part(part_number, buffer).await?; + + complete = complete.parts(uploaded); + + buffer = BytesMut::new(); + part_number += 1; + } + } else { + // TODO: cleanup incomplete upload + panic!("read error") + } + } + + // the last part + if buffer.len() > 0 { + total_uploaded += buffer.len(); + + trace!(length = buffer.len(), part_number, "upload part"); + let uploaded = upload_part(part_number, buffer).await?; + complete = complete.parts(uploaded); + } + + assert_eq!(total_in, total_uploaded); + + let _ = s3 + .complete_multipart_upload() + .bucket(s3_bucket) + .key(&s3_key) + .multipart_upload(complete.build()) + .upload_id(upload_id) + .send() + .await?; + + let hash = hash.finalize().to_hex().to_string(); + + debug!(hash, "upload complete"); + + let s3_key = if let Some(s3_key_found) = postgres::find_blob_by_hash(&pool, &hash).await? { + debug!(s3_key_found, "blob deduplicated"); + + // delete uploaded + s3.delete_object() + .bucket(s3_bucket) + .key(s3_key) + .send() + .await?; + + s3_key_found + } else { + debug!("blob created"); + postgres::insert_blob(&pool, &s3_key, &hash).await?; + s3_key + }; + + Ok(Blob { + s3_key, + size: total_uploaded as u64, + }) +} diff --git a/src/handlers.rs b/src/handlers.rs index 6f19590de7..cbb00af263 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -10,22 +10,19 @@ use actix_web::{ }, web::{Data, Header, Path, Payload}, }; -use blake3::Hasher; -use bytes::{Bytes, BytesMut}; -use futures_util::{Stream, StreamExt}; +use aws_sdk_s3::error::SdkError; +use bytes::Bytes; +use futures_util::Stream; use serde::{Deserialize, Serialize}; use tracing::*; use uuid::Uuid; +use crate::blob::upload; use crate::s3::S3Client; use crate::{ config::CONFIG, postgres::{self, Pool}, }; -use aws_sdk_s3::{ - error::SdkError, - types::{CompletedMultipartUpload, CompletedPart}, -}; #[derive(Deserialize, Debug)] pub struct ObjectPath { @@ -48,7 +45,7 @@ pub enum ApiError { Other(#[from] anyhow::Error), } -type HandlerResult = Result; +pub type HandlerResult = Result; impl actix_web::error::ResponseError for ApiError { fn error_response(&self) -> HttpResponse { @@ -73,11 +70,6 @@ impl From>, } -// upload and deduplicate blob -#[instrument(level = "debug", skip_all, fields(s3_bucket, s3_key, upload))] -async fn upload(s3: &S3Client, pool: &Pool, mut payload: Payload) -> Result { - let span = Span::current(); - - let s3_bucket = &CONFIG.s3_bucket; - let s3_key = ksuid::Ksuid::generate().to_base62(); - - span.record("s3_bucket", &s3_bucket); - span.record("s3_key", &s3_key); - - let create_multipart = s3 - .create_multipart_upload() - .bucket(s3_bucket) - .key(&s3_key) - .send() - .await?; - - let upload_id = create_multipart.upload_id().unwrap(); - - span.record("upload", &upload_id[upload_id.len().saturating_sub(16)..]); - - debug!("upload start"); - - let upload_part = async |number, buffer: BytesMut| -> HandlerResult { - let upload = s3 - .upload_part() - .bucket(s3_bucket) - .key(&s3_key) - .upload_id(upload_id) - .body(buffer.freeze().into()) - .part_number(number) - .send() - .await?; - - let part = CompletedPart::builder() - .e_tag(upload.e_tag.unwrap()) - .part_number(number) - .build(); - - Ok(part) - }; - - let mut buffer = BytesMut::with_capacity(1024 * 1024 * 6); - let mut complete = CompletedMultipartUpload::builder(); - let mut part_number = 1; - let mut hash = Hasher::new(); - let mut total_in = 0; - let mut total_uploaded = 0; - - while let Some(part) = payload.next().await { - if let Ok(part) = part { - hash.update(&part); - - total_in += part.len(); - - buffer.extend_from_slice(&part); - - // each part must be at least 5MB - if buffer.len() > 1024 * 1024 * 5 { - trace!(length = buffer.len(), part_number, "upload part"); - - total_uploaded += buffer.len(); - - let uploaded = upload_part(part_number, buffer).await?; - - complete = complete.parts(uploaded); - - buffer = BytesMut::new(); - part_number += 1; - } - } else { - // TODO: cleanup incomplete upload - panic!("read error") - } - } - - // the last part - if buffer.len() > 0 { - total_uploaded += buffer.len(); - - trace!(length = buffer.len(), part_number, "upload part"); - let uploaded = upload_part(part_number, buffer).await?; - complete = complete.parts(uploaded); - } - - assert_eq!(total_in, total_uploaded); - - let _ = s3 - .complete_multipart_upload() - .bucket(s3_bucket) - .key(&s3_key) - .multipart_upload(complete.build()) - .upload_id(upload_id) - .send() - .await?; - - let hash = hash.finalize().to_hex().to_string(); - - debug!(hash, "upload complete"); - - let s3_key = if let Some(s3_key_found) = postgres::find_blob_by_hash(&pool, &hash).await? { - debug!(s3_key_found, "blob deduplicated"); - - // delete uploaded - s3.delete_object() - .bucket(s3_bucket) - .key(s3_key) - .send() - .await?; - - s3_key_found - } else { - debug!("blob created"); - postgres::insert_blob(&pool, &s3_key, &hash).await?; - s3_key - }; - - Ok(Blob { - s3_key, - size: total_uploaded as u64, - }) -} - #[instrument(level = "debug", skip_all, fields(workspace, huly_key))] pub async fn put(request: HttpRequest, payload: Payload) -> HandlerResult { let span = Span::current(); diff --git a/src/main.rs b/src/main.rs index a1e3b7b5ad..362625cdeb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,6 +14,7 @@ use uuid::Uuid; use hulyrs::services::jwt::actix::ServiceRequestExt; +mod blob; mod config; mod handlers; mod postgres;