feat: compact compact worker (#4)

Signed-off-by: Alexander Onnikov <Alexander.Onnikov@xored.com>

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<CompactTask>,
+    ingest_handle: Arc<tokio::task::JoinHandle<()>>,
+    compact_handle: Arc<tokio::task::JoinHandle<()>>,
+}
+
+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<S3Client>, 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<CompactTask>,
+        compact_tx: mpsc::Sender<CompactTask>,
+        pending_tasks: Arc<RwLock<HashSet<CompactTask>>>,
+    ) {
+        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<CompactTask>,
+        s3: Arc<S3Client>,
+        pool: Pool,
+        lock: KeyMutex,
+        pending_tasks: Arc<RwLock<HashSet<CompactTask>>>,
+    ) {
+        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<ObjectPart<PartData>>) {
+        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<S3Client>, 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(())
+}

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<Config> = LazyLock::new(|| {
         inline_threshold = "100KB"
 
         cache_control = "public, no-cache"
+
+        compact_parts_limit = 100
+        compact_buffer_size = 1000
     "#;
 
     let mut builder =

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<String> {
         .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<Utc>,
+    pub date: DateTime<Utc>,
 
     #[serde(skip_serializing_if = "Option::is_none")]
-    headers: Option<HashMap<String, String>>,
+    pub headers: Option<HashMap<String, String>>,
 
     #[serde(skip_serializing_if = "Option::is_none")]
-    meta: Option<HashMap<String, String>>,
+    pub meta: Option<HashMap<String, String>>,
 
     #[serde(skip_serializing_if = "Option::is_none")]
     pub merge_strategy: Option<MergeStrategy>,
@@ -227,15 +226,6 @@ pub async fn put(request: HttpRequest, payload: Payload) -> HandlerResult<HttpRe
     let pool = request.app_data::<Data<Pool>>().unwrap().to_owned();
     let s3 = request.app_data::<Data<S3Client>>().unwrap().to_owned();
 
-    let lock_pool = request
-        .app_data::<Data<Arc<LockPool<String>>>>()
-        .unwrap()
-        .to_owned();
-
-    let _guard = lock_pool
-        .async_lock(format!("{}:{}", path.workspace, path.key))
-        .await;
-
     let parts = postgres::find_parts::<PartData>(&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<Http
     let pool = request.app_data::<Data<Pool>>().unwrap().to_owned();
     let s3 = request.app_data::<Data<S3Client>>().unwrap().to_owned();
 
-    let lock_pool = request
-        .app_data::<Data<Arc<LockPool<String>>>>()
-        .unwrap()
-        .to_owned();
-
-    let _guard = lock_pool
-        .async_lock(format!("{}:{}", path.workspace, path.key))
-        .await;
-
     let parts = postgres::find_parts::<PartData>(&pool, path.workspace, &path.key).await?;
 
     let mut response = if !parts.is_empty() {
@@ -449,9 +430,15 @@ pub async fn get(request: HttpRequest) -> HandlerResult<HttpResponse> {
                             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::<Data<CompactWorker>>().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?),
                 }
             }
         }

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::<String>::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<impl MessageBody>,
@@ -138,6 +142,30 @@ async fn main() -> anyhow::Result<()> {
         }
     }
 
+    async fn mutex(
+        mut request: ServiceRequest,
+        next: Next<impl MessageBody>,
+    ) -> Result<ServiceResponse<impl MessageBody>, Error> {
+        let path = request
+            .extract::<Path<handlers::ObjectPath>>()
+            .await?
+            .into_inner();
+
+        let mutex = request.app_data::<Data<KeyMutex>>().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(())
 }

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<String>,
-    pub stream: SizedStream<Pin<Box<dyn Stream<Item = Result<Bytes, IoError>>>>>,
+    pub content_length: u64,
+    pub stream: Pin<Box<dyn Stream<Item = Result<Bytes, IoError>>>>,
 }
 
 #[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<Box<dyn Stream<Item = Result<Bytes, IoError>> + Send>>,
+}
+
 #[instrument(level = "debug", skip_all)]
 pub async fn stream(
     s3: Arc<S3Client>,
     parts: Vec<ObjectPart<PartData>>,
-) -> anyhow::Result<SizedStream<Pin<Box<dyn Stream<Item = Result<Bytes, IoError>>>>>> {
+) -> anyhow::Result<StreamResponse> {
     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::<Bytes, IoError>::Ok(Bytes::from(bytes));
             };
 
-            Ok(SizedStream::new(content_length, Box::pin(stream)))
+            Ok(StreamResponse {
+                content_length,
+                stream: Box::pin(stream),
+            })
         }
     }
 }

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<LockPool<String>>,
+}
+
+impl KeyMutex {
+    pub fn new() -> Self {
+        let lock_pool = Arc::new(LockPool::<String>::new());
+
+        KeyMutex { lock_pool }
+    }
+
+    pub async fn lock(&self, workspace: Uuid, key: String) -> impl Drop + '_ {
+        self.lock_pool
+            .async_lock(format!("{}:{}", workspace, key))
+            .await
+    }
+}

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<T: DeserializeOwned + std::fmt::Debug> {
     pub inline: Option<Vec<u8>>,
     pub data: T,

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::<Value>().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::<Value>().await?;
+    assert_eq!(json, json!({ "a": count }));
+
+    Ok(())
+}

tests/src/main.rs

@@ -1,4 +1,5 @@
 mod auth;
+mod compact;
 mod config;
 mod get;
 mod head;