chore(cubestore): Realtime improvements (cube-js#5064)

* realtime improvements

* fmt

* deactivate empty chunks on repartition

* fix

* fix

* fix

* fix

* fix

* fix

* fixes and changes

* fmt

* in work

* fix

* cleanup

* fix

Author: waralexrom

rust/cubestore/cubestore/src/cluster/mod.rs

@@ -570,10 +570,26 @@ impl Cluster for ClusterImpl {
     }
 
     async fn schedule_repartition(&self, p: &IdRow<Partition>) -> Result<(), CubeError> {
+        let in_memory_node = self.node_name_by_partition(p);
+        let in_memory_job = self
+            .meta_store
+            .add_job(Job::new(
+                RowKey::Table(TableId::Partitions, p.get_id()),
+                JobType::Repartition,
+                in_memory_node.to_string(),
+            ))
+            .await?;
+        if in_memory_job.is_some() {
+            self.notify_job_runner(in_memory_node).await?;
+        }
+
         let chunks = self
             .meta_store
             .get_chunks_by_partition(p.get_id(), false)
-            .await?;
+            .await?
+            .into_iter()
+            .filter(|c| !c.get_row().in_memory())
+            .collect::<Vec<_>>();
 
         for chunk in chunks {
             let node = self.node_name_for_chunk_repartition(&chunk).await?;

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

@@ -300,6 +300,10 @@ pub trait ConfigObj: DIService {
 
     fn compaction_in_memory_chunks_count_threshold(&self) -> usize;
 
+    fn compaction_in_memory_chunks_ratio_threshold(&self) -> u64;
+
+    fn compaction_in_memory_chunks_ratio_check_threshold(&self) -> u64;
+
     fn wal_split_threshold(&self) -> u64;
 
     fn select_worker_pool_size(&self) -> usize;
@@ -316,6 +320,8 @@ pub trait ConfigObj: DIService {
 
     fn not_used_timeout(&self) -> u64;
 
+    fn in_memory_not_used_timeout(&self) -> u64;
+
     fn import_job_timeout(&self) -> u64;
 
     fn meta_store_snapshot_interval(&self) -> u64;
@@ -374,6 +380,8 @@ pub struct ConfigObjImpl {
     pub compaction_in_memory_chunks_size_limit: u64,
     pub compaction_in_memory_chunks_total_size_limit: u64,
     pub compaction_in_memory_chunks_count_threshold: usize,
+    pub compaction_in_memory_chunks_ratio_threshold: u64,
+    pub compaction_in_memory_chunks_ratio_check_threshold: u64,
     pub wal_split_threshold: u64,
     pub data_dir: PathBuf,
     pub dump_dir: Option<PathBuf>,
@@ -386,6 +394,7 @@ pub struct ConfigObjImpl {
     pub query_timeout: u64,
     /// Must be set to 2*query_timeout in prod, only for overrides in tests.
     pub not_used_timeout: u64,
+    pub in_memory_not_used_timeout: u64,
     pub import_job_timeout: u64,
     pub meta_store_log_upload_interval: u64,
     pub meta_store_snapshot_interval: u64,
@@ -449,6 +458,14 @@ impl ConfigObj for ConfigObjImpl {
         self.compaction_in_memory_chunks_count_threshold
     }
 
+    fn compaction_in_memory_chunks_ratio_threshold(&self) -> u64 {
+        self.compaction_in_memory_chunks_ratio_threshold
+    }
+
+    fn compaction_in_memory_chunks_ratio_check_threshold(&self) -> u64 {
+        self.compaction_in_memory_chunks_ratio_check_threshold
+    }
+
     fn wal_split_threshold(&self) -> u64 {
         self.wal_split_threshold
     }
@@ -481,6 +498,10 @@ impl ConfigObj for ConfigObjImpl {
         self.not_used_timeout
     }
 
+    fn in_memory_not_used_timeout(&self) -> u64 {
+        self.in_memory_not_used_timeout
+    }
+
     fn import_job_timeout(&self) -> u64 {
         self.import_job_timeout
     }
@@ -654,6 +675,14 @@ impl Config {
                     "CUBESTORE_IN_MEMORY_CHUNKS_COUNT_THRESHOLD",
                     10,
                 ),
+                compaction_in_memory_chunks_ratio_threshold: env_parse(
+                    "CUBESTORE_IN_MEMORY_CHUNKS_RATIO_THRESHOLD",
+                    3,
+                ),
+                compaction_in_memory_chunks_ratio_check_threshold: env_parse(
+                    "CUBESTORE_IN_MEMORY_CHUNKS_RATIO_CHECK_THRESHOLD",
+                    1000,
+                ),
                 store_provider: {
                     if let Ok(bucket_name) = env::var("CUBESTORE_S3_BUCKET") {
                         FileStoreProvider::S3 {
@@ -695,6 +724,7 @@ impl Config {
                 )),
                 query_timeout,
                 not_used_timeout: 2 * query_timeout,
+                in_memory_not_used_timeout: 30,
                 import_job_timeout: env_parse("CUBESTORE_IMPORT_JOB_TIMEOUT", 600),
                 meta_store_log_upload_interval: 30,
                 meta_store_snapshot_interval: 300,
@@ -755,6 +785,8 @@ impl Config {
                 compaction_in_memory_chunks_size_limit: 262_144 / 4,
                 compaction_in_memory_chunks_total_size_limit: 262_144,
                 compaction_in_memory_chunks_count_threshold: 10,
+                compaction_in_memory_chunks_ratio_threshold: 3,
+                compaction_in_memory_chunks_ratio_check_threshold: 1000,
                 store_provider: FileStoreProvider::Filesystem {
                     remote_dir: Some(
                         env::current_dir()
@@ -769,6 +801,7 @@ impl Config {
                 http_bind_address: None,
                 query_timeout,
                 not_used_timeout: 2 * query_timeout,
+                in_memory_not_used_timeout: 30,
                 import_job_timeout: 600,
                 stale_stream_timeout: 60,
                 select_workers: Vec::new(),

rust/cubestore/cubestore/src/metastore/chunks.rs

@@ -21,6 +21,7 @@ impl Chunk {
             in_memory,
             created_at: Some(Utc::now()),
             oldest_insert_at: Some(Utc::now()),
+            deactivated_at: None,
             suffix: Some(
                 String::from_utf8(thread_rng().sample_iter(&Alphanumeric).take(8).collect())
                     .unwrap()
@@ -79,6 +80,7 @@ impl Chunk {
     pub fn deactivate(&self) -> Chunk {
         let mut to_update = self.clone();
         to_update.active = false;
+        to_update.deactivated_at = Some(Utc::now());
         to_update
     }
 
@@ -102,6 +104,10 @@ impl Chunk {
         &self.oldest_insert_at
     }
 
+    pub fn deactivated_at(&self) -> &Option<DateTime<Utc>> {
+        &self.deactivated_at
+    }
+
     pub fn suffix(&self) -> &Option<String> {
         &self.suffix
     }

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

@@ -693,6 +693,8 @@ pub struct Chunk {
     #[serde(default)]
     oldest_insert_at: Option<DateTime<Utc>>,
     #[serde(default)]
+    deactivated_at: Option<DateTime<Utc>>,
+    #[serde(default)]
     suffix: Option<String>,
     #[serde(default)]
     file_size: Option<u64>
@@ -1047,12 +1049,23 @@ pub trait MetaStore: DIService + Send + Sync {
         include_inactive: bool,
     ) -> Result<Vec<IdRow<Chunk>>, CubeError>;
     async fn chunk_uploaded(&self, chunk_id: u64) -> Result<IdRow<Chunk>, CubeError>;
+    async fn chunk_update_last_inserted(
+        &self,
+        chunk_ids: Vec<u64>,
+        last_inserted_at: Option<DateTime<Utc>>,
+    ) -> Result<(), CubeError>;
     async fn deactivate_chunk(&self, chunk_id: u64) -> Result<(), CubeError>;
+    async fn deactivate_chunks(&self, chunk_ids: Vec<u64>) -> Result<(), CubeError>;
     async fn swap_chunks(
         &self,
         deactivate_ids: Vec<u64>,
         uploaded_ids_and_sizes: Vec<(u64, Option<u64>)>,
     ) -> Result<(), CubeError>;
+    async fn swap_chunks_without_check(
+        &self,
+        deactivate_ids: Vec<u64>,
+        uploaded_ids_and_sizes: Vec<(u64, Option<u64>)>,
+    ) -> Result<(), CubeError>;
     async fn activate_wal(
         &self,
         wal_id_to_delete: u64,
@@ -3442,6 +3455,7 @@ impl MetaStore for RocksMetaStore {
                 vec![(new_chunk, Some(new_chunk_file_size))],
                 db,
                 pipe,
+                false,
             )?;
             Ok(true)
         })
@@ -4120,6 +4134,27 @@ impl MetaStore for RocksMetaStore {
         })
         .await
     }
+    async fn chunk_update_last_inserted(
+        &self,
+        chunk_ids: Vec<u64>,
+        last_inserted_at: Option<DateTime<Utc>>,
+    ) -> Result<(), CubeError> {
+        self.write_operation(move |db_ref, batch_pipe| {
+            let table = ChunkRocksTable::new(db_ref.clone());
+            for chunk_id in chunk_ids {
+                let row = table.get_row_or_not_found(chunk_id)?;
+                table.update(
+                    chunk_id,
+                    row.get_row().set_oldest_insert_at(last_inserted_at),
+                    row.get_row(),
+                    batch_pipe,
+                )?;
+            }
+
+            Ok(())
+        })
+        .await
+    }
 
     async fn deactivate_chunk(&self, chunk_id: u64) -> Result<(), CubeError> {
         self.write_operation(move |db_ref, batch_pipe| {
@@ -4132,6 +4167,16 @@ impl MetaStore for RocksMetaStore {
         })
         .await
     }
+    async fn deactivate_chunks(&self, chunk_ids: Vec<u64>) -> Result<(), CubeError> {
+        self.write_operation(move |db_ref, batch_pipe| {
+            let table = ChunkRocksTable::new(db_ref.clone());
+            for chunk_id in chunk_ids {
+                table.update_with_fn(chunk_id, |row| row.deactivate(), batch_pipe)?;
+            }
+            Ok(())
+        })
+        .await
+    }
 
     async fn activate_wal(
         &self,
@@ -4217,6 +4262,30 @@ impl MetaStore for RocksMetaStore {
                 uploaded_ids_and_sizes,
                 db_ref,
                 batch_pipe,
+                true,
+            )
+        })
+        .await
+    }
+
+    async fn swap_chunks_without_check(
+        &self,
+        deactivate_ids: Vec<u64>,
+        uploaded_ids_and_sizes: Vec<(u64, Option<u64>)>,
+    ) -> Result<(), CubeError> {
+        if uploaded_ids_and_sizes.is_empty() {
+            return Err(CubeError::internal(format!(
+                "Can't swap chunks: {:?} to {:?} empty",
+                deactivate_ids, uploaded_ids_and_sizes
+            )));
+        }
+        self.write_operation(move |db_ref, batch_pipe| {
+            RocksMetaStore::swap_chunks_impl(
+                deactivate_ids,
+                uploaded_ids_and_sizes,
+                db_ref,
+                batch_pipe,
+                false,
             )
         })
         .await
@@ -6288,6 +6357,7 @@ impl RocksMetaStore {
         uploaded_ids_and_sizes: Vec<(u64, Option<u64>)>,
         db_ref: DbTableRef,
         batch_pipe: &mut BatchPipe,
+        check_rows: bool,
     ) -> Result<(), CubeError> {
         trace!(
             "Swapping chunks: deactivating ({}), activating ({})",
@@ -6340,7 +6410,7 @@ impl RocksMetaStore {
                 batch_pipe,
             )?;
         }
-        if deactivate_ids.len() > 0 && activated_row_count != deactivated_row_count {
+        if check_rows && deactivate_ids.len() > 0 && activated_row_count != deactivated_row_count {
             return Err(CubeError::internal(format!(
                 "Deactivated row count ({}) doesn't match activated row count ({}) during swap of ({}) to ({}) chunks",
                 deactivated_row_count,

rust/cubestore/cubestore/src/queryplanner/info_schema/system_chunks.rs

@@ -113,6 +113,46 @@ impl InfoSchemaTableDef for SystemChunksTableDef {
                     ))
                 }),
             ),
+            (
+                Field::new(
+                    "oldest_insert_at",
+                    DataType::Timestamp(TimeUnit::Nanosecond, None),
+                    false,
+                ),
+                Box::new(|chunks| {
+                    Arc::new(TimestampNanosecondArray::from(
+                        chunks
+                            .iter()
+                            .map(|row| {
+                                row.get_row()
+                                    .oldest_insert_at()
+                                    .as_ref()
+                                    .map(|t| t.timestamp_nanos())
+                            })
+                            .collect::<Vec<_>>(),
+                    ))
+                }),
+            ),
+            (
+                Field::new(
+                    "deactivated_at",
+                    DataType::Timestamp(TimeUnit::Nanosecond, None),
+                    false,
+                ),
+                Box::new(|chunks| {
+                    Arc::new(TimestampNanosecondArray::from(
+                        chunks
+                            .iter()
+                            .map(|row| {
+                                row.get_row()
+                                    .deactivated_at()
+                                    .as_ref()
+                                    .map(|t| t.timestamp_nanos())
+                            })
+                            .collect::<Vec<_>>(),
+                    ))
+                }),
+            ),
             (
                 Field::new("file_size", DataType::UInt64, true),
                 Box::new(|chunks| {