Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 18 additions & 10 deletions dstack/gateway/src/main_service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1100,7 +1100,7 @@ fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()>
}

for (instance_id, data) in instances {
let new_info = InstanceInfo {
let mut new_info = InstanceInfo {
id: instance_id.clone(),
app_id: data.app_id.clone(),
ip: data.ip,
Expand All @@ -1114,24 +1114,32 @@ fn reload_instances_from_kv_store(proxy: &Proxy, store: &KvStore) -> Result<()>
connections: Default::default(),
};

let old_ip = state.state.instances.get(&instance_id).map(|e| e.ip);
if let Some(existing) = state.state.instances.get(&instance_id) {
let existing = state.state.instances.get(&instance_id).cloned();
if let Some(existing) = &existing {
// Check if wg config needs update
if existing.public_key != data.public_key || existing.ip != data.ip {
wg_changed = true;
}
// Only update if remote is newer (based on reg_time)
if data.reg_time <= encode_ts(existing.reg_time) {
continue;
}
// WaveKV has already selected the winning value. Materialize it
// unconditionally instead of applying another LWW rule here.
new_info.connections = existing.connections.clone();
} else {
wg_changed = true;
}

// Release old IP if it changed (prevent IP leak)
if let Some(old_ip) = old_ip {
if old_ip != data.ip {
state.state.allocated_addresses.remove(&old_ip);
if let Some(existing) = &existing {
if existing.ip != data.ip {
state.state.allocated_addresses.remove(&existing.ip);
}
if existing.app_id != data.app_id {
if let Some(app_instances) = state.state.apps.get_mut(&existing.app_id) {
app_instances.remove(&instance_id);
if app_instances.is_empty() {
state.state.apps.remove(&existing.app_id);
state.state.top_n.remove(&existing.app_id);
}
}
}
}
state.state.allocated_addresses.insert(data.ip);
Expand Down
45 changes: 45 additions & 0 deletions dstack/gateway/src/main_service/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use crate::config::{load_config_figment, Config, MutualConfig};
use crate::kv::PortFlags;
use crate::proxy::port_policy::is_port_allowed;
use base64::Engine as _;
use std::sync::atomic::Ordering;
use tempfile::TempDir;

struct TestState {
Expand Down Expand Up @@ -557,6 +558,50 @@ async fn an_instance_deleted_on_another_node_stops_being_routable_here() {
.contains(&"10.0.0.40".parse().unwrap()));
}

#[tokio::test]
async fn proxy_state_adopts_the_wavekv_winner_regardless_of_value_reg_time() {
let state = create_test_state().await;
sync_from_peer_at(
&state,
"contended",
"10.0.0.40",
&test_pubkey("old-key"),
300,
);
reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap();
state
.lock()
.state
.instances
.get("contended")
.unwrap()
.connections
.store(7, Ordering::Relaxed);

// This is the value WaveKV selected using its own entry metadata. Its
// payload timestamp is older, so comparing reg_time again would leave the
// data plane permanently materializing a losing value.
sync_from_peer_at(
&state,
"contended",
"10.0.0.41",
&test_pubkey("winner-key"),
100,
);
reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap();

let proxy = state.lock();
let instance = &proxy.state.instances["contended"];
assert_eq!(instance.ip, "10.0.0.41".parse::<Ipv4Addr>().unwrap());
assert_eq!(instance.public_key, test_pubkey("winner-key"));
assert_eq!(encode_ts(instance.reg_time), 100);
assert_eq!(instance.num_connections(), 7);
assert!(!proxy
.state
.allocated_addresses
.contains(&"10.0.0.40".parse().unwrap()));
}

#[tokio::test]
async fn a_local_registration_survives_a_reload_that_cannot_see_it_yet() {
let state = create_test_state().await;
Expand Down
Loading