feat: support persist manifest efficientlly

horaedb/Cargo.toml

@@ -45,6 +45,7 @@ itertools = "0.3"
 lazy_static = "1"
 tracing = "0.1"
 tracing-subscriber = "0.3"
+log = "0.4"
 
 # This profile optimizes for good runtime performance.
 [profile.release]

horaedb/metric_engine/Cargo.toml

@@ -39,6 +39,7 @@ datafusion = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
 lazy_static = { workspace = true }
+log = { workspace = true }
 macros = { workspace = true }
 object_store = { workspace = true }
 parquet = { workspace = true, features = ["object_store"] }

horaedb/metric_engine/src/manifest.rs

@@ -15,11 +15,21 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::Arc;
+
 use anyhow::Context;
 use bytes::Bytes;
+use futures::StreamExt;
+use log::{error, info};
 use object_store::{path::Path, PutPayload};
 use prost::Message;
-use tokio::sync::RwLock;
+use tokio::{
+    sync::{
+        mpsc::{self, Receiver, Sender},
+        RwLock,
+    },
+    time,
+};
 
 use crate::{
     sst::{FileId, FileMeta, SstFile},
@@ -29,13 +39,16 @@ use crate::{
 
 pub const PREFIX_PATH: &str = "manifest";
 pub const SNAPSHOT_FILENAME: &str = "snapshot";
+pub const SST_PREFIX: &str = "sst";
 
 pub struct Manifest {
     path: String,
     snapshot_path: Path,
+    sst_path: Path,
     store: ObjectStoreRef,
+    sender: Sender<SstsMergeTask>,
 
-    payload: RwLock<Payload>,
+    payload: Arc<RwLock<Payload>>,
 }
 
 pub struct Payload {
@@ -69,8 +82,29 @@ impl From<Payload> for pb_types::Manifest {
 }
 
 impl Manifest {
-    pub async fn try_new(path: String, store: ObjectStoreRef) -> Result<Self> {
+    pub async fn try_new(
+        path: String,
+        store: ObjectStoreRef,
+        merge_options: SstsMergeOptions,
+    ) -> Result<Self> {
         let snapshot_path = Path::from(format!("{path}/{SNAPSHOT_FILENAME}"));
+        let sst_path = Path::from(format!("{path}/{SST_PREFIX}"));
+
+        let (tx, rx) = mpsc::channel(merge_options.channel_size);
+        let mut ssts_merge = SstsMerge::try_new(
+            snapshot_path.clone(),
+            sst_path.clone(),
+            store.clone(),
+            rx,
+            merge_options,
+        )
+        .await?;
+
+        // Start merge ssts task
+        tokio::spawn(async move {
+            ssts_merge.run().await;
+        });
+
         let payload = match store.get(&snapshot_path).await {
             Ok(v) => {
                 let bytes = v
@@ -90,43 +124,42 @@ impl Manifest {
                 }
             }
         };
+        let payload = Arc::new(RwLock::new(payload));
 
         Ok(Self {
             path,
             snapshot_path,
+            sst_path,
             store,
-            payload: RwLock::new(payload),
+            payload,
+            sender: tx,
         })
     }
 
-    // TODO: Now this functions is poorly implemented, we concat new_sst to
-    // snapshot, and upload it back in a whole.
-    // In more efficient way, we can create a new diff file, and do compaction in
-    // background to merge them to `snapshot`.
     pub async fn add_file(&self, id: FileId, meta: FileMeta) -> Result<()> {
-        let mut payload = self.payload.write().await;
-        let mut tmp_ssts = payload.files.clone();
+        let new_sst_path = Path::from(format!("{}/{id}", self.sst_path));
         let new_sst = SstFile { id, meta };
-        tmp_ssts.push(new_sst.clone());
-        let pb_manifest = pb_types::Manifest {
-            files: tmp_ssts.into_iter().map(|f| f.into()).collect::<Vec<_>>(),
-        };
 
-        let mut buf = Vec::with_capacity(pb_manifest.encoded_len());
-        pb_manifest
+        let new_sst_payload = pb_types::SstFile::from(new_sst.clone());
+        let mut buf: Vec<u8> = Vec::with_capacity(new_sst_payload.encoded_len());
+        new_sst_payload
             .encode(&mut buf)
             .context("failed to encode manifest")?;
         let put_payload = PutPayload::from_bytes(Bytes::from(buf));
 
-        // 1. Persist the snapshot
+        // 1. Persist the sst file
         self.store
-            .put(&self.snapshot_path, put_payload)
+            .put(&new_sst_path, put_payload)
             .await
             .context("Failed to update manifest")?;
 
         // 2. Update cached payload
+        let mut payload = self.payload.write().await;
         payload.files.push(new_sst);
 
+        // 3. Schedule ssts merge
+        self.sender(SstsMergeTask::MergeSsts(1)).await;
+
         Ok(())
     }
 
@@ -140,4 +173,189 @@ impl Manifest {
             .cloned()
             .collect()
     }
+
+    async fn sender(&self, task: SstsMergeTask) {
+        if let Err(err) = self.sender.send(task).await {
+            error!("Failed to send merge ssts task, err: {:?}", err);
+        }
+    }
+}
+
+enum SstsMergeTask {
+    ForceMergeSsts,
+    MergeSsts(usize),
+}
+
+pub struct SstsMergeOptions {
+    channel_size: usize,
+    max_interval_seconds: usize,
+    merge_threshold: usize,
+}
+
+impl Default for SstsMergeOptions {
+    fn default() -> Self {
+        Self {
+            channel_size: 100,
+            max_interval_seconds: 5,
+            merge_threshold: 50,
+        }
+    }
+}
+
+struct SstsMerge {
+    snapshot_path: Path,
+    sst_path: Path,
+    store: ObjectStoreRef,
+    receiver: Receiver<SstsMergeTask>,
+    sst_num: usize,
+    merge_options: SstsMergeOptions,
+}
+
+impl SstsMerge {
+    async fn try_new(
+        snapshot_path: Path,
+        sst_path: Path,
+        store: ObjectStoreRef,
+        rx: Receiver<SstsMergeTask>,
+        merge_options: SstsMergeOptions,
+    ) -> Result<Self> {
+        let ssts_merge = Self {
+            snapshot_path,
+            sst_path,
+            store,
+            receiver: rx,
+            sst_num: 0,
+            merge_options,
+        };
+        // Merge all ssts when start
+        ssts_merge.merge_ssts().await?;
+
+        Ok(ssts_merge)
+    }
+
+    async fn run(&mut self) {
+        let interval = time::Duration::from_secs(self.merge_options.max_interval_seconds as u64);
+        let threshold = self.merge_options.merge_threshold;
+
+        loop {
+            match time::timeout(interval, self.receiver.recv()).await {
+                Ok(Some(SstsMergeTask::ForceMergeSsts)) => match self.merge_ssts().await {
+                    Ok(_) => {
+                        self.sst_num = 0;
+                    }
+                    Err(err) => {
+                        error!("Failed to force merge ssts, err: {:?}", err);
+                    }
+                },
+                Ok(Some(SstsMergeTask::MergeSsts(num))) => {
+                    self.sst_num += num;
+                    if self.sst_num >= threshold {
+                        match self.merge_ssts().await {
+                            Ok(_) => {
+                                self.sst_num = 0;
+                            }
+                            Err(err) => {
+                                error!("Failed to merge ssts, err: {:?}", err);
+                            }
+                        }
+                    }
+                }
+                Ok(None) => {
+                    // The channel is disconnected.
+                    info!("Channel disconnected, merge ssts task exit");
+                    break;
+                }
+                Err(_) => {
+                    info!("Timeout receive merge ssts task");
+                }
+            }
+        }
+    }
+
+    async fn merge_ssts(&self) -> Result<()> {
+        let meta_infos = self
+            .store
+            .list(Some(&self.sst_path))
+            .collect::<Vec<_>>()
+            .await;
+        if meta_infos.is_empty() {
+            return Ok(());
+        }
+
+        let mut paths = Vec::with_capacity(meta_infos.len());
+        for meta_info in meta_infos {
+            let path = meta_info
+                .context("failed to get path of manifest sst")?
+                .location;
+            paths.push(path);
+        }
+
+        let mut sst_files = Vec::with_capacity(paths.len());
+        for path in &paths {
+            let bytes = self
+                .store
+                .get(path)
+                .await
+                .context("failed to read sst file")?
+                .bytes()
+                .await
+                .context("failed to read sst file")?;
+            let pb_sst = pb_types::SstFile::decode(bytes).context("failed to decode sst file")?;
+            sst_files.push(SstFile::try_from(pb_sst)?);
+        }
+
+        let mut payload = match self.store.get(&self.snapshot_path).await {
+            Ok(v) => {
+                let bytes = v
+                    .bytes()
+                    .await
+                    .context("failed to read manifest snapshot")?;
+                let pb_payload = pb_types::Manifest::decode(bytes)
+                    .context("failed to decode manifest snapshot")?;
+                Payload::try_from(pb_payload)?
+            }
+            Err(err) => {
+                if err.to_string().contains("not found") {
+                    Payload { files: vec![] }
+                } else {
+                    let context = format!(
+                        "Failed to get manifest snapshot, path:{}",
+                        self.snapshot_path
+                    );
+                    return Err(AnyhowError::new(err).context(context).into());
+                }
+            }
+        };
+
+        payload.files.extend(sst_files.clone());
+        let pb_manifest = pb_types::Manifest {
+            files: payload
+                .files
+                .into_iter()
+                .map(|f| f.into())
+                .collect::<Vec<_>>(),
+        };
+
+        let mut buf = Vec::with_capacity(pb_manifest.encoded_len());
+        pb_manifest
+            .encode(&mut buf)
+            .context("failed to encode manifest")?;
+        let put_payload = PutPayload::from_bytes(Bytes::from(buf));
+
+        // 1. Persist the snapshot
+        self.store
+            .put(&self.snapshot_path, put_payload)
+            .await
+            .context("Failed to update manifest")?;
+
+        // 2. Delete the old sst files
+        for path in &paths {
+            self.store
+                .delete(path)
+                .await
+                .context("failed to delete sst file")?;
+        }
+
+        Ok(())
+    }
 }

horaedb/metric_engine/src/storage.rs

@@ -47,7 +47,7 @@ use parquet::{
 };
 
 use crate::{
-    manifest::Manifest,
+    manifest::{Manifest, SstsMergeOptions},
     read::DefaultParquetFileReaderFactory,
     sst::{allocate_id, FileId, FileMeta},
     types::{ObjectStoreRef, TimeRange, Timestamp, WriteOptions, WriteResult},
@@ -112,10 +112,15 @@ impl CloudObjectStorage {
         num_primary_key: usize,
         timestamp_index: usize,
         write_options: WriteOptions,
+        merge_options: SstsMergeOptions,
     ) -> Result<Self> {
         let manifest_prefix = crate::manifest::PREFIX_PATH;
-        let manifest =
-            Manifest::try_new(format!("{root_path}/{manifest_prefix}"), store.clone()).await?;
+        let manifest = Manifest::try_new(
+            format!("{root_path}/{manifest_prefix}"),
+            store.clone(),
+            merge_options,
+        )
+        .await?;
         let df_schema = DFSchema::try_from(arrow_schema.clone()).context("build DFSchema")?;
         let write_props = Self::build_write_props(write_options, num_primary_key);
         Ok(Self {
@@ -326,6 +331,7 @@ mod tests {
     use object_store::local::LocalFileSystem;
 
     use super::*;
+    use crate::manifest::SstsMergeOptions;
 
     #[tokio::test]
     async fn test_sort_batch() {
@@ -344,6 +350,7 @@ mod tests {
             1,
             1,
             WriteOptions::default(),
+            SstsMergeOptions::default(),
         )
         .await
         .unwrap();