feat(cubestore): Streaming optimizations (#6228)

Author: waralexrom

rust/cubestore/cubestore/src/app_metrics.rs

@@ -21,3 +21,32 @@ pub static CACHE_QUERY_TIME_MS: Histogram = metrics::histogram("cs.sql.query.cac
 /// Incoming queue queries.
 pub static QUEUE_QUERIES: Counter = metrics::counter("cs.sql.query.queue");
 pub static QUEUE_QUERY_TIME_MS: Histogram = metrics::histogram("cs.sql.query.queue.ms");
+
+pub static STREAMING_ROWS_READ: Counter = metrics::counter("cs.streaming.rows");
+pub static STREAMING_CHUNKS_READ: Counter = metrics::counter("cs.streaming.chunks");
+pub static STREAMING_LASTOFFSET: Gauge = metrics::gauge("cs.streaming.lastoffset");
+pub static IN_MEMORY_CHUNKS_COUNT: Gauge = metrics::gauge("cs.workers.in_memory_chunks)");
+pub static IN_MEMORY_CHUNKS_ROWS: Gauge = metrics::gauge("cs.workers.in_memory_chunks.rows)");
+pub static IN_MEMORY_CHUNKS_MEMORY: Gauge = metrics::gauge("cs.workers.in_memory_chunks.memory)");
+pub static STREAMING_IMPORT_TIME: Histogram = metrics::histogram("cs.streaming.import_time.ms");
+pub static STREAMING_PARTITION_TIME: Histogram =
+    metrics::histogram("cs.streaming.partition_time.ms");
+pub static STREAMING_UPLOAD_TIME: Histogram = metrics::histogram("cs.streaming.upload_time.ms");
+pub static STREAMING_ROUNDTRIP_TIME: Histogram =
+    metrics::histogram("cs.streaming.roundtrip_time.ms");
+pub static STREAMING_ROUNDTRIP_ROWS: Histogram = metrics::histogram("cs.streaming.roundtrip_rows");
+pub static STREAMING_ROUNDTRIP_CHUNKS: Histogram =
+    metrics::histogram("cs.streaming.roundtrip_chunks");
+pub static STREAMING_LAG: Gauge = metrics::gauge("cs.streaming.lag");
+
+pub static METASTORE_QUEUE: Gauge = metrics::gauge("cs.metastore.queue_size");
+pub static METASTORE_READ_OPERATION: Histogram =
+    metrics::histogram("cs.metastore.read_operation.ms");
+pub static METASTORE_INNER_READ_OPERATION: Histogram =
+    metrics::histogram("cs.metastore.inner_read_operation.ms");
+pub static METASTORE_WRITE_OPERATION: Histogram =
+    metrics::histogram("cs.metastore.write_operation.ms");
+pub static METASTORE_INNER_WRITE_OPERATION: Histogram =
+    metrics::histogram("cs.metastore.inner_write_operation.ms");
+pub static METASTORE_READ_OUT_QUEUE_OPERATION: Histogram =
+    metrics::histogram("cs.metastore.read_out_queue_operation.ms");

rust/cubestore/cubestore/src/config/mod.rs

@@ -906,7 +906,7 @@ impl Config {
                 ),
                 stream_replay_check_interval_secs: env_parse(
                     "CUBESTORE_STREAM_REPLAY_CHECK_INTERVAL",
-                    0,
+                    60,
                 ),
                 check_ws_orphaned_messages_interval_secs: env_parse(
                     "CUBESTORE_CHECK_WS_ORPHANED_MESSAGES_INTERVAL",

rust/cubestore/cubestore/src/metastore/mod.rs

@@ -863,6 +863,11 @@ pub trait MetaStore: DIService + Send + Sync {
     ) -> Result<Vec<IdRow<Partition>>, CubeError>;
     async fn get_index(&self, index_id: u64) -> Result<IdRow<Index>, CubeError>;
 
+    async fn get_index_with_active_partitions_out_of_queue(
+        &self,
+        index_id: u64,
+    ) -> Result<(IdRow<Index>, Vec<IdRow<Partition>>), CubeError>;
+
     async fn create_partitioned_index(
         &self,
         schema: String,
@@ -927,6 +932,7 @@ pub trait MetaStore: DIService + Send + Sync {
         max: Option<Row>,
         in_memory: bool,
     ) -> Result<IdRow<Chunk>, CubeError>;
+    async fn insert_chunks(&self, chunks: Vec<Chunk>) -> Result<Vec<IdRow<Chunk>>, CubeError>;
     async fn get_chunk(&self, chunk_id: u64) -> Result<IdRow<Chunk>, CubeError>;
     async fn get_chunks_out_of_queue(&self, ids: Vec<u64>) -> Result<Vec<IdRow<Chunk>>, CubeError>;
     async fn get_partitions_out_of_queue(
@@ -1046,11 +1052,11 @@ pub trait MetaStore: DIService + Send + Sync {
     async fn all_replay_handles_to_merge(
         &self,
     ) -> Result<Vec<(IdRow<ReplayHandle>, bool)>, CubeError>;
-    async fn update_replay_handle_failed(
+    async fn update_replay_handle_failed_if_exists(
         &self,
         id: u64,
         failed: bool,
-    ) -> Result<IdRow<ReplayHandle>, CubeError>;
+    ) -> Result<(), CubeError>;
     async fn replace_replay_handles(
         &self,
         old_ids: Vec<u64>,
@@ -2909,6 +2915,28 @@ impl MetaStore for RocksMetaStore {
         .await
     }
 
+    async fn get_index_with_active_partitions_out_of_queue(
+        &self,
+        index_id: u64,
+    ) -> Result<(IdRow<Index>, Vec<IdRow<Partition>>), CubeError> {
+        self.read_operation_out_of_queue(move |db_ref| {
+            let index = IndexRocksTable::new(db_ref.clone()).get_row_or_not_found(index_id)?;
+            let rocks_partition = PartitionRocksTable::new(db_ref);
+
+            let partitions = rocks_partition
+                .get_rows_by_index(
+                    &PartitionIndexKey::ByIndexId(index.get_id()),
+                    &PartitionRocksIndex::IndexId,
+                )?
+                .into_iter()
+                .filter(|r| r.get_row().active)
+                .collect::<Vec<_>>();
+
+            Ok((index, partitions))
+        })
+        .await
+    }
+
     #[tracing::instrument(level = "trace", skip(self, key_columns))]
     async fn create_partitioned_index(
         &self,
@@ -3086,6 +3114,22 @@ impl MetaStore for RocksMetaStore {
         .await
     }
 
+    #[tracing::instrument(level = "trace", skip(self, chunks))]
+    async fn insert_chunks(&self, chunks: Vec<Chunk>) -> Result<Vec<IdRow<Chunk>>, CubeError> {
+        self.write_operation(move |db_ref, batch_pipe| {
+            let rocks_chunk = ChunkRocksTable::new(db_ref.clone());
+            let mut result = Vec::with_capacity(chunks.len());
+
+            for chunk in chunks.into_iter() {
+                let id_row = rocks_chunk.insert(chunk, batch_pipe)?;
+                result.push(id_row);
+            }
+
+            Ok(result)
+        })
+        .await
+    }
+
     #[tracing::instrument(level = "trace", skip(self))]
     async fn get_chunk(&self, chunk_id: u64) -> Result<IdRow<Chunk>, CubeError> {
         self.read_operation(move |db_ref| {
@@ -3723,17 +3767,18 @@ impl MetaStore for RocksMetaStore {
     }
 
     #[tracing::instrument(level = "trace", skip(self))]
-    async fn update_replay_handle_failed(
+    async fn update_replay_handle_failed_if_exists(
         &self,
         id: u64,
         failed: bool,
-    ) -> Result<IdRow<ReplayHandle>, CubeError> {
+    ) -> Result<(), CubeError> {
         self.write_operation(move |db_ref, batch_pipe| {
-            Ok(ReplayHandleRocksTable::new(db_ref.clone()).update_with_fn(
-                id,
-                |h| h.set_failed_to_persist_chunks(failed),
-                batch_pipe,
-            )?)
+            let table = ReplayHandleRocksTable::new(db_ref.clone());
+            if table.get_row(id)?.is_some() {
+                table.update_with_fn(id, |h| h.set_failed_to_persist_chunks(failed), batch_pipe)?;
+            }
+
+            Ok(())
         })
         .await
     }
@@ -3759,10 +3804,11 @@ impl MetaStore for RocksMetaStore {
                     &ChunkRocksIndex::ReplayHandleId,
                 )?;
 
-                if !chunks.is_empty() {
+                let active_chunks = chunks.iter().filter(|c| c.get_row().active() || !c.get_row().uploaded()).collect::<Vec<_>>();
+                if !active_chunks.is_empty() {
                     return Err(CubeError::internal(format!(
-                        "Can't merge replay handle with chunks: {:?}",
-                        replay_handle
+                        "Can't merge replay handle with chunks: {:?}, {}",
+                        replay_handle, active_chunks[0].get_id()
                     )))
                 }
 
@@ -3811,7 +3857,9 @@ impl MetaStore for RocksMetaStore {
                 )?;
                 result.push((
                     replay_handle,
-                    chunks.iter().filter(|c| c.get_row().active()).count() == 0,
+                    chunks
+                        .iter()
+                        .all(|c| !c.get_row().active() && c.get_row().uploaded()),
                 ));
             }
             Ok(result)

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

@@ -733,7 +733,6 @@ impl RocksStore {
                 }
             }
         }
-
         Ok(spawn_res)
     }
 

rust/cubestore/cubestore/src/metastore/rocks_table.rs

@@ -263,6 +263,7 @@ where
 
 pub struct IndexScanIter<'a, RT: RocksTable + ?Sized> {
     table: &'a RT,
+    index_id: u32,
     secondary_key_val: Vec<u8>,
     secondary_key_hash: Vec<u8>,
     iter: DBIterator<'a>,
@@ -289,7 +290,30 @@ where
                         continue;
                     }
 
-                    return Some(self.table.get_row_or_not_found(row_id));
+                    let result = match self.table.get_row(row_id) {
+                        Ok(Some(row)) => Ok(row),
+                        Ok(None) => {
+                            let index = self.table.get_index_by_id(self.index_id);
+                            match self.table.rebuild_index(&index) {
+                                Ok(_) => {
+                                    Err(CubeError::internal(format!(
+                                        "Row exists in secondary index however missing in {:?} table: {}. Repairing index.",
+                                        self.table, row_id
+                                    )))
+                                }
+                                Err(e) => {
+                                    Err(CubeError::internal(format!(
+                                        "Error while rebuilding secondary index for {:?} table: {:?}",
+                                        self.table, e
+                                                )))
+                                }
+
+                            }
+                        }
+                        Err(e) => Err(e),
+                    };
+
+                    return Some(result);
                 };
             } else {
                 return None;
@@ -1017,12 +1041,10 @@ pub trait RocksTable: BaseRocksTable + Debug + Send + Sync {
             .to_vec();
         let secondary_key_val = secondary_index.key_to_bytes(&row_key);
 
+        let index_id = RocksSecondaryIndex::get_id(secondary_index);
         let key_len = secondary_key_hash.len();
-        let key_min = RowKey::SecondaryIndex(
-            Self::index_id(RocksSecondaryIndex::get_id(secondary_index)),
-            secondary_key_hash.clone(),
-            0,
-        );
+        let key_min =
+            RowKey::SecondaryIndex(Self::index_id(index_id), secondary_key_hash.clone(), 0);
 
         let mut opts = ReadOptions::default();
         opts.set_prefix_same_as_start(true);
@@ -1033,6 +1055,7 @@ pub trait RocksTable: BaseRocksTable + Debug + Send + Sync {
 
         Ok(IndexScanIter {
             table: self,
+            index_id,
             secondary_key_val,
             secondary_key_hash,
             iter,

rust/cubestore/cubestore/src/queryplanner/test_utils.rs

@@ -622,11 +622,11 @@ impl MetaStore for MetaStoreMock {
         panic!("MetaStore mock!")
     }
 
-    async fn update_replay_handle_failed(
+    async fn update_replay_handle_failed_if_exists(
         &self,
         _id: u64,
         _failed: bool,
-    ) -> Result<IdRow<ReplayHandle>, CubeError> {
+    ) -> Result<(), CubeError> {
         panic!("MetaStore mock!")
     }
 
@@ -685,6 +685,15 @@ impl MetaStore for MetaStoreMock {
     ) -> Result<Vec<IdRow<Chunk>>, CubeError> {
         panic!("MetaStore mock!")
     }
+    async fn get_index_with_active_partitions_out_of_queue(
+        &self,
+        _index_id: u64,
+    ) -> Result<(IdRow<Index>, Vec<IdRow<Partition>>), CubeError> {
+        panic!("MetaStore mock!")
+    }
+    async fn insert_chunks(&self, _chunks: Vec<Chunk>) -> Result<Vec<IdRow<Chunk>>, CubeError> {
+        panic!("MetaStore mock!")
+    }
 }
 
 crate::di_service!(MetaStoreMock, [MetaStore]);

rust/cubestore/cubestore/src/scheduler/mod.rs

@@ -2,6 +2,7 @@ use crate::cluster::{pick_worker_by_ids, Cluster};
 use crate::config::ConfigObj;
 use crate::metastore::job::{Job, JobStatus, JobType};
 use crate::metastore::partition::partition_file_name;
+use crate::metastore::replay_handle::ReplayHandle;
 use crate::metastore::replay_handle::{
     subtract_from_right_seq_pointer_by_location, subtract_if_covers_seq_pointer_by_location,
     union_seq_pointer_by_location, SeqPointerForLocation,
@@ -206,7 +207,10 @@ impl SchedulerImpl {
         )
         .await
         {
-            error!("Error scheduling partitions compaction: {}", e);
+            error!(
+                "Error scheduling deactivation chunks without partitions: {}",
+                e
+            );
         }
 
         if let Err(e) = warn_long_fut(
@@ -314,6 +318,12 @@ impl SchedulerImpl {
     /// Otherwise we just subtract it from resulting `SeqPointer` so freshly created `ReplayHandle`
     /// can't remove failed one.
     pub async fn merge_replay_handles(&self) -> Result<(), CubeError> {
+        fn is_newest_handle(handle: &IdRow<ReplayHandle>) -> bool {
+            Utc::now()
+                .signed_duration_since(handle.get_row().created_at().clone())
+                .num_seconds()
+                < 60
+        }
         let (failed, mut without_failed) = self
             .meta_store
             .all_replay_handles_to_merge()
@@ -339,7 +349,7 @@ impl SchedulerImpl {
             let handles = handles.collect::<Vec<_>>();
             for (handle, _) in handles
                 .iter()
-                .filter(|(_, no_active_chunks)| *no_active_chunks)
+                .filter(|(handle, no_active_chunks)| !is_newest_handle(handle) && *no_active_chunks)
             {
                 union_seq_pointer_by_location(
                     &mut seq_pointer_by_location,
@@ -424,7 +434,7 @@ impl SchedulerImpl {
         for chunk in chunks_without_partitions {
             if let Some(handle_id) = chunk.get_row().replay_handle_id() {
                 self.meta_store
-                    .update_replay_handle_failed(*handle_id, true)
+                    .update_replay_handle_failed_if_exists(*handle_id, true)
                     .await?;
             }
             ids.push(chunk.get_id());
@@ -503,7 +513,7 @@ impl SchedulerImpl {
         // Using get_tables_with_path due to it's cached
         let tables = self.meta_store.get_tables_with_path(true).await?;
         for table in tables.iter() {
-            if table.table.get_row().is_ready() {
+            if table.table.get_row().is_ready() && !table.table.get_row().sealed() {
                 if let Some(locations) = table.table.get_row().locations() {
                     for location in locations.iter() {
                         if Table::is_stream_location(location) {

rust/cubestore/cubestore/src/store/compaction.rs

@@ -310,10 +310,10 @@ impl CompactionServiceImpl {
         for (failed_chunk, _) in failed {
             if let Some(handle_id) = failed_chunk.get_row().replay_handle_id() {
                 self.meta_store
-                    .update_replay_handle_failed(*handle_id, true)
+                    .update_replay_handle_failed_if_exists(*handle_id, true)
                     .await?;
-                deactivate_failed_chunk_ids.push(failed_chunk.get_id());
             }
+            deactivate_failed_chunk_ids.push(failed_chunk.get_id());
         }
         self.meta_store
             .deactivate_chunks_without_check(deactivate_failed_chunk_ids)
@@ -729,15 +729,21 @@ impl CompactionService for CompactionServiceImpl {
                                     .compaction_in_memory_chunks_max_lifetime_threshold()
                                     as i64
                         })
-                        .unwrap_or(true)
+                        .unwrap_or(false)
             });
 
-        self.compact_chunks_to_memory(mem_chunks, &partition, &index, &table)
-            .await?;
-        self.compact_chunks_to_persistent(persistent_chunks, &partition, &index, &table)
-            .await?;
-        self.deactivate_and_mark_failed_chunks_for_replay(failed)
-            .await?;
+        let deactivate_res = self
+            .deactivate_and_mark_failed_chunks_for_replay(failed)
+            .await;
+        let in_memory_res = self
+            .compact_chunks_to_memory(mem_chunks, &partition, &index, &table)
+            .await;
+        let persistent_res = self
+            .compact_chunks_to_persistent(persistent_chunks, &partition, &index, &table)
+            .await;
+        deactivate_res?;
+        in_memory_res?;
+        persistent_res?;
 
         Ok(())
     }