Merge branch 'main' into paul/dis-2207-plane-connecting-to-a-session-…

…backend-before-it-is-ready
jamsocket · Jul 9, 2024 · 9326f4d · 9326f4d
2 parents 185cb45 + f10c629
commit 9326f4d
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 28 deletions.
diff --git a/plane/src/drone/backend_manager.rs b/plane/src/drone/backend_manager.rs
@@ -178,7 +178,7 @@ impl<R: Runtime> BackendManager<R> {
                             TerminationReason::StartupTimeout,
                         )
                     } else {
-                        state.to_ready()
+                        state.to_ready(address)
                     }
                 })
             }
@@ -195,17 +195,18 @@ impl<R: Runtime> BackendManager<R> {
                             .terminate(&backend_id, termination == TerminationKind::Hard)
                             .await
                         {
-                            Ok(_) => break,
+                            Ok(false) => return state.to_terminated(None),
+                            Ok(true) => {
+                                // Return a future that never resolves, so that only the container
+                                // terminating bumps us into the next state.
+                                return pending().await;
+                            }
                             Err(err) => {
                                 tracing::error!(?err, "failed to terminate backend");
                                 backoff.wait().await;
                             }
                         }
                     }
-
-                    // Return a future that never resolves, so that only the container
-                    // terminating bumps us into the next state.
-                    pending().await
                 })
             }
             BackendState::Terminated { .. } => StepStatusResult::DoNothing,

diff --git a/plane/src/drone/state_store.rs b/plane/src/drone/state_store.rs
@@ -258,7 +258,7 @@ mod test {
             .register_event(
                 &backend_id,
                 &BackendState::Ready {
-                    address: Some(dummy_addr()),
+                    address: dummy_addr(),
                 },
                 Utc::now(),
             )
@@ -268,7 +268,7 @@ mod test {
         assert_eq!(
             result,
             BackendState::Ready {
-                address: Some(dummy_addr())
+                address: dummy_addr()
             }
         );
     }
@@ -280,7 +280,7 @@ mod test {
         let backend_id = BackendName::new_random();
 
         let ready_state = BackendState::Ready {
-            address: Some(dummy_addr()),
+            address: dummy_addr(),
         };
         {
             state_store
@@ -291,7 +291,7 @@ mod test {
             assert_eq!(
                 result,
                 BackendState::Ready {
-                    address: Some(dummy_addr())
+                    address: dummy_addr()
                 }
             );
         }
@@ -333,7 +333,7 @@ mod test {
         let backend_id = BackendName::new_random();
 
         let ready_state = BackendState::Ready {
-            address: Some(dummy_addr()),
+            address: dummy_addr(),
         };
         state_store
             .register_event(&backend_id, &ready_state, Utc::now())
@@ -348,7 +348,7 @@ mod test {
             assert_eq!(
                 event.state,
                 BackendState::Ready {
-                    address: Some(dummy_addr())
+                    address: dummy_addr()
                 }
             );
         }
@@ -391,7 +391,7 @@ mod test {
         let backend_id = BackendName::new_random();
 
         let ready_state = BackendState::Ready {
-            address: Some(dummy_addr()),
+            address: dummy_addr(),
         };
         state_store
             .register_event(&backend_id, &ready_state, Utc::now())
@@ -418,7 +418,7 @@ mod test {
             assert_eq!(
                 event.state,
                 BackendState::Ready {
-                    address: Some(dummy_addr())
+                    address: dummy_addr()
                 }
             );
         }
@@ -453,7 +453,7 @@ mod test {
             assert_eq!(
                 event.state,
                 BackendState::Ready {
-                    address: Some(dummy_addr())
+                    address: dummy_addr()
                 }
             );
         }

diff --git a/plane/src/typed_unix_socket/mod.rs b/plane/src/typed_unix_socket/mod.rs
@@ -10,7 +10,7 @@ pub struct WrappedMessage<T> {
     /// Optional ID for this message. If it is provided, this message belongs to a request/response pair
     /// (either as the request or the response). If it is not provided, this message is an event.
     id: Option<String>,
-    message: T,
+    pub message: T,
 }
 
 fn get_quick_backoff() -> ExponentialBackoff {

diff --git a/plane/src/types/backend_state.rs b/plane/src/types/backend_state.rs
@@ -90,7 +90,7 @@ pub enum BackendState {
         address: BackendAddr,
     },
     Ready {
-        address: Option<BackendAddr>,
+        address: BackendAddr,
     },
     Terminating {
         last_status: BackendStatus,
@@ -234,7 +234,7 @@ impl BackendState {
     pub fn address(&self) -> Option<BackendAddr> {
         match self {
             BackendState::Waiting { address } => Some(*address),
-            BackendState::Ready { address } => *address,
+            BackendState::Ready { address } => Some(*address),
             _ => None,
         }
     }
@@ -269,16 +269,8 @@ impl BackendState {
         }
     }
 
-    pub fn to_ready(&self) -> BackendState {
-        match self {
-            BackendState::Waiting { address } => BackendState::Ready {
-                address: Some(*address),
-            },
-            _ => {
-                tracing::warn!("to_ready called on non-waiting backend");
-                BackendState::Ready { address: None }
-            }
-        }
+    pub fn to_ready(&self, address: BackendAddr) -> BackendState {
+        BackendState::Ready { address }
     }
 
     pub fn to_terminating(