From a0e50c66c64e00d68323ca8398f557fda5457b85 Mon Sep 17 00:00:00 2001
From: Paul Butler <paul@driftingin.space>
Date: Thu, 19 Oct 2023 09:58:31 -0400
Subject: [PATCH] Pass PORT environment variable to docker backend (#464)

---
 core/src/messages/agent.rs            |  8 ++++++++
 drone/src/agent/engines/docker/mod.rs | 20 +++++++-------------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/core/src/messages/agent.rs b/core/src/messages/agent.rs
index b6a929b0a..02a5ca72c 100644
--- a/core/src/messages/agent.rs
+++ b/core/src/messages/agent.rs
@@ -20,6 +20,8 @@ use std::{
     time::Duration,
 };
 
+const DEFAULT_CONTAINER_PORT: u16 = 8080;
+
 #[derive(PartialEq, Eq, Clone, Serialize, Deserialize, Debug)]
 pub enum DockerCredentials {
     UsernamePassword { username: String, password: String },
@@ -320,6 +322,12 @@ pub struct DockerExecutableConfig {
     pub volume_mounts: Vec<Value>,
 }
 
+impl DockerExecutableConfig {
+    pub fn port(&self) -> u16 {
+        self.port.unwrap_or(DEFAULT_CONTAINER_PORT)
+    }
+}
+
 #[serde_as]
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, TypedMessage)]
 #[typed_message(subject = "cluster.#cluster.drone.#drone_id.spawn", response = "bool")]
diff --git a/drone/src/agent/engines/docker/mod.rs b/drone/src/agent/engines/docker/mod.rs
index 10fa0065a..18b6c92ee 100644
--- a/drone/src/agent/engines/docker/mod.rs
+++ b/drone/src/agent/engines/docker/mod.rs
@@ -36,7 +36,6 @@ use std::{net::SocketAddr, pin::Pin};
 use tokio_stream::{wrappers::IntervalStream, Stream, StreamExt};
 
 /// The port in the container which is exposed.
-const DEFAULT_CONTAINER_PORT: u16 = 8080;
 const DOCKER_TIMEOUT_SECONDS: u64 = 30;
 /// Interval between reporting stats of a running backend.
 /// NOTE: the minimum possible interval is 1 second.
@@ -172,11 +171,12 @@ impl DockerInterface {
         name: &str,
         executable_config: &DockerExecutableConfig,
     ) -> Result<()> {
-        let env: Vec<String> = executable_config
-            .env
-            .iter()
-            .map(|(k, v)| format!("{}={}", k, v))
-            .collect();
+        let mut env = executable_config.env.clone();
+
+        env.entry("PORT".to_string())
+            .or_insert_with(|| executable_config.port().to_string());
+
+        let env: Vec<String> = env.iter().map(|(k, v)| format!("{}={}", k, v)).collect();
 
         let device_requests = if self.gpu {
             Some(vec![DeviceRequest {
@@ -379,13 +379,7 @@ impl Engine for DockerInterface {
 
         if running {
             let ip = get_ip_of_container(&container)?;
-            let addr = SocketAddr::new(
-                ip,
-                spawn_request
-                    .executable
-                    .port
-                    .unwrap_or(DEFAULT_CONTAINER_PORT),
-            );
+            let addr = SocketAddr::new(ip, spawn_request.executable.port());
 
             Ok(EngineBackendStatus::Running { addr })
         } else {