feat(cubestore): Queue - correct sort over priority (+created) (#6094)

Author: ovr

packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.js

@@ -1,6 +1,6 @@
 import R from 'ramda';
 import { getEnv, getProcessUid } from '@cubejs-backend/shared';
-import { QueueDriverInterface } from '@cubejs-backend/base-driver';
+import { QueueDriverInterface, QueryKey, QueryKeyHash } from '@cubejs-backend/base-driver';
 import { CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver';
 
 import { TimeoutError } from './TimeoutError';
@@ -139,7 +139,7 @@ export class QueryQueue {
   /**
    * Returns stream object which will be used to pipe data from data source.
    *
-   * @param {*} queryKeyHash
+   * @param {QueryKeyHash} queryKeyHash
    */
   getQueryStream(queryKeyHash) {
     if (!this.streams.queued.has(queryKeyHash)) {
@@ -453,6 +453,10 @@ export class QueryQueue {
   async shutdown() {
     if (this.reconcilePromise) {
       await this.reconcilePromise;
+
+      return true;
+    } else {
+      return false;
     }
   }
 
@@ -741,7 +745,7 @@ export class QueryQueue {
    * Processing query specified by the `queryKey`. This method incapsulate most
    * of the logic related with the queues updates, heartbeating, etc.
    *
-   * @param {string} queryKeyHashed
+   * @param {QueryKeyHash} queryKeyHashed
    * @return {Promise<{ result: undefined | Object, error: string | undefined }>}
    */
   async processQuery(queryKeyHashed) {
@@ -971,8 +975,8 @@ export class QueryQueue {
   /**
    * Returns hash sum of the specified `queryKey`.
    *
-   * @param {*} queryKey
-   * @returns {string}
+   * @param {QueryKey} queryKey
+   * @returns {QueryKeyHash}
    */
   redisHash(queryKey) {
     return this.queueDriver.redisHash(queryKey);

packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts

@@ -61,12 +61,14 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions = {}
     });
 
     async function awaitProcessing() {
-      await queue.shutdown();
-      await Promise.all(processMessagePromises);
-      await Promise.all(processCancelPromises);
+      // process query can call reconcileQueue
+      while (await queue.shutdown() || processMessagePromises.length || processCancelPromises.length) {
+        await Promise.all(processMessagePromises);
+        processMessagePromises = [];
 
-      processMessagePromises = [];
-      processCancelPromises = [];
+        await Promise.all(processCancelPromises);
+        processCancelPromises = [];
+      }
     }
 
     afterEach(async () => {
@@ -211,7 +213,7 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions = {}
     });
 
     test('queue hash process persistent flag properly', () => {
-      const query = ['select * from table'];
+      const query: QueryKey = ['select * from table', []];
       const key1 = queue.redisHash(query);
       // @ts-ignore
       query.persistent = false;
@@ -250,7 +252,7 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions = {}
     });
 
     test('removed before reconciled', async () => {
-      const query = ['select * from'];
+      const query: QueryKey = ['select * from', []];
       const key = queue.redisHash(query);
       await queue.processQuery(key);
       const result = await queue.executeInQueue('foo', key, query);

rust/cubestore/cubestore/src/cachestore/cache_rocksstore.rs

@@ -328,7 +328,7 @@ pub trait CacheStore: DIService + Send + Sync {
         &self,
         prefix: String,
         orphaned_timeout: Option<u32>,
-        stalled_timeout: Option<u32>,
+        heartbeat_timeout: Option<u32>,
     ) -> Result<Vec<IdRow<QueueItem>>, CubeError>;
     async fn queue_list(
         &self,
@@ -555,7 +555,7 @@ impl CacheStore for RocksCacheStore {
         &self,
         prefix: String,
         orphaned_timeout: Option<u32>,
-        stalled_timeout: Option<u32>,
+        heartbeat_timeout: Option<u32>,
     ) -> Result<Vec<IdRow<QueueItem>>, CubeError> {
         self.store
             .read_operation(move |db_ref| {
@@ -568,22 +568,26 @@ impl CacheStore for RocksCacheStore {
 
                 let res = items.into_iter().filter(|item| {
                     if item.get_row().get_status() == &QueueItemStatus::Pending {
-                        if let Some(stalled_timeout) = stalled_timeout {
+                        return if let Some(orphaned_timeout) = orphaned_timeout {
                             let elapsed = now - item.get_row().get_created().clone();
-                            if elapsed.num_milliseconds() > stalled_timeout as i64 {
-                                return true;
+                            if elapsed.num_milliseconds() > orphaned_timeout as i64 {
+                                true
+                            } else {
+                                false
                             }
-                        }
+                        } else {
+                            false
+                        };
                     }
 
                     if item.get_row().get_status() == &QueueItemStatus::Active {
-                        if let Some(orphaned_timeout) = orphaned_timeout {
+                        if let Some(heartbeat_timeout) = heartbeat_timeout {
                             let elapsed = if let Some(heartbeat) = item.get_row().get_heartbeat() {
                                 now - heartbeat.clone()
                             } else {
                                 now - item.get_row().get_created().clone()
                             };
-                            if elapsed.num_milliseconds() > orphaned_timeout as i64 {
+                            if elapsed.num_milliseconds() > heartbeat_timeout as i64 {
                                 return true;
                             }
                         }
@@ -619,9 +623,7 @@ impl CacheStore for RocksCacheStore {
                 if priority_sort {
                     Ok(items
                         .into_iter()
-                        .sorted_by(|a, b| {
-                            b.get_row().get_priority().cmp(a.get_row().get_priority())
-                        })
+                        .sorted_by(|a, b| b.row.cmp(&a.row))
                         .collect())
                 } else {
                     Ok(items)
@@ -966,7 +968,7 @@ crate::di_service!(ClusterCacheStoreClient, [CacheStore]);
 
 #[cfg(test)]
 mod tests {
-    use crate::cachestore::*;
+    use super::*;
     use crate::CubeError;
 
     #[tokio::test]

rust/cubestore/cubestore/src/cachestore/lazy.rs

@@ -220,11 +220,11 @@ impl CacheStore for LazyRocksCacheStore {
         &self,
         prefix: String,
         orphaned_timeout: Option<u32>,
-        stalled_timeout: Option<u32>,
+        heartbeat_timeout: Option<u32>,
     ) -> Result<Vec<IdRow<QueueItem>>, CubeError> {
         self.init()
             .await?
-            .queue_to_cancel(prefix, orphaned_timeout, stalled_timeout)
+            .queue_to_cancel(prefix, orphaned_timeout, heartbeat_timeout)
             .await
     }
 

rust/cubestore/cubestore/src/cachestore/queue_item.rs

@@ -6,6 +6,7 @@ use crate::{base_rocks_secondary_index, rocks_table_new, CubeError};
 use chrono::serde::ts_seconds;
 use chrono::{DateTime, Duration, Utc};
 use rocksdb::WriteBatch;
+use std::cmp::Ordering;
 
 use serde::{Deserialize, Deserializer, Serialize};
 
@@ -78,6 +79,22 @@ pub struct QueueItem {
 
 impl RocksEntity for QueueItem {}
 
+impl Ord for QueueItem {
+    fn cmp(&self, other: &Self) -> Ordering {
+        if self.priority == other.priority {
+            other.created.cmp(&self.created)
+        } else {
+            self.priority.cmp(&other.priority)
+        }
+    }
+}
+
+impl PartialOrd for QueueItem {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
 impl QueueItem {
     pub fn new(path: String, value: String, status: QueueItemStatus, priority: i64) -> Self {
         let parts: Vec<&str> = path.rsplitn(2, ":").collect();
@@ -328,3 +345,82 @@ impl RocksSecondaryIndex<QueueItem, QueueItemIndexKey> for QueueItemRocksIndex {
         *self as IndexId
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::CubeError;
+    use itertools::Itertools;
+
+    #[test]
+    fn test_queue_item_sort() -> Result<(), CubeError> {
+        let priority0_1 =
+            QueueItem::new("1".to_string(), "1".to_string(), QueueItemStatus::Active, 0);
+        let priority0_2 =
+            QueueItem::new("2".to_string(), "2".to_string(), QueueItemStatus::Active, 0);
+        let priority0_3 =
+            QueueItem::new("3".to_string(), "3".to_string(), QueueItemStatus::Active, 0);
+        let priority10_4 = QueueItem::new(
+            "4".to_string(),
+            "4".to_string(),
+            QueueItemStatus::Active,
+            10,
+        );
+        let priority0_5 =
+            QueueItem::new("5".to_string(), "5".to_string(), QueueItemStatus::Active, 0);
+        let priority_n5_6 = QueueItem::new(
+            "6".to_string(),
+            "6".to_string(),
+            QueueItemStatus::Active,
+            -5,
+        );
+
+        assert_eq!(
+            vec![
+                priority0_1.clone(),
+                priority0_2.clone(),
+                priority0_3.clone(),
+                priority10_4.clone(),
+                priority_n5_6.clone(),
+                priority0_5.clone()
+            ]
+            .into_iter()
+            .sorted_by(|a, b| b.cmp(&a))
+            .map(|item| item.get_key().clone())
+            .collect::<Vec<String>>(),
+            vec![
+                "4".to_string(),
+                "1".to_string(),
+                "2".to_string(),
+                "3".to_string(),
+                "5".to_string(),
+                "6".to_string()
+            ]
+        );
+
+        assert_eq!(
+            vec![
+                priority10_4,
+                priority0_1,
+                priority0_5,
+                priority0_2,
+                priority0_3,
+                priority_n5_6
+            ]
+            .into_iter()
+            .sorted_by(|a, b| b.cmp(&a))
+            .map(|item| item.get_key().clone())
+            .collect::<Vec<String>>(),
+            vec![
+                "4".to_string(),
+                "1".to_string(),
+                "2".to_string(),
+                "3".to_string(),
+                "5".to_string(),
+                "6".to_string()
+            ]
+        );
+
+        Ok(())
+    }
+}

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

@@ -1268,12 +1268,12 @@ impl SqlService for SqlServiceImpl {
                     }
                     QueueCommand::ToCancel {
                         prefix,
+                        heartbeat_timeout,
                         orphaned_timeout,
-                        stalled_timeout,
                     } => {
                         let rows = self
                             .cachestore
-                            .queue_to_cancel(prefix.value, orphaned_timeout, stalled_timeout)
+                            .queue_to_cancel(prefix.value, orphaned_timeout, heartbeat_timeout)
                             .await?;
 
                         let columns = vec![Column::new("id".to_string(), ColumnType::String, 0)];

rust/cubestore/cubestore/src/sql/parser.rs

@@ -102,8 +102,8 @@ pub enum QueueCommand {
     },
     ToCancel {
         prefix: Ident,
+        heartbeat_timeout: Option<u32>,
         orphaned_timeout: Option<u32>,
-        stalled_timeout: Option<u32>,
     },
     List {
         prefix: Ident,
@@ -380,31 +380,31 @@ impl<'a> CubeStoreParser<'a> {
                 key: self.parser.parse_identifier()?,
             },
             "stalled" => {
-                let stalled_timeout = self.parse_integer("stalled timeout", false)?;
+                let heartbeat_timeout = Some(self.parse_integer("heartbeat timeout", false)?);
 
                 QueueCommand::ToCancel {
                     prefix: self.parser.parse_identifier()?,
                     orphaned_timeout: None,
-                    stalled_timeout: Some(stalled_timeout),
+                    heartbeat_timeout,
                 }
             }
             "orphaned" => {
-                let orphaned_timeout = self.parse_integer("orphaned timeout", false)?;
+                let orphaned_timeout = Some(self.parse_integer("orphaned timeout", false)?);
 
                 QueueCommand::ToCancel {
                     prefix: self.parser.parse_identifier()?,
-                    orphaned_timeout: Some(orphaned_timeout),
-                    stalled_timeout: None,
+                    heartbeat_timeout: None,
+                    orphaned_timeout,
                 }
             }
             "to_cancel" => {
-                let stalled_timeout = self.parse_integer("stalled timeout", false)?;
-                let orphaned_timeout = self.parse_integer("orphaned timeout", false)?;
+                let heartbeat_timeout = Some(self.parse_integer("heartbeat timeout", false)?);
+                let orphaned_timeout = Some(self.parse_integer("orphaned timeout", false)?);
 
                 QueueCommand::ToCancel {
                     prefix: self.parser.parse_identifier()?,
-                    orphaned_timeout: Some(stalled_timeout),
-                    stalled_timeout: Some(orphaned_timeout),
+                    heartbeat_timeout,
+                    orphaned_timeout,
                 }
             }
             "pending" => {

yarn.lock

@@ -10242,14 +10242,14 @@ caniuse-api@^3.0.0:
     lodash.uniq "^4.5.0"
 
 caniuse-lite@^1.0.0, caniuse-lite@^1.0.30000981, caniuse-lite@^1.0.30001032, caniuse-lite@^1.0.30001109, caniuse-lite@^1.0.30001125, caniuse-lite@^1.0.30001286:
-  version "1.0.30001373"
-  resolved "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001373.tgz"
-  integrity sha512-pJYArGHrPp3TUqQzFYRmP/lwJlj8RCbVe3Gd3eJQkAV8SAC6b19XS9BjMvRdvaS8RMkaTN8ZhoHP6S1y8zzwEQ==
+  version "1.0.30001450"
+  resolved "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001450.tgz"
+  integrity sha512-qMBmvmQmFXaSxexkjjfMvD5rnDL0+m+dUMZKoDYsGG8iZN29RuYh9eRoMvKsT6uMAWlyUUGDEQGJJYjzCIO9ew==
 
 caniuse-lite@^1.0.30001400:
-  version "1.0.30001446"
-  resolved "https://registry.yarnpkg.com/caniuse-lite/-/caniuse-lite-1.0.30001446.tgz#6d4ba828ab19f49f9bcd14a8430d30feebf1e0c5"
-  integrity sha512-fEoga4PrImGcwUUGEol/PoFCSBnSkA9drgdkxXkJLsUBOnJ8rs3zDv6ApqYXGQFOyMPsjh79naWhF4DAxbF8rw==
+  version "1.0.30001450"
+  resolved "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001450.tgz"
+  integrity sha512-qMBmvmQmFXaSxexkjjfMvD5rnDL0+m+dUMZKoDYsGG8iZN29RuYh9eRoMvKsT6uMAWlyUUGDEQGJJYjzCIO9ew==
 
 [email protected]:
   version "1.0.0"