|
8 | 8 | //! the handler sees a [`Principal::Supervisor`] and hands the stream to the |
9 | 9 | //! controller the in-process compute driver installed. |
10 | 10 |
|
| 11 | +use std::collections::HashMap; |
11 | 12 | use std::pin::Pin; |
12 | 13 | use std::sync::Arc; |
| 14 | +use std::sync::atomic::Ordering; |
13 | 15 |
|
| 16 | +use futures::StreamExt as _; |
14 | 17 | use openshell_core::proto::supervisor_host::v1::supervisor_host_server::SupervisorHost; |
15 | 18 | use openshell_core::proto::supervisor_host::v1::{HostGatewayMessage, HostMessage}; |
| 19 | +use openshell_supervisor_control::MemberView; |
| 20 | +use openshell_supervisor_control::controller::redirect; |
16 | 21 | use tonic::{Request, Response, Status, Streaming}; |
17 | 22 |
|
18 | 23 | use crate::ServerState; |
19 | 24 | use crate::auth::principal::Principal; |
| 25 | +use crate::gateway_ring::GatewayRing; |
20 | 26 |
|
21 | 27 | /// gRPC path prefix of the supervisor host service. |
22 | 28 | pub const SUPERVISOR_HOST_PATH_PREFIX: &str = "/openshell.supervisor.host.v1.SupervisorHost/"; |
@@ -65,10 +71,116 @@ impl SupervisorHost for SupervisorHostService { |
65 | 71 | if *shutdown.borrow() { |
66 | 72 | return Err(Status::unavailable("gateway is shutting down")); |
67 | 73 | } |
68 | | - let outbound = controller.attach(identity, request.into_inner()).await?; |
69 | | - let outbound = futures::StreamExt::take_until(outbound, async move { |
70 | | - let _ = shutdown.wait_for(|shutdown| *shutdown).await; |
71 | | - }); |
| 74 | + let state = self.state.clone(); |
| 75 | + let (member_id, outbound) = controller |
| 76 | + .attach(identity, request.into_inner(), |member_id| { |
| 77 | + host_home( |
| 78 | + &state, |
| 79 | + &controller.pool_view().unwrap_or_default().members, |
| 80 | + member_id, |
| 81 | + ) |
| 82 | + }) |
| 83 | + .await?; |
| 84 | + let handoff = futures::stream::once(async move { |
| 85 | + let stopping = *state.supervisor_sessions.subscribe_shutdown().borrow(); |
| 86 | + let members = controller.pool_view().unwrap_or_default().members; |
| 87 | + stopping |
| 88 | + .then(|| host_home(&state, &members, &member_id)) |
| 89 | + .flatten() |
| 90 | + .map(|peer_endpoint| Ok(redirect(peer_endpoint))) |
| 91 | + }) |
| 92 | + .filter_map(std::future::ready); |
| 93 | + let outbound = outbound |
| 94 | + .take_until(async move { |
| 95 | + let _ = shutdown.wait_for(|shutdown| *shutdown).await; |
| 96 | + }) |
| 97 | + .chain(handoff); |
72 | 98 | Ok(Response::new(Box::pin(outbound))) |
73 | 99 | } |
74 | 100 | } |
| 101 | + |
| 102 | +/// Peer endpoint of the replica that should serve the shared supervisor |
| 103 | +/// `member_id` when that is not this replica. A stopping replica leaves |
| 104 | +/// itself out. |
| 105 | +fn host_home(state: &ServerState, members: &[MemberView], member_id: &str) -> Option<String> { |
| 106 | + if state.store.is_single_replica() { |
| 107 | + return None; |
| 108 | + } |
| 109 | + let ring = state.gateway_ring.read().ok()?; |
| 110 | + let ring = if state.gateway_shutting_down.load(Ordering::Acquire) { |
| 111 | + ring.without(&state.replica_id) |
| 112 | + } else { |
| 113 | + ring.clone() |
| 114 | + }; |
| 115 | + let home = spread_home(&ring, members, member_id)?; |
| 116 | + if home == state.replica_id { |
| 117 | + return None; |
| 118 | + } |
| 119 | + state |
| 120 | + .gateway_peers |
| 121 | + .read() |
| 122 | + .ok()? |
| 123 | + .get(home) |
| 124 | + .filter(|endpoint| !endpoint.is_empty()) |
| 125 | + .cloned() |
| 126 | +} |
| 127 | + |
| 128 | +/// The first replica clockwise from `member_id` on the gateway ring that |
| 129 | +/// serves fewer than an even share of the pool's other supervisors. |
| 130 | +fn spread_home<'a>( |
| 131 | + ring: &'a GatewayRing, |
| 132 | + members: &[MemberView], |
| 133 | + member_id: &str, |
| 134 | +) -> Option<&'a str> { |
| 135 | + let mut load = HashMap::<&str, usize>::new(); |
| 136 | + let mut supervisors = 1_usize; |
| 137 | + for member in members |
| 138 | + .iter() |
| 139 | + .filter(|member| member.member_id != member_id && !member.leaving) |
| 140 | + { |
| 141 | + supervisors += 1; |
| 142 | + *load.entry(member.gateway_replica.as_str()).or_default() += 1; |
| 143 | + } |
| 144 | + let cap = supervisors.div_ceil(ring.replica_count().max(1)); |
| 145 | + ring.first_below(member_id, cap, |replica| { |
| 146 | + load.get(replica).copied().unwrap_or(0) |
| 147 | + }) |
| 148 | +} |
| 149 | + |
| 150 | +#[cfg(test)] |
| 151 | +mod tests { |
| 152 | + use super::*; |
| 153 | + |
| 154 | + fn member(member_id: &str, gateway_replica: &str) -> MemberView { |
| 155 | + MemberView { |
| 156 | + member_id: member_id.to_string(), |
| 157 | + process_id: String::new(), |
| 158 | + runtime_address: String::new(), |
| 159 | + leaving: false, |
| 160 | + gateway_replica: gateway_replica.to_string(), |
| 161 | + } |
| 162 | + } |
| 163 | + |
| 164 | + #[test] |
| 165 | + fn a_supervisor_goes_to_a_replica_below_its_even_share() { |
| 166 | + let ring = GatewayRing::new(["gw-a", "gw-b"]); |
| 167 | + let owner = ring.owner_for("sup-2").unwrap(); |
| 168 | + let other = if owner == "gw-a" { "gw-b" } else { "gw-a" }; |
| 169 | + |
| 170 | + assert_eq!( |
| 171 | + spread_home(&ring, &[member("sup-1", owner)], "sup-2"), |
| 172 | + Some(other) |
| 173 | + ); |
| 174 | + assert_eq!( |
| 175 | + spread_home(&ring, &[member("sup-1", other)], "sup-2"), |
| 176 | + Some(owner) |
| 177 | + ); |
| 178 | + // Its own record and leaving members do not count against a replica. |
| 179 | + let mut leaving = member("sup-3", owner); |
| 180 | + leaving.leaving = true; |
| 181 | + assert_eq!( |
| 182 | + spread_home(&ring, &[member("sup-2", owner), leaving], "sup-2"), |
| 183 | + Some(owner) |
| 184 | + ); |
| 185 | + } |
| 186 | +} |
0 commit comments