Pass key to backend as env var (#578)

* Pass key to backend as env var

* fix merge

plane/plane-tests/tests/metrics.rs

@@ -32,7 +32,7 @@ async fn test_get_metrics(_: TestEnvironment) {
     let executor_config =
         ExecutorConfig::from_image_with_defaults("ghcr.io/drifting-in-space/demo-image-drop-four");
     plane_docker
-        .spawn_backend(&backend_name, &container_id, executor_config)
+        .spawn_backend(&backend_name, &container_id, executor_config, None)
         .await
         .unwrap();
 

plane/src/drone/backend_manager.rs

@@ -4,7 +4,7 @@ use super::{
 };
 use crate::{
     names::BackendName,
-    protocol::BackendMetricsMessage,
+    protocol::{AcquiredKey, BackendMetricsMessage},
     typed_socket::TypedSocketSender,
     types::{
         backend_state::TerminationReason, BackendState, ExecutorConfig, PullPolicy, TerminationKind,
@@ -145,6 +145,9 @@ pub struct BackendManager {
 
     /// IP address of the drone.
     ip: IpAddr,
+
+    /// Key acquired by the backend.
+    acquired_key: AcquiredKey,
 }
 
 impl Debug for BackendManager {
@@ -165,6 +168,7 @@ impl BackendManager {
         state_callback: impl Fn(&BackendState) -> Result<(), Box<dyn Error>> + Send + Sync + 'static,
         metrics_sender: Arc<RwLock<TypedSocketSender<BackendMetricsMessage>>>,
         ip: IpAddr,
+        acquired_key: AcquiredKey,
     ) -> Arc<Self> {
         let container_id = ContainerId::from(format!("plane-{}", backend_id));
 
@@ -185,6 +189,7 @@ impl BackendManager {
             state_callback: Box::new(state_callback),
             container_id,
             ip,
+            acquired_key,
         });
 
         manager.set_state(state);
@@ -228,9 +233,15 @@ impl BackendManager {
                 let ip = self.ip;
                 self.metrics_manager.start_gathering_metrics();
 
+                let acquired_key = self.acquired_key.clone();
                 StepStatusResult::future_status(async move {
                     let spawn_result = docker
-                        .spawn_backend(&backend_id, &container_id, executor_config)
+                        .spawn_backend(
+                            &backend_id,
+                            &container_id,
+                            executor_config,
+                            Some(&acquired_key),
+                        )
                         .await;
 
                     let spawn_result = match spawn_result {

plane/src/drone/docker/commands.rs

@@ -1,5 +1,5 @@
 use super::types::ContainerId;
-use crate::{names::BackendName, types::ExecutorConfig};
+use crate::{names::BackendName, protocol::AcquiredKey, types::ExecutorConfig};
 use anyhow::Result;
 use bollard::{
     auth::DockerCredentials,
@@ -112,11 +112,21 @@ pub async fn get_port(docker: &Docker, container_id: &ContainerId) -> Result<u16
 fn get_container_config_from_executor_config(
     backend_id: &BackendName,
     exec_config: ExecutorConfig,
-    runtime: &Option<String>,
+    runtime: Option<&str>,
+    key: Option<&AcquiredKey>,
 ) -> Result<bollard::container::Config<String>> {
     let mut env = exec_config.env;
     env.insert("PORT".to_string(), CONTAINER_PORT.to_string());
-    env.insert("PLANE_BACKEND_ID".to_string(), backend_id.to_string());
+    env.insert("SESSION_BACKEND_ID".to_string(), backend_id.to_string());
+
+    if let Some(key) = key {
+        env.insert(
+            "SESSION_BACKEND_FENCING_TOKEN".to_string(),
+            key.token.to_string(),
+        );
+        env.insert("SESSION_BACKEND_KEY".to_string(), key.key.name.to_string());
+    }
+
     // TODO: set PLANE_LOCK and PLANE_FENCING_TOKEN.
     let env: Vec<String> = env
         .into_iter()
@@ -134,7 +144,7 @@ fn get_container_config_from_executor_config(
         ),
         host_config: Some(HostConfig {
             port_bindings: Some(create_port_bindings()),
-            runtime: runtime.clone(),
+            runtime: runtime.map(|s| s.to_string()),
             memory: exec_config.resource_limits.memory_limit_bytes,
             cpu_period: exec_config
                 .resource_limits
@@ -178,14 +188,16 @@ pub async fn run_container(
     backend_id: &BackendName,
     container_id: &ContainerId,
     exec_config: ExecutorConfig,
-    runtime: &Option<String>,
+    runtime: Option<&str>,
+    acquired_key: Option<&AcquiredKey>,
 ) -> Result<()> {
     let options = bollard::container::CreateContainerOptions {
         name: container_id.to_string(),
         ..Default::default()
     };
 
-    let config = get_container_config_from_executor_config(backend_id, exec_config, runtime)?;
+    let config =
+        get_container_config_from_executor_config(backend_id, exec_config, runtime, acquired_key)?;
 
     docker.create_container(Some(options), config).await?;
 
@@ -301,7 +313,8 @@ mod tests {
         let mut config = get_container_config_from_executor_config(
             &backend_name,
             executor_config.clone(),
-            &None,
+            None,
+            None,
         )
         .unwrap();
         config.cmd = Some(vec!["echo".into(), "hello world".into()]);

plane/src/drone/docker/mod.rs

@@ -2,7 +2,11 @@ use self::{
     commands::{get_port, run_container},
     types::ContainerId,
 };
-use crate::{names::BackendName, protocol::BackendMetricsMessage, types::ExecutorConfig};
+use crate::{
+    names::BackendName,
+    protocol::{AcquiredKey, BackendMetricsMessage},
+    types::ExecutorConfig,
+};
 use anyhow::Result;
 use bollard::{
     auth::DockerCredentials, container::StatsOptions, errors::Error, service::EventMessage,
@@ -115,13 +119,15 @@ impl PlaneDocker {
         backend_id: &BackendName,
         container_id: &ContainerId,
         executable: ExecutorConfig,
+        acquired_key: Option<&AcquiredKey>,
     ) -> Result<SpawnResult> {
         run_container(
             &self.docker,
             backend_id,
             container_id,
             executable,
-            &self.runtime,
+            self.runtime.as_deref(),
+            acquired_key,
         )
         .await?;
         let port = get_port(&self.docker, container_id).await?;

plane/src/drone/executor.rs

@@ -89,7 +89,7 @@ impl Executor {
         action: &BackendAction,
     ) -> Result<()> {
         match action {
-            BackendAction::Spawn { executable, .. } => {
+            BackendAction::Spawn { executable, key } => {
                 let callback = {
                     let state_store = self.state_store.clone();
                     let backend_id = backend_id.clone();
@@ -118,6 +118,7 @@ impl Executor {
                     callback,
                     metrics_sender,
                     self.ip,
+                    key.clone(),
                 );
                 tracing::info!("Inserting backend {}.", backend_id);
                 self.backends.insert(backend_id.clone(), manager);