try fifo query queue

Author: aidan-smith

src/queue.rs

@@ -6,7 +6,7 @@ use crate::task::{
 };
 use dashmap::DashMap;
 use datafusion::physical_plan::ExecutionPlan;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, VecDeque};
 use std::hash::{Hash, Hasher};
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
@@ -30,9 +30,9 @@ impl Copy for TaskId {}
 
 #[derive(Debug)]
 pub struct State {
-    // queue: Mutex<VecDeque<QueryKey>>,
     // The queue used to order queries by executor usage.
-    queue: Mutex<BTreeMap<Duration, u64>>,
+    // queue: Mutex<BTreeMap<Duration, u64>>,
+    queue: Mutex<VecDeque<u64>>,
     start_ts: SystemTime,
 
     query_id_counter: AtomicU64,
@@ -47,7 +47,8 @@ pub struct State {
 impl State {
     pub fn new(notify: Arc<Notify>) -> Self {
         Self {
-            queue: Mutex::new(BTreeMap::new()),
+            // queue: Mutex::new(BTreeMap::new()),
+            queue: Mutex::new(VecDeque::new()),
             start_ts: SystemTime::now(),
             query_id_counter: AtomicU64::new(0),
             table: DashMap::new(),
@@ -67,7 +68,8 @@ impl State {
         query.time = time;
 
         self.table.insert(id, RwLock::new(query));
-        self.queue.lock().await.insert(time, id);
+        // self.queue.lock().await.insert(time, id);
+        self.queue.lock().await.push_back(id);
 
         self.notify.notify_waiters();
         id
@@ -86,7 +88,8 @@ impl State {
     }
 
     pub async fn next_task(&self) -> Option<(TaskId, Arc<dyn ExecutionPlan>)> {
-        let Some((duration, query_id)) = self.queue.lock().await.pop_first() else {
+        // let Some((duration, query_id)) = self.queue.lock().await.pop_first() else {
+        let Some(query_id) = self.queue.lock().await.pop_front() else {
             return None;
         };
         let query = self.table.get(&query_id).unwrap();
@@ -102,7 +105,8 @@ impl State {
             .update_stage_status(task.task_id.stage_id, StageStatus::Running(0))
             .unwrap();
         if let QueryQueueStatus::Available = guard.get_queue_status() {
-            self.queue.lock().await.insert(duration, query_id);
+            // self.queue.lock().await.insert(duration, query_id);
+            self.queue.lock().await.push_back(query_id);
             self.notify.notify_waiters();
         }
 
@@ -129,14 +133,19 @@ impl State {
                 _ => unreachable!(),
             }
 
-            let new_time = guard.time + SystemTime::now().duration_since(ts).unwrap();
             let mut queue = self.queue.lock().await;
-            let _ = queue.remove(&guard.time);
-            if let QueryQueueStatus::Available = guard.get_queue_status() {
-                queue.insert(new_time, task_id.query_id);
-                self.notify.notify_waiters();
+            // let new_time = guard.time + SystemTime::now().duration_since(ts).unwrap();
+            // let _ = queue.remove(&guard.time);
+            // if let QueryQueueStatus::Available = guard.get_queue_status() {
+            //     queue.insert(new_time, task_id.query_id);
+            //     self.notify.notify_waiters();
+            // }
+            // guard.time = new_time;
+            if QueryQueueStatus::Available == guard.get_queue_status()
+                && !queue.contains(&task_id.query_id)
+            {
+                queue.push_back(task_id.query_id);
             }
-            guard.time = new_time;
         }
     }
 
@@ -148,12 +157,18 @@ impl State {
 #[cfg(test)]
 mod tests {
     use rand::Rng;
-    use std::{fs, time::{Duration, SystemTime}};
-    use tokio::{sync::{Mutex, Notify}, time::sleep};
+    use std::{
+        fs,
+        time::{Duration, SystemTime},
+    };
+    use tokio::{
+        sync::{Mutex, Notify},
+        time::sleep,
+    };
 
-    use crate::{parser::ExecutionPlanParser, query_graph::QueryGraph};
     use crate::queue::State;
     use crate::task::TaskStatus;
+    use crate::{parser::ExecutionPlanParser, query_graph::QueryGraph};
     use std::{cmp::min, sync::Arc};
 
     // Deprecated, use test_queue_conc instead
@@ -282,9 +297,7 @@ mod tests {
         // Add a bunch of queries with staggered submission time
         let start_enqueue = SystemTime::now();
         for plan in long_plans {
-            queue
-                .add_query(Arc::clone(&plan))
-                .await;
+            queue.add_query(Arc::clone(&plan)).await;
             sleep(Duration::from_millis(10)).await;
         }
         let enq_time = SystemTime::now().duration_since(start_enqueue).unwrap();