clenaer

Author: shepmaster

orchestrator/src/coordinator.rs

@@ -26,7 +26,7 @@ use crate::{
         CoordinatorMessage, JobId, Multiplexed, OneToOneResponse, ReadFileRequest,
         ReadFileResponse, SerializedError, WorkerMessage, WriteFileRequest,
     },
-    sandbox::{CompileRequest, CompileResponse, CompileResponseWithOutput, Channel},
+    sandbox::{Channel, CompileRequest, CompileResponse, CompileResponseWithOutput},
     DropErrorDetailsExt,
 };
 
@@ -37,14 +37,34 @@ enum DemultiplexCommand {
 }
 
 #[derive(Debug)]
-pub struct Container {
-    stable: ContainerCore,
-    beta: ContainerCore,
-    nightly: ContainerCore,
+pub struct Coordinator<B> {
+    backend: B,
+    // Consider making these lazily-created and/or idly time out
+    stable: Container,
+    beta: Container,
+    nightly: Container,
     token: CancellationToken,
 }
 
-impl Container {
+impl<B> Coordinator<B>
+where
+    B: Backend,
+{
+    pub fn new(backend: B) -> Result<Self, Error> {
+        let token = CancellationToken::new();
+
+        let [stable, beta, nightly] =
+            Channel::ALL.map(|channel| Container::new(channel, token.clone(), &backend));
+
+        Ok(Self {
+            backend,
+            stable: stable?,
+            beta: beta?,
+            nightly: nightly?,
+            token,
+        })
+    }
+
     pub async fn compile(
         &self,
         request: CompileRequest,
@@ -56,11 +76,14 @@ impl Container {
         &self,
         request: CompileRequest,
     ) -> Result<ActiveCompilation, CompileError> {
-        self.select_channel(request.channel).begin_compile(request).await
+        self.select_channel(request.channel)
+            .begin_compile(request)
+            .await
     }
 
-    pub async fn shutdown(self) -> Result<()> {
+    pub async fn shutdown(self) -> Result<B> {
         let Self {
+            backend,
             stable,
             beta,
             nightly,
@@ -70,16 +93,14 @@ impl Container {
 
         let (stable, beta, nightly) = join!(stable.shutdown(), beta.shutdown(), nightly.shutdown());
 
-        stable.unwrap();
+        stable.unwrap(); // ContainerTaskPanickedSnafu
         beta.unwrap();
         nightly.unwrap();
 
-        Ok(())
-//        task.await.context(ContainerTaskPanickedSnafu)?
+        Ok(backend)
     }
 
-
-    fn select_channel(&self, channel: Channel) -> &ContainerCore {
+    fn select_channel(&self, channel: Channel) -> &Container {
         match channel {
             Channel::Stable => &self.stable,
             Channel::Beta => &self.beta,
@@ -88,13 +109,19 @@ impl Container {
     }
 }
 
+impl Coordinator<DockerBackend> {
+    pub fn new_docker() -> Result<Self, Error> {
+        Self::new(DockerBackend(()))
+    }
+}
+
 #[derive(Debug)]
-pub struct ContainerCore {
+pub struct Container {
     task: JoinHandle<Result<()>>,
     commander: Commander,
 }
 
-impl ContainerCore {
+impl Container {
     fn new(channel: Channel, token: CancellationToken, backend: &impl Backend) -> Result<Self> {
         let (mut child, stdin, stdout) = backend.run_worker_in_background(channel)?;
         let IoQueue {
@@ -124,11 +151,7 @@ impl ContainerCore {
             id: Default::default(),
         };
 
-        Ok(ContainerCore {
-            task,
-            commander,
-        })
-
+        Ok(Container { task, commander })
     }
 
     pub async fn compile(
@@ -234,10 +257,7 @@ impl ContainerCore {
     }
 
     pub async fn shutdown(self) -> Result<()> {
-        let Self {
-            task,
-            commander,
-        } = self;
+        let Self { task, commander } = self;
         drop(commander);
         task.await.context(ContainerTaskPanickedSnafu)?
     }
@@ -495,40 +515,11 @@ pub enum CommanderError {
     WorkerOperationFailed { text: String },
 }
 
-#[derive(Debug)]
-pub struct Coordinator<B> {
-    backend: B,
-}
-
-impl<B: Backend> Coordinator<B> {
-    pub fn new(backend: B) -> Self {
-        Self { backend }
-    }
-
-    pub fn allocate(&mut self) -> Result<Container, Error> {
-        let token = CancellationToken::new();
-
-        let [stable, beta, nightly] = Channel::ALL.map(|channel| {
-            ContainerCore::new(channel, token.clone(), &self.backend)
-        });
-
-        Ok(Container {
-            stable: stable?,
-            beta: beta?,
-            nightly: nightly?,
-            token,
-        })
-    }
-}
-
-impl Coordinator<DockerBackend> {
-    pub fn new_docker() -> Self {
-        Self::new(DockerBackend(()))
-    }
-}
-
 pub trait Backend {
-    fn run_worker_in_background(&self, channel: Channel) -> Result<(Child, ChildStdin, ChildStdout)> {
+    fn run_worker_in_background(
+        &self,
+        channel: Channel,
+    ) -> Result<(Child, ChildStdin, ChildStdout)> {
         let mut child = self
             .prepare_worker_command(channel)
             .stdin(Stdio::piped())
@@ -544,6 +535,15 @@ pub trait Backend {
     fn prepare_worker_command(&self, channel: Channel) -> Command;
 }
 
+impl<B> Backend for &B
+where
+    B: Backend,
+{
+    fn prepare_worker_command(&self, channel: Channel) -> Command {
+        B::prepare_worker_command(self, channel)
+    }
+}
+
 macro_rules! docker_command {
     ($($arg:expr),* $(,)?) => ({
         let mut cmd = Command::new("docker");
@@ -721,6 +721,7 @@ mod tests {
 
     use super::*;
 
+    #[derive(Debug)]
     struct TestBackend {
         project_dir: TempDir,
     }
@@ -754,10 +755,7 @@ mod tests {
 
     impl Backend for TestBackend {
         fn prepare_worker_command(&self, channel: Channel) -> Command {
-            let toolchain_file = format!(r#"
-                [toolchain]
-                channel = "{}"
-            "#, channel.to_str());
+            let toolchain_file = format!(r#"[toolchain]\nchannel = "{}""#, channel.to_str());
             let path = self.project_dir.path().join("rust-toolchain.toml");
             std::fs::write(path, toolchain_file).expect("Couldn't write toolchain file");
 
@@ -767,7 +765,7 @@ mod tests {
         }
     }
 
-    fn new_coordinator() -> Coordinator<impl Backend> {
+    fn new_coordinator() -> Result<Coordinator<impl Backend>> {
         Coordinator::new(TestBackend::new())
         // Coordinator::new_docker()
     }
@@ -788,35 +786,33 @@ mod tests {
     #[tokio::test]
     #[snafu::report]
     async fn test_compile_response() -> Result<()> {
-        let mut coordinator = new_coordinator();
+        let coordinator = new_coordinator()?;
 
-        let container = coordinator.allocate()?;
         let response = tokio::time::timeout(
             Duration::from_millis(5000),
-            container.compile(new_compile_request()),
+            coordinator.compile(new_compile_request()),
         )
         .await
         .expect("Failed to receive streaming from container in time")
         .unwrap();
 
         assert!(response.success);
 
-        container.shutdown().await?;
+        coordinator.shutdown().await?;
 
         Ok(())
     }
 
     #[tokio::test]
     #[snafu::report]
     async fn test_compile_streaming() -> Result<()> {
-        let mut coordinator = new_coordinator();
+        let coordinator = new_coordinator()?;
 
-        let container = coordinator.allocate()?;
         let ActiveCompilation {
             task,
             stdout_rx,
             stderr_rx,
-        } = container
+        } = coordinator
             .begin_compile(new_compile_request())
             .await
             .unwrap();
@@ -841,7 +837,7 @@ mod tests {
         );
         assert!(stderr.contains("Finished"), "Missing `Finished`: {stderr}");
 
-        container.shutdown().await?;
+        coordinator.shutdown().await?;
 
         Ok(())
     }

orchestrator/src/sandbox.rs

@@ -38,6 +38,7 @@ pub enum Channel {
 impl Channel {
     pub(crate) const ALL: [Self; 3] = [Self::Stable, Self::Beta, Self::Nightly];
 
+    #[cfg(test)]
     pub(crate) fn to_str(&self) -> &'static str {
         match self {
             Channel::Stable => "stable",