From be1046877b6ffaae1a04099f8d2a3c0115de888d Mon Sep 17 00:00:00 2001 From: Alexey Aristov Date: Mon, 1 Sep 2025 23:37:22 +0200 Subject: [PATCH] merge strategy (wip) Signed-off-by: Alexey Aristov --- Cargo.lock | 1 + server/Cargo.toml | 1 + server/src/blob.rs | 83 +++++++++-------- server/src/config.rs | 3 + server/src/handlers.rs | 204 +++++++++++++++++++++-------------------- server/src/main.rs | 1 + server/src/merge.rs | 57 ++++++++++++ server/src/postgres.rs | 9 +- server/src/s3.rs | 13 ++- tests/src/put.rs | 163 +++++++++++++++++++++++++++++++- tests/src/util.rs | 10 ++ 11 files changed, 398 insertions(+), 147 deletions(-) create mode 100644 server/src/merge.rs diff --git a/Cargo.lock b/Cargo.lock index 721b52acc7..0a466a4931 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2315,6 +2315,7 @@ dependencies = [ "serde", "serde_json", "size", + "strum 0.27.2", "thiserror 2.0.16", "tokio", "tokio-postgres", diff --git a/server/Cargo.toml b/server/Cargo.toml index bfe22cf7cb..9bac4c7443 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -46,3 +46,4 @@ mime = "0.3.17" async-stream = "0.3.6" blake3 = "1.8.2" futures = "0.3.31" +strum = { version = "0.27.2", features = ["derive"] } diff --git a/server/src/blob.rs b/server/src/blob.rs index ff4264bf66..f79a080c97 100644 --- a/server/src/blob.rs +++ b/server/src/blob.rs @@ -1,6 +1,6 @@ use std::error::Error as StdError; -use aws_sdk_s3::primitives::ByteStream; +use aws_sdk_s3::primitives::{ByteStream, SdkBody}; use blake3::Hasher; use bytes::{Bytes, BytesMut}; use futures::stream::StreamExt; @@ -15,10 +15,13 @@ use crate::{ postgres::{self, Pool}, }; +#[derive(Debug)] pub struct Blob { pub s3_key: String, pub length: usize, - pub inline: Option>, + pub inline: Option, + pub parts_count: Option, + pub deduplicated: bool, } fn random_key() -> String { @@ -59,40 +62,41 @@ where return Err(actix_web::error::ErrorBadRequest("payload size mismatch").into()); } + let buffer = buffer.freeze(); + let hash = hash.update(&buffer).finalize().to_hex(); let length = buffer.len(); - let inline = if length < CONFIG.inline_threshold.bytes() as usize { - Some(buffer.to_vec()) - } else { - None - }; + let inline = Some(buffer.clone()); - let s3_key = if let Some(s3_key_found) = postgres::find_blob_by_hash(&pool, &hash).await? { - span.record("s3_key", &s3_key_found); - debug!(s3_key_found, "blob deduplicated"); - s3_key_found - } else { - let s3_key = random_key(); - span.record("s3_key", &s3_key); + let (s3_key, deduplicated) = + if let Some(s3_key_found) = postgres::find_blob_by_hash(&pool, &hash).await? { + span.record("s3_key", &s3_key_found); + debug!(s3_key_found, "blob deduplicated"); + (s3_key_found, true) + } else { + let s3_key = random_key(); + span.record("s3_key", &s3_key); - s3.put_object() - .bucket(s3_bucket) - .key(&s3_key) - .body(ByteStream::from(buffer.freeze())) - .send() - .await?; + s3.put_object() + .bucket(s3_bucket) + .key(&s3_key) + .body(ByteStream::from(buffer)) + .send() + .await?; - postgres::insert_blob(&pool, &s3_key, &hash).await?; + postgres::insert_blob(&pool, &s3_key, &hash).await?; - debug!("blob created"); - s3_key - }; + debug!("blob created"); + (s3_key, false) + }; Blob { s3_key, length, inline, + parts_count: None, + deduplicated, } } else { let s3_key = random_key(); @@ -102,27 +106,30 @@ where let hash = upload.hash.to_hex().to_string(); - let s3_key = if let Some(s3_key_found) = postgres::find_blob_by_hash(&pool, &hash).await? { - debug!(s3_key_found, "blob deduplicated"); + let (s3_key, deduplicated) = + 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?; + // 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 - }; + (s3_key_found, true) + } else { + debug!("blob created"); + postgres::insert_blob(&pool, &s3_key, &hash).await?; + (s3_key, false) + }; Blob { s3_key, length: upload.length, inline: None, + parts_count: Some(upload.parts_count), + deduplicated, } }; diff --git a/server/src/config.rs b/server/src/config.rs index 1d954208cb..67739638b6 100644 --- a/server/src/config.rs +++ b/server/src/config.rs @@ -32,7 +32,10 @@ pub struct Config { pub s3_bucket: String, + // use multipart upload if blob size is greater than this pub multipart_threshold: Size, + + // store blobs inline if size is less than this pub inline_threshold: Size, } diff --git a/server/src/handlers.rs b/server/src/handlers.rs index 64f0c3ebab..0cf7a2832c 100644 --- a/server/src/handlers.rs +++ b/server/src/handlers.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, fmt::Display, sync::Arc}; +use std::{collections::HashMap, fmt::Display, str::FromStr, sync::Arc}; use actix_web::{ HttpRequest, HttpResponse, @@ -14,13 +14,13 @@ use aws_sdk_s3::error::SdkError; use bytes::Bytes; use futures_util::Stream; use serde::{Deserialize, Serialize}; -use serde_json::{Value, from_slice}; use size::Size; use tracing::*; use uuid::Uuid; -use crate::blob; +use crate::merge::MergeStrategy; use crate::s3::S3Client; +use crate::{blob, merge}; use crate::{ config::CONFIG, postgres::{self, Pool}, @@ -72,41 +72,80 @@ impl From HandlerResult; - async fn content_type(&mut self) -> Option; - async fn merge_strategy(&mut self) -> Option; +fn random_etag() -> String { + ksuid::Ksuid::generate().to_base62() } -impl ServiceRequestExt for ServiceRequest { - async fn content_length(&mut self) -> HandlerResult { - self.extract::>() - .await - .map(|header| Size::from_bytes(*header.0)) - .map_err(|_| actix_web::error::ErrorBadRequest("invalid content length").into()) - } - - async fn content_type(&mut self) -> Option { - self.extract::>() - .await - .map(|header| header.0.to_string()) - .ok() - } - - async fn merge_strategy(&mut self) -> Option { - self.headers() - .get("Huly-Merge-Strategy") - .and_then(|v| v.to_str().ok()) - .and_then(|v| serde_json::from_str::(v).ok()) - } +#[derive(Debug, Clone)] +pub struct Headers { + pub content_length: Size, + pub content_type: Option, + pub merge_strategy: MergeStrategy, + pub huly_headers: Vec<(String, String)>, + pub meta: Vec<(String, String)>, } -#[derive(Clone, Copy, Debug, Serialize, Deserialize, Default)] -#[serde(rename_all = "lowercase")] -enum MergeStrategy { - JsonPatch, - #[default] - Concatenate, +async fn extract_headers(request: &mut ServiceRequest) -> HandlerResult { + let content_length = request + .extract::>() + .await + .map(|header| Size::from_bytes(*header.0)) + .map_err(|_| actix_web::error::ErrorBadRequest("invalid content length"))?; + + let content_type = request + .extract::>() + .await + .map(|header| header.0.to_string()) + .ok(); + + let merge_strategy = request + .headers() + .get("Huly-Merge-Strategy") + .and_then(|v| v.to_str().ok()) + .map(|v| { + MergeStrategy::from_str(v).map_err(|_| { + actix_web::error::ErrorBadRequest(format!("invalid merge strategy: {v}")) + }) + }) + .transpose()? + .unwrap_or(MergeStrategy::Concatenate); + + let mut huly_headers = Vec::new(); + for (key, value) in request.headers().iter() { + if let Some(header) = key.as_str().strip_prefix("huly-header-") { + if let Ok(value) = value.to_str() { + huly_headers.push((header.to_owned(), value.to_owned())); + } + } + } + if let Some(content_type) = &content_type { + huly_headers.push(( + http::header::CONTENT_TYPE.as_str().to_owned(), + content_type.to_owned(), + )); + } + + let mut meta = Vec::new(); + for (key, value) in request.headers().iter() { + if let Some(header) = key.as_str().strip_prefix("huly-meta-") { + if let Ok(value) = value.to_str() { + meta.push((header.to_owned(), value.to_owned())); + } + } + } + + meta.push(( + "merge-strategy".to_owned(), + serde_json::to_string(&merge_strategy).unwrap(), + )); + + Ok(Headers { + content_length, + content_type, + merge_strategy, + huly_headers, + meta, + }) } #[derive(Serialize, Deserialize, Debug)] @@ -138,23 +177,9 @@ pub async fn put(request: HttpRequest, payload: Payload) -> HandlerResult>().await?.into_inner(); - let content_length = request.content_length().await?; - let content_type = request.content_type().await; - let merge_strategy = request.merge_strategy().await.unwrap_or_default(); + let headers = extract_headers(&mut request).await?; - match (merge_strategy, &content_type) { - (MergeStrategy::JsonPatch, Some(ct)) - if ct != "application/json" || content_length > CONFIG.inline_threshold => - { - return Err( - actix_web::error::ErrorBadRequest("invalid content type and length").into(), - ); - } - - _ => { - // - } - } + merge::validate_put_request(&headers)?; span.record("workspace", path.workspace.to_string()); span.record("huly_key", &path.key); @@ -162,40 +187,9 @@ pub async fn put(request: HttpRequest, payload: Payload) -> HandlerResult>().unwrap().to_owned(); let s3 = request.app_data::>().unwrap().to_owned(); - let mut headers = Vec::new(); - for (key, value) in request.headers().iter() { - if let Some(header) = key.as_str().strip_prefix("huly-header-") { - if let Ok(value) = value.to_str() { - headers.push((header.to_owned(), value.to_owned())); - } - } - } + let uploaded = blob::upload(&s3, &pool, headers.content_length, payload).await?; - if let Some(content_type) = content_type { - headers.push((http::header::CONTENT_TYPE.as_str().to_owned(), content_type)); - } - - let mut meta = Vec::new(); - for (key, value) in request.headers().iter() { - if let Some(header) = key.as_str().strip_prefix("huly-meta-") { - if let Ok(value) = value.to_str() { - meta.push((header.to_owned(), value.to_owned())); - } - } - } - - let uploaded = blob::upload(&s3, &pool, content_length, payload).await?; - - match merge_strategy { - MergeStrategy::JsonPatch => { - from_slice::(uploaded.inline.as_ref().unwrap()) - .map_err(|x| actix_web::error::ErrorBadRequest(x.to_string()))?; - } - - _ => { - // - } - } + merge::validate_put_body(&headers, &uploaded)?; let part_data = PartData { workspace: path.workspace, @@ -203,27 +197,37 @@ pub async fn put(request: HttpRequest, payload: Payload) -> HandlerResult HandlerResult>().unwrap().to_owned(); let s3 = request.app_data::>().unwrap().to_owned(); - let content_length = request.content_length().await?; + let headers = extract_headers(&mut request).await?; - let uploaded = blob::upload(&s3, &pool, content_length, payload).await?; + let uploaded = blob::upload(&s3, &pool, headers.content_length, payload).await?; let parts = postgres::find_parts::(&pool, path.workspace, &path.key).await?; @@ -261,7 +265,7 @@ pub async fn patch(request: HttpRequest, payload: Payload) -> HandlerResult HandlerResult<()> { + dbg!(&headers); + + match headers.merge_strategy { + MergeStrategy::JsonPatch => { + if headers.content_type != Some("application/json".to_string()) + || headers.content_length > CONFIG.inline_threshold + { + return Err( + actix_web::error::ErrorBadRequest("invalid content type and length").into(), + ); + } + } + + _ => { + // + } + } + + Ok(()) +} + +pub fn validate_put_body(headers: &Headers, blob: &Blob) -> HandlerResult<()> { + dbg!(&headers); + dbg!(&blob); + + match headers.merge_strategy { + MergeStrategy::JsonPatch => { + from_slice::(blob.inline.as_ref().unwrap()) + .map_err(|x| actix_web::error::ErrorBadRequest(x.to_string()))?; + } + + _ => { + // + } + } + + Ok(()) +} diff --git a/server/src/postgres.rs b/server/src/postgres.rs index 7c97d311fd..bdae14f528 100644 --- a/server/src/postgres.rs +++ b/server/src/postgres.rs @@ -1,6 +1,7 @@ use std::pin::Pin; use bb8_postgres::PostgresConnectionManager; +use bytes::Bytes; use serde::de::DeserializeOwned; use tokio_postgres::NoTls; use tokio_postgres::{self as pg}; @@ -134,12 +135,13 @@ pub async fn append_part( workspace: uuid::Uuid, key: &str, part: u32, - inline: Option>, + inline: Option, data: &D, ) -> anyhow::Result<()> { let connection = pool.get().await?; let data = serde_json::to_value(data)?; + let inline = inline.map(|b| b.to_vec()); connection .execute( @@ -155,7 +157,7 @@ pub async fn set_part( pool: &Pool, workspace: uuid::Uuid, key: &str, - inline: Option>, + inline: Option, data: &D, ) -> anyhow::Result<()> { let mut connection = pool.get().await?; @@ -170,13 +172,14 @@ pub async fn set_part( .await?; let data = serde_json::to_value(data)?; + let inline = inline.map(|b| b.to_vec()); transaction .execute( r#" insert into object (workspace, key, part, inline, data) values ($1, $2, 0, $3, $4) on conflict (workspace, key, part) do update set - inline = $3, + inline = $3, data = $4 "#, &[&workspace, &key, &inline, &data], diff --git a/server/src/s3.rs b/server/src/s3.rs index 0ca106f43c..8d19a01120 100644 --- a/server/src/s3.rs +++ b/server/src/s3.rs @@ -32,6 +32,7 @@ pub async fn client() -> S3Client { pub struct Upload { pub hash: Hash, pub length: usize, + pub parts_count: usize, } async fn multipart_upload_stream( @@ -110,7 +111,17 @@ where let hash = hash.finalize(); - Ok((complete.build(), Upload { hash, length })) + let complete = complete.build(); + let parts_count = complete.parts().len(); + + Ok(( + complete, + Upload { + hash, + length, + parts_count, + }, + )) } #[tracing::instrument(level = "debug", skip_all)] diff --git a/tests/src/put.rs b/tests/src/put.rs index a4900e31f8..a002c4d1f5 100644 --- a/tests/src/put.rs +++ b/tests/src/put.rs @@ -1,3 +1,4 @@ +use hulyrs::StatusCode; use tanu::{ check, check_eq, eyre, http::{self, Client}, @@ -6,9 +7,39 @@ use tanu::{ use crate::util::*; #[tanu::test] -pub async fn put_new() -> eyre::Result<()> { +pub async fn put_new_deduplicated_no_multipart() -> eyre::Result<()> { + let text = random_text(1024); + + let http = Client::new(); + + let res = http + .key_put(&random_key()) + .body(text.clone()) + .send() + .await?; + check!(res.status().is_success()); + check_eq!(None, res.header("huly-deduplicated")); + + let res = http + .key_put(&random_key()) + .body(text.clone()) + .send() + .await?; + check!(res.status().is_success()); + check_eq!(Some("true"), res.header("huly-deduplicated")); + + Ok(()) +} + +#[tanu::test(1)] +#[tanu::test(2)] +#[tanu::test(10)] +#[tanu::test(100)] +#[tanu::test(1024)] +#[tanu::test(1024 * 10)] +pub async fn put_new_sized_no_multipart(size: usize) -> eyre::Result<()> { let key = random_key(); - let text = random_text(1024 * 10); + let text = random_text(size); let http = Client::new(); @@ -23,6 +54,7 @@ pub async fn put_new() -> eyre::Result<()> { check!(res.status().is_success()); check!(res.headers().get(http::header::ETAG).is_some()); check!(res.headers().get(http::header::CONTENT_LOCATION).is_some()); + check_eq!(None, res.header("huly-parts-count")); // check content let res = http.key_get(&key).send().await?; @@ -31,6 +63,49 @@ pub async fn put_new() -> eyre::Result<()> { "text/plain", res.headers().get(http::header::CONTENT_TYPE).unwrap() ); + check_eq!( + res.header("content-length"), + Some(size.to_string().as_str()) + ); + check_eq!(text, res.text().await?); + + Ok(()) +} + +#[tanu::test] +pub async fn put_new_sized_multipart() -> eyre::Result<()> { + let text = random_text(1024 * 1024 * 5); // above multipart threshold + let http = Client::new(); + + let key1 = random_key(); + + let res = http.key_put(&key1).body(text.clone()).send().await?; + check!(res.status().is_success()); + + check!(res.status().is_success()); + check_eq!(Some("1"), res.header("huly-s3-parts-count")); + + let res = http.key_get(&key1).send().await?; + check!(res.status().is_success()); + check_eq!( + res.header("content-length"), + Some(text.len().to_string().as_str()) + ); + check_eq!(text, res.text().await?); + + let key2 = random_key(); + + let res = http.key_put(&key2).body(text.clone()).send().await?; + check!(res.status().is_success()); + check_eq!(None, res.header("huly-s3-parts-count")); + check_eq!(Some("true"), res.header("huly-deduplicated")); + + let res = http.key_get(&key2).send().await?; + check!(res.status().is_success()); + check_eq!( + res.header("content-length"), + Some(text.len().to_string().as_str()) + ); check_eq!(text, res.text().await?); Ok(()) @@ -134,7 +209,7 @@ pub async fn put_with_header_case( // check header is returned let res = http.key_get(&key).send().await?; check!(res.status().is_success()); - check_eq!(value, res.headers().get(res_header).unwrap()); + check_eq!(Some(value), res.header(res_header)); Ok(()) } @@ -152,6 +227,8 @@ pub async fn put_with_headers() -> eyre::Result<()> { .body(body) .header("huly-header-header1", "foo") .header("huly-header-header2", "bar") + .header("huly-meta-meta1", "baz") + .header("content-type", "application/json") .send() .await?; check!(res.status().is_success()); @@ -159,8 +236,10 @@ pub async fn put_with_headers() -> eyre::Result<()> { // check headers are returned let res = http.key_get(&key).send().await?; check!(res.status().is_success()); - check_eq!("foo", res.headers().get("header1").unwrap()); - check_eq!("bar", res.headers().get("header2").unwrap()); + check_eq!(Some("foo"), res.header("header1")); + check_eq!(Some("bar"), res.header("header2")); + check_eq!(Some("application/json"), res.header("content-type")); + check_eq!(None, res.header("meta1")); Ok(()) } @@ -190,3 +269,77 @@ pub async fn put_with_meta() -> eyre::Result<()> { Ok(()) } + +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum Body { + Random(usize), + Text(&'static str), +} + +#[tanu::test(1, None, None, Body::Text("{}"), StatusCode::CREATED)] +#[tanu::test(2, Some("invalid"), None, Body::Text("{}"), StatusCode::BAD_REQUEST)] +#[tanu::test(3, Some("jsonpatch"), None, Body::Text("{}"), StatusCode::BAD_REQUEST)] +#[tanu::test( + 4, + Some("jsonpatch"), + Some("application/xml"), + Body::Text("{}"), + StatusCode::BAD_REQUEST +)] +#[tanu::test( + 5, + Some("jsonpatch"), + Some("application/json"), + Body::Text("{"), + StatusCode::BAD_REQUEST +)] +#[tanu::test( + 6, + Some("jsonpatch"), + Some("application/json"), + Body::Random(1024 * 101), + StatusCode::BAD_REQUEST +)] +#[tanu::test( + 7, + Some("jsonpatch"), + Some("application/json"), + Body::Text("{}"), + StatusCode::CREATED +)] + +pub async fn put_merge_patch( + _: usize, + strategy: Option<&str>, + content_type: Option<&str>, + body: Body, + status: StatusCode, +) -> eyre::Result<()> { + let http = Client::new(); + let key = random_key(); + + let mut req = http.key_put(&key); + + match body { + Body::Random(size) => { + req = req.body(random_body(size)); + } + Body::Text(text) => { + req = req.body(text.to_string()); + } + } + + if let Some(strategy) = strategy { + req = req.header("huly-merge-strategy", strategy); + } + + if let Some(content_type) = content_type { + req = req.header("content-type", content_type); + } + + let res = req.send().await?; + + check!(res.status() == status); + + Ok(()) +} diff --git a/tests/src/util.rs b/tests/src/util.rs index 1812c66c89..f318812894 100644 --- a/tests/src/util.rs +++ b/tests/src/util.rs @@ -55,6 +55,16 @@ impl ClientExt for Client { } } +pub trait ResponseExt { + fn header(&self, key: &str) -> Option<&str>; +} + +impl ResponseExt for tanu::http::Response { + fn header(&self, key: &str) -> Option<&str> { + self.headers().get(key).and_then(|v| v.to_str().ok()) + } +} + pub fn random_text(length: usize) -> String { let mut rng = rand::rng(); let charset: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789 ";