Route production control through lifecycle ownership

crates/lazaret-host-coordinator/src/bootstrap.rs

@@ -1,6 +1,7 @@
 use std::fmt;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::sync::Mutex;
 
 use lazaret_core::configuration::ResolvedDomainConfiguration;
 use lazaret_core::identity::{BundleDigest, ContainerId, MountIdentity};
@@ -15,7 +16,9 @@ use lazaret_state_anchor::StateAnchorLayout;
 use lazaret_storage_catalog::Catalog;
 
 use crate::control::ControlSessionAdmission;
-use crate::control::{PublishedControlFace, SessionRegistry, StagedControlFace, TrustAnchor};
+use crate::control::{
+    ControlError, PublishedControlFace, SessionRegistry, StagedControlFace, TrustAnchor,
+};
 use crate::domain_store::HostCoordinatorDomainStore;
 use crate::lifecycle::{
     HostCoordinatorLifecycle, LifecycleAction, LifecycleAdapter, SessionAdmissionIdentity,
@@ -63,7 +66,7 @@ pub struct ProductionLifecycleAdapter<'a> {
     registry: Arc<SessionRegistry>,
     trust_anchor: Option<TrustAnchor>,
     staged_control: Option<StagedControlFace>,
-    published_control: Option<PublishedControlFace>,
+    published_control: Option<Arc<PublishedControlFace>>,
     awaiting_runtime: Option<AwaitingAttachedRunsc>,
     runtime: Option<CommittedSandboxRuntime>,
     state_anchor_mount: Option<MountIdentity>,
@@ -169,8 +172,8 @@ impl<'a> ProductionLifecycleAdapter<'a> {
     }
 
     #[must_use]
-    pub fn published_control(&self) -> Option<&PublishedControlFace> {
-        self.published_control.as_ref()
+    pub fn published_control(&self) -> Option<Arc<PublishedControlFace>> {
+        self.published_control.as_ref().map(Arc::clone)
     }
 
     fn require_generation(&self, generation: &Generation) -> Result<(), ProductionLifecycleError> {
@@ -347,13 +350,13 @@ impl<'a> ProductionLifecycleAdapter<'a> {
     }
 
     fn expose_publication(&mut self) -> Result<(), ProductionLifecycleError> {
-        self.published_control = Some(
+        self.published_control = Some(Arc::new(
             self.staged_control
                 .take()
                 .ok_or_else(|| ProductionLifecycleError("staged control face is absent".into()))?
                 .expose()
                 .map_err(failure)?,
-        );
+        ));
         Ok(())
     }
 
@@ -445,6 +448,12 @@ impl LifecycleAdapter for ProductionLifecycleAdapter<'_> {
             LifecycleAction::UnmountStateAnchor => self.unmount_state_anchor(),
             LifecycleAction::ReleasePreparedLowerLease => self.release_prepared_lower(),
             LifecycleAction::GracefulSupervisorExit => self.graceful_exit(),
+            LifecycleAction::BeginTeardown => {
+                if let Some(control) = self.published_control.take() {
+                    control.withdraw();
+                }
+                Ok(())
+            }
             LifecycleAction::FinishTeardown => {
                 self.published_control.take();
                 Ok(())
@@ -457,8 +466,7 @@ impl LifecycleAdapter for ProductionLifecycleAdapter<'_> {
             | LifecycleAction::AdvanceWaitingSession
             | LifecycleAction::RejectConfigurationMismatch
             | LifecycleAction::RegisterSession
-            | LifecycleAction::DisconnectSession
-            | LifecycleAction::BeginTeardown => Ok(()),
+            | LifecycleAction::DisconnectSession => Ok(()),
         }
     }
 
@@ -476,6 +484,37 @@ impl LifecycleAdapter for ProductionLifecycleAdapter<'_> {
     }
 }
 
+/// Serves one production control connection through the exact retained lifecycle owner.
+///
+/// # Errors
+/// Refuses absent publication, poisoned lifecycle ownership, or control/session failure.
+pub fn serve_production_control(
+    lifecycle: &Mutex<ProductionHostCoordinatorLifecycle<'_, '_>>,
+) -> Result<(), ControlError> {
+    let face = lifecycle
+        .lock()
+        .map_err(|_| ControlError::LockPoisoned)?
+        .adapter()
+        .published_control()
+        .ok_or(ControlError::StaleGeneration)?;
+    face.serve_next_coordinated(
+        |admission| {
+            lifecycle
+                .lock()
+                .map_err(|_| ControlError::LockPoisoned)?
+                .admit_session(admission)
+                .map_err(|error| ControlError::Lifecycle(error.to_string()))
+        },
+        |session_id| {
+            lifecycle
+                .lock()
+                .map_err(|_| ControlError::LockPoisoned)?
+                .disconnect_session(session_id)
+                .map_err(|error| ControlError::Lifecycle(error.to_string()))
+        },
+    )
+}
+
 fn path_from_host(bytes: &[u8]) -> &Path {
     use std::os::unix::ffi::OsStrExt;