feat(supervisor): spread shared supervisors evenly across gateway replicas

.config/nextest.toml

@@ -36,7 +36,7 @@ test-group = "kubernetes-ha"
 # they kill members or gateways, or scale the pool. Never run them at once.
 # Test names, not binary(), for the reason above.
 [[profile.e2e-kubernetes.overrides]]
-filter = "test(/^(losing_a_supervisor_pod_moves_its_sandboxes_without_restarting_them|scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running|pool_sessions_follow_their_supervisors_gateway)$/)"
+filter = "test(/^(losing_a_supervisor_pod_moves_its_sandboxes_without_restarting_them|scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running|pool_sessions_follow_their_supervisors_gateway|supervisors_spread_evenly_across_gateways)$/)"
 test-group = "kubernetes-shared-supervisor-pool"
 
 # Scaling the pool runs two bounded Helm upgrades (up to 6 minutes each), waits
@@ -51,6 +51,12 @@ slow-timeout = { period = "60s", terminate-after = 20 }
 filter = "test(/^pool_sessions_follow_their_supervisors_gateway$/)"
 slow-timeout = { period = "60s", terminate-after = 20 }
 
+# One bounded rollout per supervisor and one for the gateway, each followed by
+# a settle round.
+[[profile.e2e-kubernetes.overrides]]
+filter = "test(/^supervisors_spread_evenly_across_gateways$/)"
+slow-timeout = { period = "60s", terminate-after = 20 }
+
 # Relative to the profile store dir (`e2e/rust/target/nextest/e2e-kubernetes/`).
 [profile.e2e-kubernetes.junit]
 path = "../../../../../results/e2e-kubernetes.xml"

crates/openshell-server/src/gateway_metrics.rs

@@ -23,6 +23,8 @@ use tonic::{Code, Status};
 
 // Gauges
 pub const SUPERVISOR_SESSIONS: &str = "openshell_server_supervisor_sessions";
+pub const SUPERVISOR_HOSTS: &str = "openshell_server_supervisor_hosts";
+pub const SUPERVISOR_HOST_INSTANCES_MAX: &str = "openshell_server_supervisor_host_instances_max";
 pub const RELAY_PENDING: &str = "openshell_server_relay_pending";
 pub const RELAY_PENDING_CAPACITY: &str = "openshell_server_relay_pending_capacity";
 pub const MUTATION_LOCK_CONNECTIONS_IN_USE: &str =
@@ -262,6 +264,16 @@ pub fn describe_and_initialize(relay: RelayCapacity, mutation_lock_connections:
         Unit::Count,
         "Supervisor control sessions registered on this gateway replica."
     );
+    describe_gauge!(
+        SUPERVISOR_HOSTS,
+        Unit::Count,
+        "Shared supervisors whose assignment streams this gateway replica serves."
+    );
+    describe_gauge!(
+        SUPERVISOR_HOST_INSTANCES_MAX,
+        Unit::Count,
+        "Most sandbox instances placed on one shared supervisor this gateway replica serves."
+    );
     describe_gauge!(
         RELAY_PENDING,
         Unit::Count,
@@ -335,6 +347,8 @@ pub fn describe_and_initialize(relay: RelayCapacity, mutation_lock_connections:
 
     // `increment(0)` registers a series without overwriting a value recorded earlier.
     gauge!(SUPERVISOR_SESSIONS).increment(0.0);
+    gauge!(SUPERVISOR_HOSTS).increment(0.0);
+    gauge!(SUPERVISOR_HOST_INSTANCES_MAX).increment(0.0);
     gauge!(RELAY_PENDING).increment(0.0);
     gauge!(RELAY_PENDING_CAPACITY).set(count_as_f64(relay.per_replica));
     for kind in RelayKind::ALL {
@@ -433,6 +447,11 @@ impl Drop for GaugeSlot {
     }
 }
 
+pub fn record_supervisor_hosts(hosts: usize, most_instances: usize) {
+    gauge!(SUPERVISOR_HOSTS).set(count_as_f64(hosts));
+    gauge!(SUPERVISOR_HOST_INSTANCES_MAX).set(count_as_f64(most_instances));
+}
+
 pub fn record_relay_rejected(reason: RelayRejection) {
     counter!(RELAY_REJECTED_TOTAL, LABEL_REASON => reason.label()).increment(1);
 }
@@ -642,6 +661,8 @@ mod tests {
             ("openshell_server_mutation_lock_connections_capacity", 4),
             ("openshell_server_mutation_lock_connections_in_use", 0),
             ("openshell_server_supervisor_sessions", 0),
+            ("openshell_server_supervisor_hosts", 0),
+            ("openshell_server_supervisor_host_instances_max", 0),
             ("openshell_server_relay_pending", 0),
             (
                 "openshell_server_relay_rejected_total{reason=\"replica_capacity\"}",

crates/openshell-server/src/gateway_ring.rs

@@ -84,6 +84,23 @@ impl GatewayRing {
         Some(self.points[index].1.as_str())
     }
 
+    /// The first replica clockwise from `key` whose `load` is below `cap`,
+    /// or the owner of `key` when every replica is full.
+    pub fn first_below(&self, key: &str, cap: usize, load: impl Fn(&str) -> usize) -> Option<&str> {
+        let owner = self.owner_for(key)?;
+        let start = self
+            .points
+            .partition_point(|(candidate, _)| *candidate < hash_point(key));
+        let mut seen = std::collections::HashSet::new();
+        self.points[start..]
+            .iter()
+            .chain(&self.points[..start])
+            .map(|(_, id)| id.as_str())
+            .filter(|id| seen.insert(*id))
+            .find(|id| load(id) < cap)
+            .or(Some(owner))
+    }
+
     /// The same ring with `replica_id` taken out.
     pub fn without(&self, replica_id: &str) -> Self {
         Self {
@@ -250,4 +267,44 @@ mod tests {
             "adding one replica moved {moved} of {total} sandboxes"
         );
     }
+
+    #[test]
+    fn first_below_skips_full_replicas_and_falls_back_to_the_owner() {
+        let ring = GatewayRing::new(["gw-0", "gw-1"]);
+        let owner = ring.owner_for("sup-0").unwrap();
+        let other = if owner == "gw-0" { "gw-1" } else { "gw-0" };
+
+        assert_eq!(ring.first_below("sup-0", 1, |_| 0), Some(owner));
+        assert_eq!(
+            ring.first_below("sup-0", 1, |id| usize::from(id == owner)),
+            Some(other)
+        );
+        assert_eq!(ring.first_below("sup-0", 1, |_| 1), Some(owner));
+        assert_eq!(GatewayRing::default().first_below("sup-0", 1, |_| 0), None);
+    }
+
+    /// Supervisors arriving one by one under a cap of their even share fill
+    /// every replica to within that cap.
+    #[test]
+    fn first_below_spreads_arrivals_evenly() {
+        for (replicas, supervisors) in [(2, 2), (3, 7), (4, 40)] {
+            let ids: Vec<String> = (0..replicas).map(|i| format!("gw-{i}")).collect();
+            let ring = GatewayRing::new(&ids);
+            let cap = usize::div_ceil(supervisors, replicas);
+            let mut load: HashMap<String, usize> = HashMap::new();
+            for i in 0..supervisors {
+                let home = ring
+                    .first_below(&format!("sup-{i}"), cap, |id| {
+                        load.get(id).copied().unwrap_or(0)
+                    })
+                    .unwrap()
+                    .to_string();
+                *load.entry(home).or_default() += 1;
+            }
+            assert!(
+                load.values().all(|count| *count <= cap) && load.len() == replicas,
+                "{supervisors} supervisors on {replicas} replicas: {load:?}"
+            );
+        }
+    }
 }

crates/openshell-server/src/grpc/supervisor_host.rs

@@ -8,15 +8,21 @@
 //! the handler sees a [`Principal::Supervisor`] and hands the stream to the
 //! controller the in-process compute driver installed.
 
+use std::collections::HashMap;
 use std::pin::Pin;
 use std::sync::Arc;
+use std::sync::atomic::Ordering;
 
+use futures::StreamExt as _;
 use openshell_core::proto::supervisor_host::v1::supervisor_host_server::SupervisorHost;
 use openshell_core::proto::supervisor_host::v1::{HostGatewayMessage, HostMessage};
+use openshell_supervisor_control::MemberView;
+use openshell_supervisor_control::controller::redirect;
 use tonic::{Request, Response, Status, Streaming};
 
 use crate::ServerState;
 use crate::auth::principal::Principal;
+use crate::gateway_ring::GatewayRing;
 
 /// gRPC path prefix of the supervisor host service.
 pub const SUPERVISOR_HOST_PATH_PREFIX: &str = "/openshell.supervisor.host.v1.SupervisorHost/";
@@ -65,10 +71,116 @@ impl SupervisorHost for SupervisorHostService {
         if *shutdown.borrow() {
             return Err(Status::unavailable("gateway is shutting down"));
         }
-        let outbound = controller.attach(identity, request.into_inner()).await?;
-        let outbound = futures::StreamExt::take_until(outbound, async move {
-            let _ = shutdown.wait_for(|shutdown| *shutdown).await;
-        });
+        let state = self.state.clone();
+        let (member_id, outbound) = controller
+            .attach(identity, request.into_inner(), |member_id| {
+                host_home(
+                    &state,
+                    &controller.pool_view().unwrap_or_default().members,
+                    member_id,
+                )
+            })
+            .await?;
+        let handoff = futures::stream::once(async move {
+            let stopping = *state.supervisor_sessions.subscribe_shutdown().borrow();
+            let members = controller.pool_view().unwrap_or_default().members;
+            stopping
+                .then(|| host_home(&state, &members, &member_id))
+                .flatten()
+                .map(|peer_endpoint| Ok(redirect(peer_endpoint)))
+        })
+        .filter_map(std::future::ready);
+        let outbound = outbound
+            .take_until(async move {
+                let _ = shutdown.wait_for(|shutdown| *shutdown).await;
+            })
+            .chain(handoff);
         Ok(Response::new(Box::pin(outbound)))
     }
 }
+
+/// Peer endpoint of the replica that should serve the shared supervisor
+/// `member_id` when that is not this replica. A stopping replica leaves
+/// itself out.
+fn host_home(state: &ServerState, members: &[MemberView], member_id: &str) -> Option<String> {
+    if state.store.is_single_replica() {
+        return None;
+    }
+    let ring = state.gateway_ring.read().ok()?;
+    let ring = if state.gateway_shutting_down.load(Ordering::Acquire) {
+        ring.without(&state.replica_id)
+    } else {
+        ring.clone()
+    };
+    let home = spread_home(&ring, members, member_id)?;
+    if home == state.replica_id {
+        return None;
+    }
+    state
+        .gateway_peers
+        .read()
+        .ok()?
+        .get(home)
+        .filter(|endpoint| !endpoint.is_empty())
+        .cloned()
+}
+
+/// The first replica clockwise from `member_id` on the gateway ring that
+/// serves fewer than an even share of the pool's other supervisors.
+fn spread_home<'a>(
+    ring: &'a GatewayRing,
+    members: &[MemberView],
+    member_id: &str,
+) -> Option<&'a str> {
+    let mut load = HashMap::<&str, usize>::new();
+    let mut supervisors = 1_usize;
+    for member in members
+        .iter()
+        .filter(|member| member.member_id != member_id && !member.leaving)
+    {
+        supervisors += 1;
+        *load.entry(member.gateway_replica.as_str()).or_default() += 1;
+    }
+    let cap = supervisors.div_ceil(ring.replica_count().max(1));
+    ring.first_below(member_id, cap, |replica| {
+        load.get(replica).copied().unwrap_or(0)
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn member(member_id: &str, gateway_replica: &str) -> MemberView {
+        MemberView {
+            member_id: member_id.to_string(),
+            process_id: String::new(),
+            runtime_address: String::new(),
+            leaving: false,
+            gateway_replica: gateway_replica.to_string(),
+        }
+    }
+
+    #[test]
+    fn a_supervisor_goes_to_a_replica_below_its_even_share() {
+        let ring = GatewayRing::new(["gw-a", "gw-b"]);
+        let owner = ring.owner_for("sup-2").unwrap();
+        let other = if owner == "gw-a" { "gw-b" } else { "gw-a" };
+
+        assert_eq!(
+            spread_home(&ring, &[member("sup-1", owner)], "sup-2"),
+            Some(other)
+        );
+        assert_eq!(
+            spread_home(&ring, &[member("sup-1", other)], "sup-2"),
+            Some(owner)
+        );
+        // Its own record and leaving members do not count against a replica.
+        let mut leaving = member("sup-3", owner);
+        leaving.leaving = true;
+        assert_eq!(
+            spread_home(&ring, &[member("sup-2", owner), leaving], "sup-2"),
+            Some(owner)
+        );
+    }
+}

crates/openshell-server/src/supervisor_pool.rs

@@ -182,6 +182,7 @@ pub struct LiveMember {
     pub process_id: String,
     pub runtime_address: String,
     pub leaving: bool,
+    pub gateway_replica: String,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -1168,6 +1169,7 @@ impl SupervisorPool {
                 process_id: payload.process_id,
                 runtime_address: payload.runtime_address,
                 leaving: payload.leaving,
+                gateway_replica: payload.gateway_replica,
             });
         }
         members.sort_by(|left, right| left.member_id.cmp(&right.member_id));
@@ -2671,6 +2673,7 @@ impl SupervisorPool {
                 process_id: payload.process_id,
                 runtime_address: payload.runtime_address,
                 leaving: payload.leaving,
+                gateway_replica: payload.gateway_replica,
             });
         Ok(takeover_cause(owner, member.as_ref()))
     }
@@ -2825,6 +2828,7 @@ impl SupervisorPool {
                     process_id: member.process_id.clone(),
                     runtime_address: member.runtime_address.clone(),
                     leaving: member.leaving,
+                    gateway_replica: member.gateway_replica.clone(),
                 })
                 .collect(),
             owners: BTreeMap::new(),
@@ -2868,6 +2872,19 @@ impl SupervisorPool {
             );
         }
         lock(&self.minted).retain(|term, _| current_terms.contains(term));
+        let mut placed: BTreeMap<&str, usize> = served
+            .iter()
+            .map(|(member, _)| (member.as_str(), 0))
+            .collect();
+        for owner in view.owners.values() {
+            if let Some(count) = placed.get_mut(owner.member_id.as_str()) {
+                *count += 1;
+            }
+        }
+        crate::gateway_metrics::record_supervisor_hosts(
+            placed.len(),
+            placed.values().copied().max().unwrap_or(0),
+        );
         Some(view)
     }
 
@@ -5519,6 +5536,7 @@ mod tests {
             process_id: process_id.to_string(),
             runtime_address: "10.0.0.1:5500".to_string(),
             leaving,
+            gateway_replica: String::new(),
         };
         assert_eq!(takeover_cause(&owner, None), Some("expired"));
         assert_eq!(
@@ -5543,6 +5561,7 @@ mod tests {
             process_id: format!("p-{id}"),
             runtime_address: "10.0.0.1:5500".to_string(),
             leaving,
+            gateway_replica: String::new(),
         };
         let members = vec![member("m1", false), member("m2", false), member("m3", true)];
         let ring = ring_of(&members);

crates/openshell-supervisor-control/README.md

@@ -48,7 +48,10 @@ re-implementing the control contracts.
   `pool_timing::GATEWAY_SILENCE_TIMEOUT` of silence, and drops a member that
   stops reading as soon. `HostAccepted` carries this replica's peer endpoint
   (`set_peer_endpoint`); the host dials its instances' sandbox sessions
-  there, so each session runs on the replica serving its supervisor.
+  there, so each session runs on the replica serving its supervisor. A host
+  whose hello does not set `stay` is sent a `HostRedirect` instead when the
+  gateway places its member on another replica; the gateway picks the first
+  replica on its ring below an even share of supervisors.
 - `pool_timing`: the timing contract between pool hosts and the gateway: the
   15 s membership TTL, the 2 s record refresh, 1 s heartbeats, host silence
   detection and bounded reconnect attempts, and `reconnect_budget`, the worst