Initial working 'State'

aidan-smith · aidan-smith · commit dc1bce079bd7 · 2024-05-02T04:25:43.000-04:00
diff --git a/src/query_graph.rs b/src/query_graph.rs
@@ -3,7 +3,6 @@ use crate::composable_database::{QueryStatus, TaskId};
 use crate::task::{Task, TaskStatus};
 use crate::task_queue::TaskQueue;
 use datafusion::arrow::datatypes::Schema;
-use datafusion::common::JoinSide;
 use datafusion::physical_plan::aggregates::AggregateExec;
 use datafusion::physical_plan::joins::{
     CrossJoinExec, HashJoinExec, NestedLoopJoinExec, SortMergeJoinExec, SymmetricHashJoinExec,
@@ -15,6 +14,7 @@ use datafusion::physical_plan::{with_new_children_if_necessary, ExecutionPlan};
 use std::collections::HashMap;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
+use std::time::Duration;
 
 // TODO Change to Waiting, Ready, Running(vec[taskid]), Finished(vec[locations?])
 #[derive(Clone, Debug, Default)]
@@ -48,10 +48,11 @@ pub struct QueryGraph {
     tid_counter: AtomicU64, // TODO: add mutex to stages and make elements pointers to avoid copying
     pub stages: Vec<QueryStage>, // Can be a vec since all stages in a query are enumerated from 0.
     task_queue: TaskQueue,  // Ready tasks in this graph
+    pub time: Duration,
 }
 
 impl QueryGraph {
-    pub async fn new(query_id: u64, plan: Arc<dyn ExecutionPlan>) -> Self {
+    pub fn new(query_id: u64, plan: Arc<dyn ExecutionPlan>) -> Self {
         // Build stages.
         let mut builder = GraphBuilder::new();
         let stages = builder.build(plan.clone());
@@ -63,6 +64,7 @@ impl QueryGraph {
             tid_counter: AtomicU64::new(0),
             stages,
             task_queue: TaskQueue::new(),
+            time: Duration::new(0, 0),
         };
 
         // Build tasks for leaf stages.
@@ -121,6 +123,7 @@ impl QueryGraph {
                     let outputs = stage.outputs.clone();
 
                     if outputs.is_empty() {
+                        println!("QueryGraph::update_stage_status: Query {} is done.", self.query_id);
                         self.status = QueryStatus::Done;
                         return Ok(());
                     }
@@ -161,6 +164,10 @@ impl QueryGraph {
     }
 
     // fn build_tasks(&mut self)
+    pub fn get_plan(&self, stage_id: u64) -> Arc<dyn ExecutionPlan> {
+        let plan = self.stages[stage_id as usize].plan.clone();
+        plan
+    }
 }
 
 #[derive(Clone, Debug)]
diff --git a/src/queue.rs b/src/queue.rs
@@ -5,12 +5,15 @@ use crate::task::{
     TaskStatus::{self, *},
 };
 use crate::SchedulerError;
-use std::collections::{BTreeSet, HashMap};
+use dashmap::DashMap;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion_proto::bytes::physical_plan_to_bytes;
+use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::hash::{Hash, Hasher};
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
-use datafusion_proto::bytes::physical_plan_to_bytes;
-use tokio::sync::{Mutex, Notify};
+use tokio::sync::{Mutex, Notify, RwLock};
 
 // Must implement here since generated TaskId does not derive Hash.
 impl Hash for TaskId {
@@ -48,6 +51,114 @@ pub struct Queue {
     avail: Arc<Notify>,
 }
 
+pub struct State {
+    // queue: Mutex<VecDeque<QueryKey>>,
+    queue: Mutex<BTreeMap<Duration, u64>>,
+    start_ts: SystemTime,
+
+    query_id_counter: AtomicU64,
+    table: DashMap<u64, RwLock<QueryGraph>>,
+    running_tasks: DashMap<TaskId, Task>,
+    notify: Arc<Notify>,
+}
+
+impl State {
+    pub fn new(notify: Arc<Notify>) -> Self {
+        Self {
+            queue: Mutex::new(BTreeMap::new()),
+            start_ts: SystemTime::now(),
+            query_id_counter: AtomicU64::new(0),
+            table: DashMap::new(),
+            running_tasks: DashMap::new(),
+            notify,
+        }
+    }
+
+    fn next_query_id(&self) -> u64 {
+        self.query_id_counter.fetch_add(1, Ordering::SeqCst)
+    }
+
+    pub async fn add_query(&self, plan: Arc<dyn ExecutionPlan>) -> u64 {
+        let id = self.next_query_id();
+        let mut query = QueryGraph::new(id, plan);
+        let time = SystemTime::now().duration_since(self.start_ts).unwrap();
+        query.time = time;
+
+        self.table.insert(id, RwLock::new(query));
+        self.queue.lock().await.insert(time, id);
+
+        self.notify.notify_waiters();
+        id
+    }
+
+    pub async fn get_query_status(&self, query_id: u64) -> Option<QueryStatus> {
+        let status = self.table.get(&query_id)?.read().await.status;
+        if status == QueryStatus::Done {
+            self.table.remove(&query_id);
+        }
+        Some(status)
+    }
+
+    pub async fn abort_query(&self, query_id: u64) {
+        todo!()
+    }
+
+    pub async fn next_task(&self) -> Option<(TaskId, Arc<dyn ExecutionPlan>)> {
+        let Some((duration, query_id)) = self.queue.lock().await.pop_first() else {
+            return None;
+        };
+        let query = self.table.get(&query_id).unwrap();
+        let mut guard = query.write().await;
+
+        let mut task = guard.next_task();
+        task.status = Running(SystemTime::now());
+        let task_id = task.task_id;
+        let plan = guard.get_plan(task.task_id.stage_id);
+
+        // Update query to reflect running task. Requeue if more tasks are available.
+        guard
+            .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.notify.notify_waiters();
+        }
+
+        self.running_tasks.insert(task_id, task);
+        Some((task_id, plan))
+    }
+
+    pub async fn report_task(&self, task_id: TaskId, status: TaskStatus) {
+        if let Some((_, task)) = self.running_tasks.remove(&task_id) {
+            println!("Updating {:?} status to {:?}", task_id, status);
+            let TaskStatus::Running(ts) = task.status else {
+                println!("Task removed with status {:?}", task.status);
+                panic!("Task removed but is not running.");
+            };
+            let query = self.table.get(&task_id.query_id).unwrap();
+            let mut guard = query.write().await;
+
+            match status {
+                TaskStatus::Finished => guard
+                    .update_stage_status(task_id.stage_id, StageStatus::Finished(0))
+                    .unwrap(),
+                TaskStatus::Failed => todo!(),
+                TaskStatus::Aborted => todo!(),
+                _ => 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();
+            }
+            guard.time = new_time;
+        }
+    }
+}
+
 // Notify variable is shared with scheduler service to control task dispatch.
 impl Queue {
     pub fn new(avail: Arc<Notify>) -> Self {
diff --git a/src/server.rs b/src/server.rs
@@ -4,24 +4,22 @@ use crate::composable_database::{
     QueryJobStatusArgs, QueryJobStatusRet, QueryStatus, ScheduleQueryArgs, ScheduleQueryRet,
     TaskId,
 };
-use crate::intermediate_results::{get_results, TaskKey};
 use crate::mock_catalog::load_catalog;
-use crate::parser::ExecutionPlanParser;
 use crate::query_graph::{QueryGraph, StageStatus};
 use crate::queue::Queue;
+use crate::queue::State;
+use crate::task::TaskStatus;
 use crate::SchedulerError;
-use datafusion::arrow::util::pretty::print_batches;
 use datafusion::execution::context::SessionContext;
-use datafusion_proto::bytes::physical_plan_from_bytes;
+use datafusion_proto::bytes::{physical_plan_from_bytes, physical_plan_to_bytes};
 use std::fmt;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use tokio::sync::{Mutex, Notify};
-use tokio::time::{sleep, Duration};
-use tonic::transport::Server;
 use tonic::{Request, Response, Status};
 
 pub struct SchedulerService {
+    state: Arc<State>,
     queue: Arc<Mutex<Queue>>,
     ctx: Arc<SessionContext>, // If we support changing the catalog at runtime, this should be a RwLock.
     query_id_counter: AtomicU64,
@@ -30,18 +28,15 @@ pub struct SchedulerService {
 
 impl fmt::Debug for SchedulerService {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(
-            f,
-            "SchedulerService {{ queue: {:?} }}",
-            self.queue,
-        )
+        write!(f, "SchedulerService {{ queue: {:?} }}", self.queue,)
     }
 }
 
 impl SchedulerService {
     pub async fn new(catalog_path: &str) -> Self {
         let avail = Arc::new(Notify::new());
         Self {
+            state: Arc::new(State::new(avail.clone())),
             queue: Arc::new(Mutex::new(Queue::new(Arc::clone(&avail)))),
             ctx: load_catalog(catalog_path).await,
             query_id_counter: AtomicU64::new(0),
@@ -59,19 +54,28 @@ impl SchedulerService {
         task_id_opt: Option<TaskId>,
     ) -> Result<(TaskId, Vec<u8>), SchedulerError> {
         if let Some(task_id) = task_id_opt {
-            let mut queue = self.queue.lock().await;
-            // Remove the current task from the queue.
-            queue.remove_task(task_id, StageStatus::Finished(0)).await;
+            // let mut queue = self.queue.lock().await;
+            // // Remove the current task from the queue.
+            // queue.remove_task(task_id, StageStatus::Finished(0)).await;
+            self.state.report_task(task_id, TaskStatus::Finished).await;
         }
         loop {
-            let mut queue = self.queue.lock().await;
-            if let Some(new_task_id) = queue.next_task().await {
-                let stage = queue
-                    .get_plan_bytes(new_task_id.query_id, new_task_id.stage_id)
-                    .await?;
-                return Ok((new_task_id, stage));
+            // let mut queue = self.queue.lock().await;
+            // if let Some(new_task_id) = queue.next_task().await {
+            //     let stage = queue
+            //         .get_plan_bytes(new_task_id.query_id, new_task_id.stage_id)
+            //         .await?;
+            //     return Ok((new_task_id, stage));
+            // }
+            // drop(queue);
+            if let Some((task_id, plan)) = self.state.next_task().await {
+                let bytes = physical_plan_to_bytes(plan)
+                    .expect("Failed to serialize physical plan")
+                    .to_vec();
+                println!("SchedulerService: Sending task {:?}", task_id);
+                return Ok((task_id, bytes));
             }
-            drop(queue);
+            println!("SchedulerService: Waiting for new tasks.");
             self.avail.notified().await;
         }
     }
@@ -94,13 +98,13 @@ impl SchedulerApi for SchedulerService {
 
         let plan = physical_plan_from_bytes(bytes.as_slice(), &self.ctx)
             .expect("Failed to deserialize physical plan");
-        // println!("schedule_query: received plan {:?}", plan);
 
-        // Build a query graph, store in query table, enqueue new tasks.
-        let qid = self.next_query_id();
-        let query = QueryGraph::new(qid, plan).await;
-        self.queue.lock().await.add_query(qid, Arc::new(Mutex::new(query))).await;
+        let qid = self.state.add_query(plan).await;
 
+        // Build a query graph, store in query table, enqueue new tasks.
+        // let qid = self.next_query_id();
+        // let query = QueryGraph::new(qid, plan);
+        // self.queue.lock().await.add_query(qid, Arc::new(Mutex::new(query))).await;
 
         let response = ScheduleQueryRet { query_id: qid };
         Ok(Response::new(response))
@@ -113,8 +117,14 @@ impl SchedulerApi for SchedulerService {
     ) -> Result<Response<QueryJobStatusRet>, Status> {
         let QueryJobStatusArgs { query_id } = request.into_inner();
 
-        let status = self.queue.lock().await.get_query_status(query_id).await;
+        let status = self
+            .state
+            .get_query_status(query_id)
+            .await
+            .unwrap_or(QueryStatus::NotFound);
+        // let status = self.queue.lock().await.get_query_status(query_id).await;
         if status == QueryStatus::Done {
+            println!("SchedulerService: Query {} is done.", query_id);
             let stage_id = 0;
             // let final_result = get_results(&TaskKey { stage_id, query_id })
             //     .await
@@ -127,8 +137,8 @@ impl SchedulerApi for SchedulerService {
 
             return Ok(Response::new(QueryJobStatusRet {
                 query_status: QueryStatus::Done.into(),
-                stage_id: stage_id,
-                query_id: query_id
+                stage_id,
+                query_id,
             }));
             // ****************** END CHANGES FROM INTEGRATION TESTING****************//
         }
@@ -145,7 +155,8 @@ impl SchedulerApi for SchedulerService {
     ) -> Result<Response<AbortQueryRet>, Status> {
         // TODO: Actually call executor API to abort query.
         let AbortQueryArgs { query_id } = request.into_inner();
-        self.queue.lock().await.abort_query(query_id).await;
+        // self.queue.lock().await.abort_query(query_id).await;
+        self.state.abort_query(query_id).await;
         let response = AbortQueryRet { aborted: true };
         Ok(Response::new(response))
     }
@@ -243,10 +254,10 @@ mod tests {
                 test_file
             );
         }
-        println!(
-            "test_scheduler: queued {} tasks.",
-            scheduler_service.queue.lock().await.size()
-        );
+        // println!(
+        //     "test_scheduler: queued {} tasks.",
+        //     scheduler_service.queue.lock().await.size()
+        // );
 
         // TODO: add concurrent test eventually
         let mut send_task = NotifyTaskStateArgs {
