chore(cubestore): Small metastore refactoring (#5790)

Author: waralexrom

rust/cubestore/cubestore/src/metastore/rocks_fs.rs

@@ -24,14 +24,25 @@ pub trait MetaStoreFs: Send + Sync {
     ) -> Result<Arc<RocksStore>, CubeError>;
     async fn upload_log(
         &self,
-        log_name: &str,
+        dir: &str,
+        seq_number: u64,
         serializer: &WriteBatchContainer,
     ) -> Result<u64, CubeError>;
     async fn upload_checkpoint(
         &self,
         remote_path: String,
         checkpoint_path: PathBuf,
     ) -> Result<(), CubeError>;
+    async fn check_rocks_store(
+        &self,
+        rocks_store: Arc<RocksStore>,
+        snapshot: Option<u128>,
+    ) -> Result<Arc<RocksStore>, CubeError>;
+    async fn load_metastore_logs(
+        &self,
+        snapshot: u128,
+        rocks_store: &Arc<RocksStore>,
+    ) -> Result<(), CubeError>;
 }
 
 #[derive(Clone)]
@@ -201,48 +212,22 @@ impl BaseRocksStoreFs {
         };
         Ok(last_metastore_snapshot)
     }
-    pub async fn load_metastore_logs(
-        &self,
-        snapshot: u128,
-        rocks_store: &Arc<RocksStore>,
-    ) -> Result<(), CubeError> {
-        let logs_to_batch = self
+
+    pub async fn files_to_load(&self, snapshot: u128) -> Result<Vec<(String, u64)>, CubeError> {
+        let res = self
             .remote_fs
-            .list(&format!("{}-{}-logs", self.name, snapshot))
-            .await?;
-        let mut logs_to_batch_to_seq = logs_to_batch
+            .list_with_metadata(&format!("{}-{}", self.name, snapshot))
+            .await?
             .into_iter()
-            .map(|f| -> Result<_, CubeError> {
-                let last = f
-                    .split("/")
-                    .last()
-                    .ok_or(CubeError::internal(format!("Can't split path: {}", f)))?;
-                let result = last.replace(".flex", "").parse::<usize>().map_err(|e| {
-                    CubeError::internal(format!("Can't parse flex path {}: {}", f, e))
-                })?;
-                Ok((f, result))
-            })
-            .collect::<Result<Vec<_>, _>>()?;
-        logs_to_batch_to_seq.sort_unstable_by_key(|(_, seq)| *seq);
-
-        for (log_file, _) in logs_to_batch_to_seq.iter() {
-            let path_to_log = self.remote_fs.local_file(log_file).await?;
-            let batch = WriteBatchContainer::read_from_file(&path_to_log).await;
-            if let Ok(batch) = batch {
-                let db = rocks_store.db.clone();
-                db.write(batch.write_batch())?;
-            } else if let Err(e) = batch {
-                error!(
-                    "Corrupted {} WAL file. Discarding: {:?} {}",
-                    self.name, log_file, e
-                );
-                break;
-            }
-        }
-        Ok(())
+            .map(|f| (f.remote_path, f.file_size))
+            .collect::<Vec<_>>();
+        Ok(res)
     }
+}
 
-    pub async fn check_rocks_store(
+#[async_trait]
+impl MetaStoreFs for BaseRocksStoreFs {
+    async fn check_rocks_store(
         &self,
         rocks_store: Arc<RocksStore>,
         snapshot: Option<u128>,
@@ -255,21 +240,6 @@ impl BaseRocksStoreFs {
 
         Ok(rocks_store)
     }
-
-    pub async fn files_to_load(&self, snapshot: u128) -> Result<Vec<(String, u64)>, CubeError> {
-        let res = self
-            .remote_fs
-            .list_with_metadata(&format!("{}-{}", self.name, snapshot))
-            .await?
-            .into_iter()
-            .map(|f| (f.remote_path, f.file_size))
-            .collect::<Vec<_>>();
-        Ok(res)
-    }
-}
-
-#[async_trait]
-impl MetaStoreFs for BaseRocksStoreFs {
     async fn load_from_remote(
         self: Arc<Self>,
         path: &str,
@@ -321,10 +291,12 @@ impl MetaStoreFs for BaseRocksStoreFs {
 
     async fn upload_log(
         &self,
-        log_name: &str,
+        dir: &str,
+        seq_number: u64,
         serializer: &WriteBatchContainer,
     ) -> Result<u64, CubeError> {
-        let file_name = self.remote_fs.local_file(log_name).await?;
+        let log_name = format!("{}/{}.flex", dir, seq_number);
+        let file_name = self.remote_fs.local_file(&log_name).await?;
         serializer.write_to_file(&file_name).await?;
         // TODO persist file size
         self.remote_fs.upload_file(&file_name, &log_name).await
@@ -344,6 +316,46 @@ impl MetaStoreFs for BaseRocksStoreFs {
 
         Ok(())
     }
+    async fn load_metastore_logs(
+        &self,
+        snapshot: u128,
+        rocks_store: &Arc<RocksStore>,
+    ) -> Result<(), CubeError> {
+        let logs_to_batch = self
+            .remote_fs
+            .list(&format!("{}-{}-logs", self.name, snapshot))
+            .await?;
+        let mut logs_to_batch_to_seq = logs_to_batch
+            .into_iter()
+            .map(|f| -> Result<_, CubeError> {
+                let last = f
+                    .split("/")
+                    .last()
+                    .ok_or(CubeError::internal(format!("Can't split path: {}", f)))?;
+                let result = last.replace(".flex", "").parse::<usize>().map_err(|e| {
+                    CubeError::internal(format!("Can't parse flex path {}: {}", f, e))
+                })?;
+                Ok((f, result))
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+        logs_to_batch_to_seq.sort_unstable_by_key(|(_, seq)| *seq);
+
+        for (log_file, _) in logs_to_batch_to_seq.iter() {
+            let path_to_log = self.remote_fs.local_file(log_file).await?;
+            let batch = WriteBatchContainer::read_from_file(&path_to_log).await;
+            if let Ok(batch) = batch {
+                let db = rocks_store.db.clone();
+                db.write(batch.write_batch())?;
+            } else if let Err(e) = batch {
+                error!(
+                    "Corrupted {} WAL file. Discarding: {:?} {}",
+                    self.name, log_file, e
+                );
+                break;
+            }
+        }
+        Ok(())
+    }
 }
 
 crate::di_service!(BaseRocksStoreFs, [MetaStoreFs]);

rust/cubestore/cubestore/src/metastore/rocks_store.rs

@@ -625,12 +625,10 @@ impl RocksStore {
 
         if max.is_some() {
             let checkpoint_time = self.last_checkpoint_time.read().await;
-            let log_name = format!(
-                "{}-logs/{}.flex",
-                self.get_store_path(&checkpoint_time),
-                min.unwrap()
-            );
-            self.metastore_fs.upload_log(&log_name, &serializer).await?;
+            let dir_name = format!("{}-logs", self.get_store_path(&checkpoint_time));
+            self.metastore_fs
+                .upload_log(&dir_name, min.unwrap(), &serializer)
+                .await?;
             let mut seq = self.last_upload_seq.write().await;
             *seq = max.unwrap();
             self.write_completed_notify.notify_waiters();