send all requests as a whole to worker

Author: adwinwhite

ui/src/coordinator.rs

@@ -2,17 +2,18 @@ use std::collections::HashMap;
 use std::convert::{TryFrom, TryInto};
 
 use std::str::from_utf8;
+use std::time::SystemTime;
 
-use tokio::sync::{broadcast, mpsc};
-use worker_message::{CoordinatorMessage, ExecuteOutput};
+
+use worker_message::{ExecuteOutput, Job};
 
 use crate::sandbox::{
     self, Channel, CompileRequest, CrateType, Edition, ExecuteRequest, FormatRequest, Mode,
 };
 use crate::{parse_channel, parse_crate_type, parse_edition, parse_mode, parse_target, Error};
 
 struct JobBatch {
-    jobs: Vec<worker_message::Request>,
+    job: Job,
     extra: serde_json::Value
 }
 
@@ -21,14 +22,14 @@ fn split_request(req: WSRequest) -> JobBatch {
         WSRequest::Compile(req) => {
             let (req, extra) = req.try_into().unwrap();
             JobBatch {
-                jobs: compile_request_to_batch(req),
+                job: worker_message::Job { reqs: compile_request_to_batch(req) },
                 extra
             }
         }
         WSRequest::Format(req) => {
             let (req, extra) = req.try_into().unwrap();
             JobBatch {
-                jobs: format_request_to_batch(req),
+                job: worker_message::Job { reqs: format_request_to_batch(req) },
                 extra
             }
         }
@@ -356,17 +357,16 @@ fn worker_response_to_websocket_response(resp: BatchResponse) -> WSResponse {
 // }
 // response
 // }
-
-async fn work(
-    req: worker_message::Request,
-    sender: &mpsc::Sender<CoordinatorMessage>,
-    receiver: &mut broadcast::Receiver<worker_message::Response>,
-) -> worker_message::Response {
-    let coordinator_msg = CoordinatorMessage::Request(0, req);
-    sender.send(coordinator_msg).await.unwrap();
-    receiver.recv().await.unwrap()
+//
+// It's safe unless a leap second happens.
+fn now_as_uid() -> u128 {
+    SystemTime::now()
+        .duration_since(SystemTime::UNIX_EPOCH)
+        .unwrap()
+        .as_nanos()
 }
 
+
 #[derive(serde::Deserialize)]
 #[serde(rename_all = "camelCase")]
 struct WSExecuteRequest {
@@ -566,6 +566,7 @@ enum WSResponse {
 #[cfg(test)]
 mod tests {
 
+    use std::collections::HashMap;
     use std::process::Stdio;
 
     use std::sync::Arc;
@@ -574,17 +575,17 @@ mod tests {
 
     const WORKER_FILEPATH: &str = "../worker-message/target/debug/worker";
     use crate::coordinator::{
-        split_request, work, worker_response_to_websocket_response, BatchResponse,
-        PlaygroundMessage, ResponseKind, WSMessage, WSRequest, WSResponse,
+        split_request, worker_response_to_websocket_response, BatchResponse,
+        PlaygroundMessage, ResponseKind, WSMessage, WSRequest, WSResponse, now_as_uid,
     };
 
     use serde_json::json;
     use tokio::io::{AsyncReadExt, AsyncWriteExt};
     use tokio::process::{ChildStdin, ChildStdout, Command};
-    use tokio::sync::{broadcast, mpsc};
+    use tokio::sync::{mpsc, oneshot};
     use tokio::task::JoinHandle;
     use tokio::time::error::Elapsed;
-    use worker_message::{CoordinatorMessage, WorkerMessage};
+    use worker_message::{CoordinatorMessage, WorkerMessage, JobReport};
 
     use super::{WSCompileRequest, WSFormatRequest};
 
@@ -717,8 +718,9 @@ mod tests {
         worker_sender: mpsc::Sender<CoordinatorMessage>,
         mut worker_receiver: mpsc::Receiver<WorkerMessage>,
     ) {
-        let mut current_task: Option<JoinHandle<()>> = None;
-        let (worker_response_tx, _worker_response_rx) = broadcast::channel(32);
+        let mut worker_response_senders = HashMap::new();
+        let (worker_response_sender_tx, mut worker_response_sender_rx) = mpsc::channel(32);
+        let mut current_job: Option<JoinHandle<()>> = None;
         loop {
             tokio::select! {
                 ws_request = ws_receiver.recv() => {
@@ -730,54 +732,60 @@ mod tests {
                             let ws_msg = serde_json::from_str(&txt).expect("Failed to deserialize websocket message from json string");
                             match ws_msg {
                                 WSMessage::Request(ws_request) => {
-                                    // Abort current task.
+                                    // Abort current job.
                                     // Lower request into low-level operations.
                                     // Execute operations in sequence until a failure is reached.
-                                    if let Some(task) = current_task {
-                                        task.abort();
+                                    if let Some(job) = current_job {
+                                        job.abort();
                                     }
-                                    let mut worker_response_rx = worker_response_tx.subscribe();
                                     let worker_sender = worker_sender.clone();
                                     let ws_sender = ws_sender.clone();
-                                    let kind = ResponseKind::kind(&ws_request);
-                                    let batch = split_request(ws_request);
-                                    let jobs = batch.jobs;
-                                    let extra = batch.extra;
-                                    current_task = Some(tokio::spawn(async move {
-                                        assert!(jobs.len() != 0);
-                                        let mut responses = Vec::with_capacity(jobs.len());
-                                        for job in jobs {
-                                            let resp = work(job, &worker_sender, &mut worker_response_rx).await;
-                                            if !resp.is_ok() {
-                                                responses.push(resp);
-                                                break;
-                                            }
-                                            responses.push(resp);
-                                        }
+                                    let worker_response_sender_tx = worker_response_sender_tx.clone();
+                                    current_job = Some(tokio::spawn(async move {
+                                        let kind = ResponseKind::kind(&ws_request);
+                                        let batch = split_request(ws_request);
+                                        let job = batch.job;
+                                        let extra = batch.extra;
+                                        let id = now_as_uid();
+                                        let (resp_tx, resp_rx) = oneshot::channel();
+                                        worker_response_sender_tx.send((id, resp_tx)).await.unwrap();
+                                        let coordinator_msg = CoordinatorMessage::Request(id, job);
+                                        worker_sender.send(coordinator_msg).await.unwrap();
+                                        let job_report: JobReport = resp_rx.await.unwrap();
+                                        let worker_response = job_report.resps;
                                         let batch_response = BatchResponse {
                                             kind,
-                                            responses,
+                                            responses: worker_response,
                                             extra
                                         };
                                         let ws_response = worker_response_to_websocket_response(batch_response);
                                         ws_sender.send(serde_json::to_string(&PlaygroundMessage::Response(ws_response)).expect("Failed to serialize websocket response")).await.expect("WebSocket failed to send response to user");
                                     }));
+
                                 }
                                 WSMessage::StdinPacket(packet) => {}
                             }
                         }
                     }
                 },
+                worker_response_sender = worker_response_sender_rx.recv() => {
+                    if let Some((uid, resp_tx)) = worker_response_sender {
+                        worker_response_senders.insert(uid, resp_tx);
+                    } else {
+                        break;
+                    }
+                },
                 worker_msg = worker_receiver.recv() => {
                     match worker_msg {
                         None => {
                             break;
                         }
                         Some(msg) => {
                             match msg {
-                                WorkerMessage::Response(uid, resp) => {
-                                    // TODO: drop response that belongs to aborted job.
-                                    worker_response_tx.send(resp).unwrap();
+                                WorkerMessage::Response(id, resp) => {
+                                    if let Some(tx) = worker_response_senders.remove(&id) {
+                                        tx.send(resp).unwrap();
+                                    }
                                 }
                                 WorkerMessage::StdoutPacket(pid, packet) => {}
                                 WorkerMessage::StderrPacket(pid, packet) => {}
@@ -791,12 +799,12 @@ mod tests {
 
     async fn setup_coordinator() -> (mpsc::Sender<String>, mpsc::Receiver<String>) {
         let (worker_sender, worker_receiver) = make_worker_channel().await;
-        let (tx, ws_receiver) = mpsc::channel(32);
-        let (ws_sender, rx) = mpsc::channel(32);
+        let (client_tx, ws_receiver) = mpsc::channel(32);
+        let (ws_sender, client_rx) = mpsc::channel(32);
         tokio::spawn(async move {
             pair_websocket_worker(ws_sender, ws_receiver, worker_sender, worker_receiver).await;
         });
-        (tx, rx)
+        (client_tx, client_rx)
     }
 
     async fn check_websocket_request_response(

worker-message/src/message.rs

@@ -2,20 +2,32 @@ use serde::{Deserialize, Serialize};
 use std::{collections::HashMap, fmt};
 
 pub type Pid = u32;
-pub type RequestId = u64;
+pub type JobId = u128;
 pub type Path = String;
 pub type ResponseError = String;
 pub type Result<T> = std::result::Result<T, ResponseError>;
 
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Job {
+    pub reqs: Vec<Request>
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct JobReport {
+    pub resps: Vec<Response>
+}
+
+
 #[derive(Debug, Serialize, Deserialize)]
 pub enum CoordinatorMessage {
-    Request(RequestId, Request),
+    Request(JobId, Job),
     StdinPacket(Pid, Vec<u8>),
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 pub enum WorkerMessage {
-    Response(RequestId, Response),
+    Response(JobId, JobReport),
     StdoutPacket(Pid, Vec<u8>),
     StderrPacket(Pid, Vec<u8>),
 }