Skip to content

Commit 136f6d8

Browse files
committed
refactor(gateway): trim routing-lock comments and redundant tests
1 parent e04528c commit 136f6d8

4 files changed

Lines changed: 17 additions & 79 deletions

File tree

‎dstack/gateway/src/admin_service.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,8 @@ impl AdminRpcHandler {
6161
.kv_store()
6262
.get_best_zt_domain()
6363
.unwrap_or_default();
64-
// Snapshot routing fields only. KV readers can wait behind a sync merge;
65-
// neither those waits nor a cold wg query may hold the routing mutex.
64+
// KV reads can wait behind a sync merge, so do them after releasing
65+
// the routing lock.
6666
let mut hosts = {
6767
let state = self.state.lock();
6868
state
@@ -97,6 +97,8 @@ impl AdminRpcHandler {
9797
nodes: self.state.get_all_nodes(),
9898
hosts,
9999
num_connections: NUM_CONNECTIONS.load(Ordering::Relaxed),
100+
// Reads the post-probe config, so this is what the data path is
101+
// running rather than what the file asked for.
100102
accel: Some(accel_status(&config.proxy)),
101103
health_gating: config.proxy.health_check.enabled,
102104
})

‎dstack/gateway/src/main_service.rs‎

Lines changed: 5 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -291,13 +291,6 @@ impl ProxyInner {
291291
self.state.lock().or_panic("Failed to lock AppState")
292292
}
293293

294-
/// For tests that assert a handler does its slow work off this lock: the
295-
/// only way to observe how long it is held is to compete for it.
296-
#[cfg(test)]
297-
pub(crate) fn routing_lock_is_free(&self) -> bool {
298-
self.state.try_lock().is_ok()
299-
}
300-
301294
/// WireGuard handshake ages, without taking the routing lock.
302295
///
303296
/// The cache is its own synchronization, so a caller that only needs
@@ -317,13 +310,8 @@ impl ProxyInner {
317310

318311
/// Publish this node's WireGuard observations to the KV store.
319312
///
320-
/// The routing lock is taken only to copy the public-key-to-instance-id
321-
/// map, never to write. Each write takes the KV store's own write lock,
322-
/// which a sync round holds for as long as it takes to merge an inbound
323-
/// envelope; composing the two would let one peer's merge stall every
324-
/// proxied connection for the length of it. `Admin.Status` calls this, and
325-
/// `web_routes::route_index` calls `Admin.Status`, so the composition is
326-
/// reachable from a dashboard page load.
313+
/// The KV writes wait on the store's write lock, which a sync merge holds,
314+
/// so they must not run under the routing lock.
327315
pub(crate) fn refresh_state(&self) -> Result<()> {
328316
let handshakes = self.latest_handshakes(None)?;
329317
let instance_ids: Vec<(String, u64)> = {
@@ -2570,11 +2558,8 @@ impl ProxyState {
25702558
Ok(())
25712559
}
25722560

2573-
/// Drop instances the cluster has stopped seeing.
2574-
///
2575-
/// Reads the handshake observations `ProxyInner::refresh_state` publishes,
2576-
/// so the caller refreshes first -- off this lock, because the refresh
2577-
/// writes to the KV store.
2561+
/// Drop instances the cluster has stopped seeing. The caller runs
2562+
/// `ProxyInner::refresh_state` first, off this lock.
25782563
fn recycle(&mut self) -> Result<()> {
25792564
// Note: Gateway nodes are not removed from KvStore, only marked offline/retired
25802565

@@ -2804,4 +2789,4 @@ impl RpcCall<Proxy> for RpcHandler {
28042789
}
28052790

28062791
#[cfg(test)]
2807-
pub(crate) mod tests;
2792+
mod tests;

‎dstack/gateway/src/main_service/handshakes.rs‎

Lines changed: 3 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -51,24 +51,9 @@ impl LatestHandshakesCache {
5151
self.cell.set(timestamps);
5252
}
5353

54-
/// The cached snapshot, never a freshly produced one.
55-
///
56-
/// The routing path reaches this from `select_top_n_hosts` once per
57-
/// proxied connection, holding the `ProxyState` mutex every tenant's
58-
/// traffic takes. Producing a value here means forking `wg show` under
59-
/// that lock -- and the cell only ever fills on success, so a host where
60-
/// `wg show` cannot work (wrong `core.wg.interface`, no WireGuard module,
61-
/// a container without `CAP_NET_ADMIN`) would fork once per inbound
62-
/// connection, for as long as the traffic lasts. The producer runs in
63-
/// [`Self::refresh`] and in the periodic task instead, both off the
64-
/// routing path and on the blocking pool.
65-
///
66-
/// Serving an expired snapshot rather than nothing is deliberate: every
67-
/// consumer compares handshake *age* against its own threshold, so a stale
68-
/// snapshot ages the whole fleet uniformly, while an empty map reads as
69-
/// "no instance has ever handshaked" and would drain routing on the first
70-
/// refresh that fails. Empty is only what a node that has never had a
71-
/// snapshot actually knows.
54+
/// The cached snapshot, stale if need be; never runs `wg show`. Callers
55+
/// hold the routing lock, and a host where `wg show` always fails would
56+
/// otherwise fork once per connection.
7257
pub(crate) fn latest(&self, stale_timeout: Option<Duration>) -> Result<HandshakesWithAge> {
7358
let timestamps = match self.cell.get_allow_stale() {
7459
Ok(snapshot) => snapshot.into_value(),
@@ -146,14 +131,6 @@ fn add_elapsed_time(
146131
mod tests {
147132
use super::*;
148133

149-
/// A cold cache must not fork a process on the routing path.
150-
///
151-
/// `latest` is reached from `select_top_n_hosts` while `ProxyState`'s mutex
152-
/// is held, once per proxied connection. A host where `wg show` can never
153-
/// succeed -- a misconfigured `core.wg.interface`, no WireGuard module, a
154-
/// container without `CAP_NET_ADMIN` -- never fills the cell, so a `latest`
155-
/// that produces its own value forks once per inbound connection while
156-
/// holding the lock every tenant's traffic takes.
157134
#[test]
158135
fn a_cold_cache_does_not_shell_out_on_the_routing_path() {
159136
const CONNECTIONS: usize = 500;
@@ -181,23 +158,6 @@ mod tests {
181158
);
182159
}
183160

184-
/// An expired snapshot still routes.
185-
///
186-
/// Consumers threshold on handshake *age*, so serving a stale map ages the
187-
/// whole fleet uniformly. Returning nothing instead would read as "no
188-
/// instance has ever handshaked" and drain routing on the first refresh
189-
/// that fails.
190-
#[test]
191-
fn an_expired_snapshot_is_still_served() {
192-
let cache = LatestHandshakesCache::new("dstack-no-such-iface0".to_string(), Duration::ZERO);
193-
cache.set_for_test(BTreeMap::from([("pubkey-a".to_string(), 1730190589)]));
194-
195-
assert!(cache
196-
.latest(None)
197-
.expect("an expired snapshot still answers")
198-
.contains_key("pubkey-a"));
199-
}
200-
201161
#[test]
202162
fn parses_latest_handshake_timestamps() {
203163
let handshakes = parse_latest_handshake_timestamps(

‎dstack/gateway/src/main_service/tests.rs‎

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@ use base64::Engine as _;
1212
use std::sync::atomic::Ordering;
1313
use tempfile::TempDir;
1414

15-
pub(crate) struct TestState {
16-
pub(crate) proxy: Proxy,
15+
struct TestState {
16+
proxy: Proxy,
1717
_temp_dir: TempDir,
1818
}
1919

@@ -24,7 +24,7 @@ impl std::ops::Deref for TestState {
2424
}
2525
}
2626

27-
pub(crate) async fn create_test_state() -> TestState {
27+
async fn create_test_state() -> TestState {
2828
create_test_state_with(|_| {}).await
2929
}
3030

@@ -2680,23 +2680,14 @@ fn tombstone_collection_triggers_on_write_count_boundaries_not_on_time() {
26802680
assert!(!tombstone_collection_due(Some(w(100, 0)), w(90, 0), 100));
26812681
}
26822682

2683-
/// `Admin.Status` must not park the data plane on the KV store's write lock.
2684-
///
2685-
/// `refresh_state` publishes this node's handshake observations: one ephemeral
2686-
/// write per instance plus one for the node itself. Those take the KV store's
2687-
/// own write lock, which a sync round holds for as long as it takes to merge
2688-
/// an inbound envelope -- megabytes, from a peer. Running the writes under
2689-
/// `ProxyState` composes the two locks, so one slow merge stalls every proxied
2690-
/// connection, and `web_routes::route_index` reaches this on a dashboard page
2691-
/// load.
2683+
/// A sync merge holding the KV store must not stall routing via `refresh_state`.
26922684
#[tokio::test]
26932685
async fn refreshing_state_leaves_the_routing_lock_free_while_it_writes_to_the_kv_store() {
26942686
use std::sync::atomic::AtomicBool;
26952687

26962688
let state = create_test_state().await;
26972689

2698-
// Stand in for a sync round holding the store: every write in
2699-
// `refresh_state` now blocks.
2690+
// Stand in for a sync merge holding the store.
27002691
let _store_held = state.kv_store().ephemeral().write();
27012692

27022693
let finished = std::sync::Arc::new(AtomicBool::new(false));

0 commit comments

Comments
 (0)