From ef5a5ebcec53b7cc96c6f95b92d54e3f7d688d86 Mon Sep 17 00:00:00 2001 From: Alexander Onnikov Date: Tue, 21 Oct 2025 14:11:15 +0700 Subject: [PATCH] feat: compact compact worker (#4) Signed-off-by: Alexander Onnikov --- server/src/compact.rs | 188 +++++++++++++++++++++++++++++++++++++++++ server/src/config.rs | 6 ++ server/src/handlers.rs | 53 +++++------- server/src/main.rs | 41 ++++++++- server/src/merge.rs | 24 ++++-- server/src/mutex.rs | 23 +++++ server/src/postgres.rs | 2 +- tests/src/compact.rs | 59 +++++++++++++ tests/src/main.rs | 1 + 9 files changed, 353 insertions(+), 44 deletions(-) create mode 100644 server/src/compact.rs create mode 100644 server/src/mutex.rs create mode 100644 tests/src/compact.rs diff --git a/server/src/compact.rs b/server/src/compact.rs new file mode 100644 index 0000000000..9d9fc94e32 --- /dev/null +++ b/server/src/compact.rs @@ -0,0 +1,188 @@ +use std::collections::HashSet; +use std::sync::Arc; + +use size::Size; +use tokio::sync::{RwLock, mpsc}; +use tracing::*; +use uuid::Uuid; + +use crate::config::CONFIG; +use crate::handlers::{ApiError, PartData}; +use crate::merge; +use crate::mutex::KeyMutex; +use crate::postgres::{ObjectPart, Pool}; +use crate::s3::S3Client; +use crate::{blob, postgres, recovery}; + +#[derive(Debug, Clone, Hash, Eq, PartialEq)] +pub struct CompactTask { + pub workspace: Uuid, + pub key: String, +} + +pub struct CompactWorker { + ingest_tx: mpsc::Sender, + ingest_handle: Arc>, + compact_handle: Arc>, +} + +impl Clone for CompactWorker { + fn clone(&self) -> Self { + CompactWorker { + ingest_tx: self.ingest_tx.clone(), + ingest_handle: self.ingest_handle.clone(), + compact_handle: self.compact_handle.clone(), + } + } +} + +impl CompactWorker { + pub fn new(s3: Arc, pool: Pool, lock: KeyMutex, buffer_size: usize) -> Self { + let (ingest_tx, ingest_rx) = mpsc::channel(buffer_size); + let (compact_tx, compact_rx) = mpsc::channel(buffer_size); + + let pending_tasks = Arc::new(RwLock::new(HashSet::new())); + let pending_tasks_ingest = pending_tasks.clone(); + let pending_tasks_compact = pending_tasks.clone(); + + let ingest_handle = tokio::spawn(async move { + debug!(buffer_size, "started ingest worker"); + Self::run_ingest_worker(ingest_rx, compact_tx, pending_tasks_ingest).await + }); + + let compact_handle = tokio::spawn(async move { + debug!(buffer_size, "started compact worker"); + Self::run_compact_worker( + compact_rx, + s3.clone(), + pool, + lock.clone(), + pending_tasks_compact, + ) + .await; + }); + + Self { + ingest_tx, + ingest_handle: Arc::new(ingest_handle), + compact_handle: Arc::new(compact_handle), + } + } + + async fn run_ingest_worker( + mut ingest_rx: mpsc::Receiver, + compact_tx: mpsc::Sender, + pending_tasks: Arc>>, + ) { + loop { + while let Some(task) = ingest_rx.recv().await { + let is_new = pending_tasks.write().await.insert(task.clone()); + if !is_new { + continue; + } + + if let Err(err) = compact_tx.send(task.clone()).await { + error!(%err, "failed to send compact task"); + pending_tasks.write().await.remove(&task); + } + } + } + } + + async fn run_compact_worker( + mut rx: mpsc::Receiver, + s3: Arc, + pool: Pool, + lock: KeyMutex, + pending_tasks: Arc>>, + ) { + loop { + while let Some(task) = rx.recv().await { + let CompactTask { workspace, key } = task.clone(); + + let _guard = lock.lock(workspace, key).await; + + pending_tasks.write().await.remove(&task); + + let res = compact(s3.clone(), pool.clone(), task.clone()).await; + match res { + Ok(_) => debug!(workspace = %task.workspace, key = %task.key, "blob compacted"), + Err(err) => error!(%err, "failed to compact"), + } + } + } + } + + pub async fn send(&self, parts: &Vec>) { + if parts.len() > CONFIG.compact_parts_limit { + let task = CompactTask { + workspace: parts[0].data.workspace, + key: parts[0].data.key.clone(), + }; + + let res = self.ingest_tx.send(task.clone()).await; + if let Err(err) = res { + warn!(%err, "failed to schedule compact"); + } + } + } + + pub async fn stop(&self) { + self.ingest_handle.abort(); + self.compact_handle.abort(); + } +} + +#[instrument(level = "debug", skip_all, fields(workspace, huly_key))] +async fn compact(s3: Arc, pool: Pool, task: CompactTask) -> anyhow::Result<(), ApiError> { + let pool = pool.clone(); + + let workspace = task.workspace; + let key = task.key; + + Span::current() + .record("workspace", workspace.to_string()) + .record("huly_key", &key); + + let parts = postgres::find_parts(&pool, task.workspace, &key).await?; + let first = &parts.first().unwrap().data; + let last = &parts.last().unwrap().data; + + let stream = merge::stream(s3.clone(), parts.to_vec()).await?; + + let uploaded = blob::upload( + &s3, + &pool, + Size::from_bytes(stream.content_length), + stream.stream, + ) + .await?; + + let inline = uploaded.inline.and_then(|inline| { + if inline.len() < CONFIG.inline_threshold.bytes() as usize { + Some(inline) + } else { + None + } + }); + + let part_data = PartData { + workspace, + key: key.to_owned(), + part: 0, + blob: uploaded.s3_key, + size: uploaded.length, + etag: last.etag.clone(), + date: last.date.clone(), + + headers: first.headers.clone(), + meta: first.meta.clone(), + merge_strategy: first.merge_strategy, + }; + let obj_parts = vec![&part_data]; + + postgres::set_part(&pool, workspace, &key, inline, &part_data).await?; + recovery::set_object(&s3, workspace, &key, obj_parts, None).await?; + + Ok(()) +} diff --git a/server/src/config.rs b/server/src/config.rs index 8376d0417d..1eee9f3886 100644 --- a/server/src/config.rs +++ b/server/src/config.rs @@ -39,6 +39,9 @@ pub struct Config { pub inline_threshold: Size, pub cache_control: String, + + pub compact_parts_limit: usize, + pub compact_buffer_size: usize, } pub mod hulyrs { @@ -69,6 +72,9 @@ pub static CONFIG: LazyLock = LazyLock::new(|| { inline_threshold = "100KB" cache_control = "public, no-cache" + + compact_parts_limit = 100 + compact_buffer_size = 1000 "#; let mut builder = diff --git a/server/src/handlers.rs b/server/src/handlers.rs index 1f82ea0f3e..43d86e4a90 100644 --- a/server/src/handlers.rs +++ b/server/src/handlers.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, fmt::Display, io, str::FromStr, sync::Arc, time::SystemTime}; +use std::{collections::HashMap, fmt::Display, io, str::FromStr, time::SystemTime}; use actix_web::{ HttpRequest, HttpResponse, @@ -13,13 +13,11 @@ use actix_web::{ use aws_sdk_s3::error::SdkError; use chrono::{DateTime, Utc}; use futures::{StreamExt, stream}; -use lockable::LockPool; use serde::{Deserialize, Serialize}; use size::Size; use tracing::*; use uuid::Uuid; -use crate::conditional; use crate::s3::S3Client; use crate::{ blob, @@ -27,6 +25,7 @@ use crate::{ merge, postgres::ObjectPart, }; +use crate::{compact::CompactWorker, conditional}; use crate::{ config::CONFIG, postgres::{self, Pool}, @@ -35,8 +34,8 @@ use crate::{merge::MergeStrategy, recovery}; #[derive(Deserialize, Debug)] pub struct ObjectPath { - workspace: Uuid, - key: String, + pub workspace: Uuid, + pub key: String, } #[derive(thiserror::Error, Debug)] @@ -188,23 +187,23 @@ async fn extract_range_header(request: &mut ServiceRequest) -> Option { .map(|header| header.0.to_string()) .ok() } -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize, Debug, Clone)] pub struct PartData { - workspace: Uuid, - key: String, - part: u32, + pub workspace: Uuid, + pub key: String, + pub part: u32, pub size: usize, pub blob: String, - etag: String, + pub etag: String, #[serde(default)] - date: DateTime, + pub date: DateTime, #[serde(skip_serializing_if = "Option::is_none")] - headers: Option>, + pub headers: Option>, #[serde(skip_serializing_if = "Option::is_none")] - meta: Option>, + pub meta: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub merge_strategy: Option, @@ -227,15 +226,6 @@ pub async fn put(request: HttpRequest, payload: Payload) -> HandlerResult>().unwrap().to_owned(); let s3 = request.app_data::>().unwrap().to_owned(); - let lock_pool = request - .app_data::>>>() - .unwrap() - .to_owned(); - - let _guard = lock_pool - .async_lock(format!("{}:{}", path.workspace, path.key)) - .await; - let parts = postgres::find_parts::(&pool, path.workspace, &path.key).await?; let conditionals = validate_put_conditionals(request.request(), &parts)?; @@ -305,15 +295,6 @@ pub async fn patch(request: HttpRequest, payload: Payload) -> HandlerResult>().unwrap().to_owned(); let s3 = request.app_data::>().unwrap().to_owned(); - let lock_pool = request - .app_data::>>>() - .unwrap() - .to_owned(); - - let _guard = lock_pool - .async_lock(format!("{}:{}", path.workspace, path.key)) - .await; - let parts = postgres::find_parts::(&pool, path.workspace, &path.key).await?; let mut response = if !parts.is_empty() { @@ -449,9 +430,15 @@ pub async fn get(request: HttpRequest) -> HandlerResult { response.insert_header((header::CONTENT_RANGE, content_range)); } - response.body(partial.stream) + response.body(SizedStream::new(partial.content_length, partial.stream)) + } + None => { + let compact = request.app_data::>().unwrap(); + compact.send(&parts).await; + + let stream = merge::stream(s3.clone(), parts).await?; + response.body(SizedStream::new(stream.content_length, stream.stream)) } - None => response.body(merge::stream(s3, parts).await?), } } } diff --git a/server/src/main.rs b/server/src/main.rs index 96947ebb92..baa343756b 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -8,7 +8,6 @@ use actix_web::{ middleware::{Next, from_fn}, web::{self, Data, Path}, }; -use lockable::LockPool; use tracing::*; use tracing_actix_web::TracingLogger; use uuid::Uuid; @@ -17,10 +16,12 @@ use hulyrs::services::jwt::actix::ServiceRequestExt; use hulyrs::services::otel; mod blob; +mod compact; mod conditional; mod config; mod handlers; mod merge; +mod mutex; mod patch; mod postgres; mod recovery; @@ -28,6 +29,8 @@ mod s3; use config::CONFIG; +use crate::mutex::KeyMutex; + fn initialize_tracing() { use opentelemetry::trace::TracerProvider; use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge; @@ -96,7 +99,7 @@ async fn main() -> anyhow::Result<()> { "configuration" ); - let lock = Arc::new(LockPool::::new()); + let lock = mutex::KeyMutex::new(); let postgres = postgres::pool().await?; let s3 = s3::client().await; @@ -110,6 +113,7 @@ async fn main() -> anyhow::Result<()> { let bind_to = SocketAddr::new(CONFIG.bind_host.as_str().parse()?, CONFIG.bind_port); + #[allow(dead_code)] async fn auth( mut request: ServiceRequest, next: Next, @@ -138,6 +142,30 @@ async fn main() -> anyhow::Result<()> { } } + async fn mutex( + mut request: ServiceRequest, + next: Next, + ) -> Result, Error> { + let path = request + .extract::>() + .await? + .into_inner(); + + let mutex = request.app_data::>().unwrap().to_owned(); + + let _guard = mutex.lock(path.workspace, path.key).await; + + next.call(request).await + } + + let compactor = compact::CompactWorker::new( + Arc::new(s3.clone()), + postgres.clone(), + lock.clone(), + CONFIG.compact_buffer_size, + ); + let compactor_handle = compactor.clone(); + let server = HttpServer::new(move || { let cors = Cors::default() .allow_any_origin() @@ -152,6 +180,7 @@ async fn main() -> anyhow::Result<()> { .app_data(Data::new(postgres.clone())) .app_data(Data::new(s3.clone())) .app_data(Data::new(lock.clone())) + .app_data(Data::new(compactor.clone())) .wrap(TracingLogger::default()) .wrap(cors) .service( @@ -159,8 +188,11 @@ async fn main() -> anyhow::Result<()> { .wrap(from_fn(auth)) .route(KEY_PATH, web::head().to(handlers::head)) .route(KEY_PATH, web::get().to(handlers::get)) - .route(KEY_PATH, web::put().to(handlers::put)) - .route(KEY_PATH, web::patch().to(handlers::patch)) + .route(KEY_PATH, web::put().to(handlers::put).wrap(from_fn(mutex))) + .route( + KEY_PATH, + web::patch().to(handlers::patch).wrap(from_fn(mutex)), + ) .route(KEY_PATH, web::delete().to(handlers::delete)), ) .route("/status", web::get().to(async || "ok")) @@ -171,6 +203,7 @@ async fn main() -> anyhow::Result<()> { info!("http listener on {}", bind_to); server.await?; + compactor_handle.stop().await; Ok(()) } diff --git a/server/src/merge.rs b/server/src/merge.rs index 604f2488a3..954f2c3c2f 100644 --- a/server/src/merge.rs +++ b/server/src/merge.rs @@ -1,6 +1,5 @@ use std::{io::Error as IoError, pin::Pin, sync::Arc}; -use actix_web::body::SizedStream; use actix_web::error::ErrorBadRequest; use async_stream::stream; use bytes::Bytes; @@ -89,7 +88,8 @@ pub fn validate_patch_body(merge_strategy: MergeStrategy, blob: &Blob) -> Handle pub struct PartialResponse { pub partial: bool, pub content_range: Option, - pub stream: SizedStream>>>>, + pub content_length: u64, + pub stream: Pin>>>, } #[instrument(level = "debug", skip_all)] @@ -120,15 +120,21 @@ pub async fn partial( Ok(PartialResponse { partial: part.data.size != content_length as usize, content_range, - stream: SizedStream::new(content_length, Box::pin(stream)), + content_length, + stream: Box::pin(stream), }) } +pub struct StreamResponse { + pub content_length: u64, + pub stream: Pin> + Send>>, +} + #[instrument(level = "debug", skip_all)] pub async fn stream( s3: Arc, parts: Vec>, -) -> anyhow::Result>>>>> { +) -> anyhow::Result { let first = parts.first().unwrap(); let merge_strategy = first.data.merge_strategy.unwrap(); @@ -164,7 +170,10 @@ pub async fn stream( } }; - Ok(SizedStream::new(content_length as u64, Box::pin(stream))) + Ok(StreamResponse { + content_length: content_length as u64, + stream: Box::pin(stream), + }) } MergeStrategy::JsonPatch => { @@ -197,7 +206,10 @@ pub async fn stream( yield Result::::Ok(Bytes::from(bytes)); }; - Ok(SizedStream::new(content_length, Box::pin(stream))) + Ok(StreamResponse { + content_length, + stream: Box::pin(stream), + }) } } } diff --git a/server/src/mutex.rs b/server/src/mutex.rs new file mode 100644 index 0000000000..261910dea5 --- /dev/null +++ b/server/src/mutex.rs @@ -0,0 +1,23 @@ +use std::sync::Arc; + +use lockable::LockPool; +use uuid::Uuid; + +#[derive(Clone)] +pub struct KeyMutex { + lock_pool: Arc>, +} + +impl KeyMutex { + pub fn new() -> Self { + let lock_pool = Arc::new(LockPool::::new()); + + KeyMutex { lock_pool } + } + + pub async fn lock(&self, workspace: Uuid, key: String) -> impl Drop + '_ { + self.lock_pool + .async_lock(format!("{}:{}", workspace, key)) + .await + } +} diff --git a/server/src/postgres.rs b/server/src/postgres.rs index fef62ee9a0..bd52e0b769 100644 --- a/server/src/postgres.rs +++ b/server/src/postgres.rs @@ -122,7 +122,7 @@ pub async fn insert_blob(pool: &Pool, key: &str, hash: &str) -> anyhow::Result<( Ok(()) } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct ObjectPart { pub inline: Option>, pub data: T, diff --git a/tests/src/compact.rs b/tests/src/compact.rs new file mode 100644 index 0000000000..be8ae4e68e --- /dev/null +++ b/tests/src/compact.rs @@ -0,0 +1,59 @@ +use serde_json::{self as json, Value, json}; +use tanu::{check, eyre, http::Client}; + +use crate::util::*; + +#[tanu::test(50)] +#[tanu::test(100)] +#[tanu::test(200)] +#[tanu::test(500)] +pub async fn compact_json(count: usize) -> eyre::Result<()> { + let key = random_key(); + + let http = Client::new(); + + let initial = json!({ + "a": 0 + }); + + // create new blob + let res = http + .key_put(&key) + .body(json::to_string(&initial)?) + .header("huly-merge-strategy", "jsonpatch") + .header("content-type", "application/json") + .send() + .await?; + + check!(res.status().is_success(), "{:#?}", res); + + for i in 0..count { + let patch = json!([ + { "op": "replace", "path": "/a", "value": i + 1}, + ]); + + let body: String = json::to_string(&patch)?; + let res = http + .key_patch(&key) + .body(body) + .header("content-type", "application/json-patch+json") + .send() + .await?; + + check!(res.status().is_success(), "{:#?}", res); + } + + let res = http.key_get(&key).send().await?; + check!(res.status().is_success(), "{:#?}", res); + let json = res.json::().await?; + assert_eq!(json, json!({ "a": count })); + + tokio::time::sleep(std::time::Duration::from_millis(1000)).await; + + let res = http.key_get(&key).send().await?; + check!(res.status().is_success(), "{:#?}", res); + let json = res.json::().await?; + assert_eq!(json, json!({ "a": count })); + + Ok(()) +} diff --git a/tests/src/main.rs b/tests/src/main.rs index 9a30d6798b..c6fe12ba78 100644 --- a/tests/src/main.rs +++ b/tests/src/main.rs @@ -1,4 +1,5 @@ mod auth; +mod compact; mod config; mod get; mod head;