From 124a42edac766db49d4fe63be79d5502acab64a9 Mon Sep 17 00:00:00 2001
From: Paul Butler <paul@driftingin.space>
Date: Wed, 17 Jan 2024 16:43:12 -0500
Subject: [PATCH] Fix bug spawning with the same key (#556)

* Fix bug spawning with the same key

* Add key reuse test
---
 plane/plane-tests/Cargo.toml         | 22 ++++----
 plane/plane-tests/tests/reuse_key.rs | 82 ++++++++++++++++++++++++++++
 plane/src/database/backend_key.rs    |  2 +-
 3 files changed, 94 insertions(+), 12 deletions(-)
 create mode 100644 plane/plane-tests/tests/reuse_key.rs

diff --git a/plane/plane-tests/Cargo.toml b/plane/plane-tests/Cargo.toml
index d32c2467..03710696 100644
--- a/plane/plane-tests/Cargo.toml
+++ b/plane/plane-tests/Cargo.toml
@@ -4,19 +4,19 @@ version = "0.2.0"
 edition = "2021"
 
 [dependencies]
-plane = { path = "../plane-dynamic", package = "plane-dynamic" }
-chrono = { version = "0.4.31", features = ["serde"] }
-tracing = "0.1.40"
-tracing-appender = "0.2.2"
-tokio = { version = "1.33.0", features = ["macros", "rt-multi-thread", "signal"] }
+anyhow = "1.0.75"
 async-trait = "0.1.74"
-tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
-serde_json = "1.0.107"
 bollard = "0.15.0"
+chrono = { version = "0.4.31", features = ["serde"] }
+futures-util = "0.3.29"
+hyper = { version = "0.14.27", features = ["server"] }
+plane = { path = "../plane-dynamic", package = "plane-dynamic" }
+plane-test-macro = { path = "plane-test-macro" }
 reqwest = { version = "0.11.22", features = ["json", "rustls-tls"], default-features = false }
+serde_json = "1.0.107"
 thiserror = "1.0.50"
-hyper = { version = "0.14.27", features = ["server"] }
-anyhow = "1.0.75"
+tokio = { version = "1.33.0", features = ["macros", "rt-multi-thread", "signal"] }
+tracing = "0.1.40"
+tracing-appender = "0.2.2"
+tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
 url = "2.4.1"
-futures-util = "0.3.29"
-plane-test-macro = { path = "plane-test-macro" }
diff --git a/plane/plane-tests/tests/reuse_key.rs b/plane/plane-tests/tests/reuse_key.rs
new file mode 100644
index 00000000..b55c9ca8
--- /dev/null
+++ b/plane/plane-tests/tests/reuse_key.rs
@@ -0,0 +1,82 @@
+use crate::common::timeout::WithTimeout;
+use common::test_env::TestEnvironment;
+use plane::{
+    types::{BackendStatus, ConnectRequest, ExecutorConfig, PullPolicy, SpawnConfig},
+    types::{KeyConfig, ResourceLimits},
+};
+use plane_test_macro::plane_test;
+use serde_json::Map;
+use std::collections::HashMap;
+
+mod common;
+
+#[plane_test]
+async fn reuse_key(env: TestEnvironment) {
+    let controller = env.controller().await;
+    let client = controller.client();
+    let _drone = env.drone(&controller).await;
+
+    // Wait for the drone to register. TODO: this seems long.
+    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
+
+    tracing::info!("Requesting backend.");
+    let connect_request = ConnectRequest {
+        spawn_config: Some(SpawnConfig {
+            executable: ExecutorConfig {
+                image: "ghcr.io/drifting-in-space/demo-image-drop-four".to_string(),
+                pull_policy: PullPolicy::IfNotPresent,
+                env: HashMap::default(),
+                resource_limits: ResourceLimits::default(),
+                credentials: None,
+            },
+            lifetime_limit_seconds: Some(5),
+            max_idle_seconds: None,
+        }),
+        key: Some(KeyConfig {
+            name: "reuse-key".to_string(),
+            namespace: "".to_string(),
+            tag: "".to_string(),
+        }),
+        user: None,
+        auth: Map::default(),
+    };
+
+    let response = client
+        .connect(&env.cluster, &connect_request)
+        .await
+        .unwrap();
+    tracing::info!("Got response.");
+
+    assert!(response.spawned);
+
+    let backend_id = response.backend_id.clone();
+
+    let mut backend_status_stream = client
+        .backend_status_stream(&env.cluster, &backend_id)
+        .with_timeout(10)
+        .await
+        .unwrap()
+        .unwrap();
+
+    let response2 = client
+        .connect(&env.cluster, &connect_request)
+        .await
+        .unwrap();
+
+    assert!(!response2.spawned);
+    assert_eq!(response2.backend_id, backend_id);
+
+    loop {
+        let message = backend_status_stream
+            .next()
+            .with_timeout(10)
+            .await
+            .unwrap()
+            .unwrap();
+
+        tracing::info!("Got status: {:?}", message);
+        if message.status == BackendStatus::Terminated {
+            break;
+        }
+    }
+}
diff --git a/plane/src/database/backend_key.rs b/plane/src/database/backend_key.rs
index 8e91ce12..80b60e5d 100644
--- a/plane/src/database/backend_key.rs
+++ b/plane/src/database/backend_key.rs
@@ -134,6 +134,6 @@ pub struct BackendKeyResult {
 
 impl BackendKeyResult {
     pub fn is_live(&self) -> bool {
-        self.as_of > self.expires_at
+        self.as_of < self.expires_at
     }
 }