feat(cubestore): Max disk space limit per worker (#6085)

Author: waralexrom

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

@@ -442,11 +442,7 @@ impl Cluster for ClusterImpl {
     }
 
     fn node_name_by_partition(&self, p: &IdRow<Partition>) -> String {
-        if let Some(id) = p.get_row().multi_partition_id() {
-            pick_worker_by_ids(self.config_obj.as_ref(), [id]).to_string()
-        } else {
-            pick_worker_by_partitions(self.config_obj.as_ref(), [p]).to_string()
-        }
+        node_name_by_partition(self.config_obj.as_ref(), p)
     }
 
     async fn node_name_for_chunk_repartition(
@@ -1749,6 +1745,13 @@ fn is_self_reference(name: &str) -> bool {
     name.starts_with("@loop:")
 }
 
+pub fn node_name_by_partition<'a>(config: &'a dyn ConfigObj, p: &IdRow<Partition>) -> String {
+    if let Some(id) = p.get_row().multi_partition_id() {
+        pick_worker_by_ids(config, [id]).to_string()
+    } else {
+        pick_worker_by_partitions(config, [p]).to_string()
+    }
+}
 /// Picks a worker by opaque id for any distributing work in a cluster.
 /// Ids usually come from multi-partitions of the metastore.
 pub fn pick_worker_by_ids<'a>(

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

@@ -401,6 +401,9 @@ pub trait ConfigObj: DIService {
     fn metastore_snapshots_lifetime(&self) -> u64;
 
     fn max_disk_space(&self) -> u64;
+    fn max_disk_space_per_worker(&self) -> u64;
+
+    fn disk_space_cache_duration_secs(&self) -> u64;
 }
 
 #[derive(Debug, Clone)]
@@ -459,6 +462,8 @@ pub struct ConfigObjImpl {
     pub minimum_metastore_snapshots_count: u64,
     pub metastore_snapshots_lifetime: u64,
     pub max_disk_space: u64,
+    pub max_disk_space_per_worker: u64,
+    pub disk_space_cache_duration_secs: u64,
 }
 
 crate::di_service!(ConfigObjImpl, [ConfigObj]);
@@ -666,6 +671,14 @@ impl ConfigObj for ConfigObjImpl {
     fn max_disk_space(&self) -> u64 {
         self.max_disk_space
     }
+
+    fn max_disk_space_per_worker(&self) -> u64 {
+        self.max_disk_space_per_worker
+    }
+
+    fn disk_space_cache_duration_secs(&self) -> u64 {
+        self.disk_space_cache_duration_secs
+    }
 }
 
 lazy_static! {
@@ -881,6 +894,11 @@ impl Config {
                     24 * 60 * 60,
                 ),
                 max_disk_space: env_parse("CUBESTORE_MAX_DISK_SPACE_GB", 0) * 1024 * 1024 * 1024,
+                max_disk_space_per_worker: env_parse("CUBESTORE_MAX_DISK_SPACE_PER_WORKER_GB", 0)
+                    * 1024
+                    * 1024
+                    * 1024,
+                disk_space_cache_duration_secs: 300,
             }),
         }
     }
@@ -951,6 +969,8 @@ impl Config {
                 minimum_metastore_snapshots_count: 3,
                 metastore_snapshots_lifetime: 24 * 3600,
                 max_disk_space: 0,
+                max_disk_space_per_worker: 0,
+                disk_space_cache_duration_secs: 0,
             }),
         }
     }

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

@@ -18,6 +18,7 @@ pub use rocks_fs::*;
 pub use rocks_store::*;
 pub use rocks_table::*;
 
+use crate::cluster::node_name_by_partition;
 use async_trait::async_trait;
 use log::info;
 use rocksdb::{MergeOperands, Options, DB};
@@ -81,10 +82,11 @@ use std::str::FromStr;
 use crate::cachestore::{CacheItem, QueueItem, QueueItemStatus, QueueResult, QueueResultAckEvent};
 use crate::remotefs::LocalDirRemoteFs;
 use snapshot_info::SnapshotInfo;
-use std::time::Duration;
+use std::time::{Duration, SystemTime};
 use table::Table;
 use table::{TableRocksIndex, TableRocksTable};
 use tokio::sync::broadcast::Sender;
+use tokio::sync::RwLock;
 use wal::WALRocksTable;
 
 #[macro_export]
@@ -908,7 +910,10 @@ pub trait MetaStore: DIService + Send + Sync {
         partition_id: u64,
         include_inactive: bool,
     ) -> Result<Vec<IdRow<Chunk>>, CubeError>;
-    async fn get_used_disk_space_out_of_queue(&self) -> Result<u64, CubeError>;
+    async fn get_used_disk_space_out_of_queue(
+        &self,
+        node: Option<String>,
+    ) -> Result<u64, CubeError>;
     async fn get_all_partitions_and_chunks_out_of_queue(
         &self,
     ) -> Result<(Vec<IdRow<Partition>>, Vec<IdRow<Chunk>>), CubeError>;
@@ -1134,6 +1139,7 @@ impl RocksStoreDetails for RocksMetaStoreDetails {
 
 pub struct RocksMetaStore {
     store: Arc<RocksStore>,
+    disk_space_cache: Arc<RwLock<Option<(HashMap<String, u64>, SystemTime)>>>,
     upload_loop: Arc<WorkerLoop>,
 }
 
@@ -1155,6 +1161,7 @@ impl RocksMetaStore {
     fn new_from_store(store: Arc<RocksStore>) -> Arc<Self> {
         Arc::new(Self {
             store,
+            disk_space_cache: Arc::new(RwLock::new(None)),
             upload_loop: Arc::new(WorkerLoop::new("Metastore upload")),
         })
     }
@@ -2236,17 +2243,73 @@ impl MetaStore for RocksMetaStore {
         .await
     }
 
-    async fn get_used_disk_space_out_of_queue(&self) -> Result<u64, CubeError> {
-        let (partitions, chunks) = self.get_all_partitions_and_chunks_out_of_queue().await?;
-        let partitions_size: u64 = partitions
-            .into_iter()
-            .map(|p| p.get_row().file_size().unwrap_or(0))
-            .sum();
-        let chunks_size: u64 = chunks
-            .into_iter()
-            .map(|c| c.get_row().file_size().unwrap_or(0))
-            .sum();
-        Ok(partitions_size + chunks_size)
+    async fn get_used_disk_space_out_of_queue(
+        &self,
+        node: Option<String>,
+    ) -> Result<u64, CubeError> {
+        let cached = if let Some((sizes, time)) = self.disk_space_cache.read().await.as_ref() {
+            let cache_duration =
+                Duration::from_secs(self.store.config.disk_space_cache_duration_secs());
+            if time.elapsed()? < cache_duration {
+                Some(sizes.clone())
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+        let sizes_map = if let Some(sizes) = cached {
+            sizes
+        } else {
+            let (partitions, chunks) = self.get_all_partitions_and_chunks_out_of_queue().await?;
+            let mut partitions_map = partitions
+                .into_iter()
+                .map(|p| {
+                    (
+                        p.get_id(),
+                        (
+                            p.get_row().file_size().unwrap_or(0),
+                            node_name_by_partition(self.store.config.as_ref(), &p),
+                        ),
+                    )
+                })
+                .collect::<HashMap<u64, (u64, String)>>();
+            for c in chunks.into_iter() {
+                if let Some((ref mut size, _)) =
+                    partitions_map.get_mut(&c.get_row().get_partition_id())
+                {
+                    *size = c.get_row().file_size().unwrap_or(0);
+                }
+            }
+
+            let workers = if self.store.config.select_workers().is_empty() {
+                vec![self.store.config.server_name().clone()]
+            } else {
+                self.store.config.select_workers().clone()
+            };
+
+            let mut map = workers
+                .into_iter()
+                .map(|n| (n, 0))
+                .collect::<HashMap<String, u64>>();
+
+            for (_, (size, node)) in partitions_map.into_iter() {
+                map.entry(node).and_modify(|s| *s += size).or_insert(0);
+            }
+
+            let mut cache = self.disk_space_cache.write().await;
+            *cache = Some((map.clone(), SystemTime::now()));
+
+            map
+        };
+
+        let res = if let Some(node_name) = node {
+            sizes_map.get(&node_name).unwrap_or(&0).clone()
+        } else {
+            sizes_map.values().sum::<u64>()
+        };
+
+        Ok(res)
     }
 
     async fn get_all_partitions_and_chunks_out_of_queue(

rust/cubestore/cubestore/src/sql/mod.rs

@@ -320,7 +320,7 @@ impl SqlServiceImpl {
 
         let max_disk_space = self.config_obj.max_disk_space();
         if max_disk_space > 0 {
-            let used_space = self.db.get_used_disk_space_out_of_queue().await?;
+            let used_space = self.db.get_used_disk_space_out_of_queue(None).await?;
             if max_disk_space < used_space {
                 return Err(CubeError::user(format!(
                     "Exceeded available storage space: {:.3} GB out of {} GB allowed. Please consider changing pre-aggregations build range, reducing index count or pre-aggregations granularity.",
@@ -3602,18 +3602,105 @@ mod tests {
                 c.partition_split_threshold = 1000000;
                 c.compaction_chunks_count_threshold = 100;
                 c.max_disk_space = 300_000;
-                c.select_workers = vec!["127.0.0.1:24306".to_string()];
-                c.metastore_bind_address = Some("127.0.0.1:25312".to_string());
+                c.select_workers = vec!["127.0.0.1:24308".to_string()];
+                c.metastore_bind_address = Some("127.0.0.1:25314".to_string());
                 c
             })
             .start_test(async move |services| {
                 let service = services.sql_service;
 
                 Config::test("disk_space_limit_worker_1")
                     .update_config(|mut c| {
-                        c.worker_bind_address = Some("127.0.0.1:24306".to_string());
-                        c.server_name = "127.0.0.1:24306".to_string();
-                        c.metastore_remote_address = Some("127.0.0.1:25312".to_string());
+                        c.worker_bind_address = Some("127.0.0.1:24308".to_string());
+                        c.server_name = "127.0.0.1:24308".to_string();
+                        c.max_disk_space = 300_000;
+                        c.metastore_remote_address = Some("127.0.0.1:25314".to_string());
+                        c.store_provider = FileStoreProvider::Filesystem {
+                            remote_dir: Some(env::current_dir()
+                                .unwrap()
+                                .join("disk_space_limit-upstream")),
+                        };
+                        c
+                    })
+                    .start_test_worker(async move |_| {
+                        let paths = {
+                            let dir = env::temp_dir();
+
+                            let path_1 = dir.clone().join("foo-cluster-1.csv");
+                            let path_2 = dir.clone().join("foo-cluster-2.csv.gz");
+                            let mut file = File::create(path_1.clone()).unwrap();
+
+                            file.write_all("id,city,arr,t\n".as_bytes()).unwrap();
+                            for i in 0..100000
+                            {
+                                file.write_all(format!("{},\"New York\",\"[\"\"\"\"]\",2021-01-24 19:12:23.123 UTC\n", i).as_bytes()).unwrap();
+                            }
+
+
+                            let mut file = GzipEncoder::new(BufWriter::new(tokio::fs::File::create(path_2.clone()).await.unwrap()));
+
+                            file.write_all("id,city,arr,t\n".as_bytes()).await.unwrap();
+                            for i in 0..100000
+                            {
+                                file.write_all(format!("{},San Francisco,\"[\"\"Foo\"\",\"\"Bar\"\",\"\"FooBar\"\"]\",\"2021-01-24 12:12:23 UTC\"\n", i).as_bytes()).await.unwrap();
+                            }
+
+                            file.shutdown().await.unwrap();
+
+                            vec![path_1, path_2]
+                        };
+
+                        let _ = service.exec_query("CREATE SCHEMA IF NOT EXISTS Foo").await.unwrap();
+                        let _ = service.exec_query(
+                            &format!(
+                                "CREATE TABLE Foo.Persons (id int, city text, t timestamp, arr text) INDEX persons_city (`city`, `id`) LOCATION {}",
+                                paths.iter().map(|p| format!("'{}'", p.to_string_lossy())).join(",")
+                            )
+                        ).await.unwrap();
+
+                        let res = service.exec_query(
+                            &format!(
+                                "CREATE TABLE Foo.Persons2 (id int, city text, t timestamp, arr text) INDEX persons_city (`city`, `id`) LOCATION {}",
+                                paths.iter().map(|p| format!("'{}'", p.to_string_lossy())).join(",")
+                            )
+                        ).await;
+                        if let Err(err) = res {
+                            assert!(err.message.starts_with("Exceeded available storage space:"));
+                        } else {
+                            assert!(false);
+                        }
+
+                    })
+                    .await;
+            })
+            .await;
+    }
+
+    #[tokio::test]
+    async fn disk_space_limit_per_worker() {
+        Config::test("disk_space_limit_per_worker")
+            .update_config(|mut c| {
+                c.partition_split_threshold = 1000000;
+                c.compaction_chunks_count_threshold = 100;
+                c.max_disk_space_per_worker = 6_000_000;
+                c.select_workers = vec!["127.0.0.1:24309".to_string()];
+                c.metastore_bind_address = Some("127.0.0.1:25315".to_string());
+                c
+            })
+            .start_test(async move |services| {
+                let service = services.sql_service;
+
+                Config::test("disk_space_limit_per_worker_worker_1")
+                    .update_config(|mut c| {
+                        c.worker_bind_address = Some("127.0.0.1:24309".to_string());
+                        c.server_name = "127.0.0.1:24309".to_string();
+                        c.max_disk_space_per_worker = 6_000_000;
+                        c.metastore_remote_address = Some("127.0.0.1:25315".to_string());
+                        c.store_provider = FileStoreProvider::Filesystem {
+                            remote_dir: Some(env::current_dir()
+                                .unwrap()
+                                .join("disk_space_limit_per_worker-upstream")),
+                        };
                         c
                     })
                     .start_test_worker(async move |_| {
@@ -3659,7 +3746,7 @@ mod tests {
                             )
                         ).await;
                         if let Err(err) = res {
-                            assert!(err.message.starts_with("Exceeded available storage "));
+                            assert!(err.message.contains("Exceeded available storage space on worker"));
                         } else {
                             assert!(false);
                         }