mod state_machine;
mod store;
mod types;
use types::ConnectionStore;
pub use types::{ConnectionHealth, ConnectionRecord, ConnectionState, PeerSnapshot};
use crate::heartbeat::state_machine::aggregate_peer_health;
use state_machine::{default_health_for_state, merge_upsert_reason, merge_upsert_state};
use std::collections::HashSet;
use store::{
add_index, add_indexes, build_peer_snapshot, canonical_peer_id, collect_peer_ids,
migrate_peer_state, remove_index, resolve_peer_id, unix_ms_now,
};
use tokio::sync::RwLock;
#[derive(Default)]
pub struct ConnectionManager {
store: RwLock<ConnectionStore>,
}
impl ConnectionManager {
pub fn new() -> Self {
Self::default()
}
pub async fn upsert_pending(
&self,
connection_id: String,
node_id: Option<String>,
device_id_hint: Option<String>,
endpoint_id: Option<String>,
) -> ConnectionRecord {
self.upsert(
connection_id,
node_id,
device_id_hint,
endpoint_id,
ConnectionState::Pending,
None,
)
.await
}
pub async fn set_connecting(&self, connection_id: &str) -> Option<ConnectionRecord> {
self.transition_state(connection_id, ConnectionState::Connecting, None)
.await
}
pub async fn set_connected(
&self,
connection_id: &str,
endpoint_id: Option<String>,
) -> Option<ConnectionRecord> {
self.set_connected_with_transport(connection_id, endpoint_id, None, None)
.await
}
pub async fn set_connected_with_transport(
&self,
connection_id: &str,
endpoint_id: Option<String>,
transport_stable_id: Option<u64>,
transport_source: Option<String>,
) -> Option<ConnectionRecord> {
let mut store = self.store.write().await;
let previous_endpoint = store
.by_id
.get(connection_id)
.and_then(|record| record.endpoint_id.clone());
let new_endpoint = endpoint_id.clone();
if endpoint_id.is_some() {
remove_index(
&mut store.by_endpoint_id,
previous_endpoint.as_deref(),
connection_id,
);
add_index(
&mut store.by_endpoint_id,
new_endpoint.as_deref(),
connection_id,
);
}
let previous = store.by_id.get(connection_id)?.clone();
let previous_peer_id = canonical_peer_id(&previous);
let record = store.by_id.get_mut(connection_id)?;
if endpoint_id.is_some() {
record.endpoint_id = endpoint_id;
}
let transport_changed = transport_stable_id
.map(|stable_id| record.transport_stable_id != Some(stable_id))
.unwrap_or(false);
if record.state != ConnectionState::Connected {
record.transition_count = record.transition_count.saturating_add(1);
}
if transport_changed {
record.transport_generation = record.transport_generation.saturating_add(1);
record.transport_stable_id = transport_stable_id;
record.transport_source = transport_source;
record.last_transport_change_at_ms = unix_ms_now();
record.replacement_count = record.replacement_count.saturating_add(1);
} else {
if transport_stable_id.is_some() {
record.transport_stable_id = transport_stable_id;
}
if transport_source.is_some() {
record.transport_source = transport_source;
}
}
record.state = ConnectionState::Connected;
record.status_reason = None;
record.last_disconnect_reason = None;
record.last_reconnect_reason = record.transport_source.clone();
record.updated_at_ms = unix_ms_now();
let updated = record.clone();
let next_peer_id = canonical_peer_id(&updated);
if transport_changed {
store
.health_by_peer
.insert(next_peer_id.clone(), ConnectionHealth::Unknown);
} else {
store
.health_by_peer
.entry(next_peer_id.clone())
.or_insert_with(|| default_health_for_state(&updated.state));
}
migrate_peer_state(&mut store, &previous_peer_id, &next_peer_id);
Some(updated)
}
pub async fn mark_transport_replaced(
&self,
connection_id: &str,
transport_stable_id: Option<u64>,
transport_source: Option<String>,
status_reason: Option<String>,
) -> Option<ConnectionRecord> {
let mut store = self.store.write().await;
let updated = {
let record = store.by_id.get_mut(connection_id)?;
let now = unix_ms_now();
let transport_changed = transport_stable_id
.map(|stable_id| record.transport_stable_id != Some(stable_id))
.unwrap_or(false);
if transport_changed {
record.transport_generation = record.transport_generation.saturating_add(1);
record.replacement_count = record.replacement_count.saturating_add(1);
}
record.transport_stable_id = transport_stable_id;
record.transport_source = transport_source;
record.last_transport_change_at_ms = now;
record.status_reason = status_reason;
record.last_reconnect_reason = record.transport_source.clone();
record.updated_at_ms = now;
record.clone()
};
let peer_id = canonical_peer_id(&updated);
store
.health_by_peer
.insert(peer_id, ConnectionHealth::Unknown);
Some(updated)
}
pub async fn current_transport_matches(
&self,
connection_id: &str,
transport_stable_id: Option<u64>,
) -> bool {
let Some(record) = self.store.read().await.by_id.get(connection_id).cloned() else {
return false;
};
match transport_stable_id {
Some(stable_id) => record.transport_stable_id == Some(stable_id),
None => record.transport_stable_id.is_none(),
}
}
pub async fn set_device_id(
&self,
connection_id: &str,
device_id: String,
) -> Option<ConnectionRecord> {
let mut store = self.store.write().await;
let existing = store.by_id.get(connection_id)?.clone();
let previous_peer_id = canonical_peer_id(&existing);
remove_index(
&mut store.by_device_id,
existing.device_id.as_deref(),
connection_id,
);
remove_index(
&mut store.by_device_hint_id,
existing.device_id_hint.as_deref(),
connection_id,
);
let mut updated = existing;
let prior_hint = updated
.device_id_hint
.clone()
.or_else(|| updated.device_id.clone());
updated.device_id = Some(device_id.clone());
if prior_hint.as_deref() == Some(device_id.as_str()) {
updated.device_id_hint = None;
} else {
updated.device_id_hint = prior_hint;
}
updated.updated_at_ms = unix_ms_now();
add_indexes(&mut store, &updated);
let next_peer_id = canonical_peer_id(&updated);
migrate_peer_state(&mut store, &previous_peer_id, &next_peer_id);
store
.by_id
.insert(connection_id.to_string(), updated.clone());
Some(updated)
}
pub async fn set_closing(
&self,
connection_id: &str,
reason: Option<String>,
) -> Option<ConnectionRecord> {
self.transition_state(connection_id, ConnectionState::Closing, reason)
.await
}
pub async fn set_closed(
&self,
connection_id: &str,
reason: Option<String>,
) -> Option<ConnectionRecord> {
self.transition_state(connection_id, ConnectionState::Closed, reason)
.await
}
pub async fn set_failed(
&self,
connection_id: &str,
reason: Option<String>,
) -> Option<ConnectionRecord> {
self.transition_state(connection_id, ConnectionState::Failed, reason)
.await
}
pub async fn remove(&self, connection_id: &str) -> Option<ConnectionRecord> {
let mut store = self.store.write().await;
let removed = store.by_id.remove(connection_id)?;
store.history_by_connection_id.insert(
connection_id.to_string(),
types::ConnectionHistory::from(&removed),
);
let peer_id = canonical_peer_id(&removed);
remove_index(
&mut store.by_node_id,
removed.node_id.as_deref(),
connection_id,
);
remove_index(
&mut store.by_device_id,
removed.device_id.as_deref(),
connection_id,
);
remove_index(
&mut store.by_device_hint_id,
removed.device_id_hint.as_deref(),
connection_id,
);
remove_index(
&mut store.by_endpoint_id,
removed.endpoint_id.as_deref(),
connection_id,
);
let peer_has_records = store
.by_id
.values()
.any(|r| canonical_peer_id(r) == peer_id);
if !peer_has_records {
store.health_by_peer.remove(&peer_id);
store.scopes_by_peer.remove(&peer_id);
}
Some(removed)
}
pub async fn get_by_connection_id(&self, connection_id: &str) -> Option<ConnectionRecord> {
self.store.read().await.by_id.get(connection_id).cloned()
}
#[cfg(test)]
pub async fn set_last_transport_change_at_ms_for_test(
&self,
connection_id: &str,
timestamp_ms: i64,
) -> Option<ConnectionRecord> {
let mut store = self.store.write().await;
let record = store.by_id.get_mut(connection_id)?;
record.last_transport_change_at_ms = timestamp_ms;
record.updated_at_ms = timestamp_ms.max(record.updated_at_ms);
Some(record.clone())
}
pub async fn get_by_node_id(&self, node_id: &str) -> Vec<ConnectionRecord> {
self.lookup_indexed(node_id, IndexType::Node).await
}
pub async fn get_by_device_id(&self, device_id: &str) -> Vec<ConnectionRecord> {
self.lookup_indexed(device_id, IndexType::Device).await
}
pub async fn get_by_endpoint_id(&self, endpoint_id: &str) -> Vec<ConnectionRecord> {
self.lookup_indexed(endpoint_id, IndexType::Endpoint).await
}
pub async fn list_all(&self) -> Vec<ConnectionRecord> {
self.store.read().await.by_id.values().cloned().collect()
}
pub async fn list_active(&self) -> Vec<ConnectionRecord> {
self.store
.read()
.await
.by_id
.values()
.filter(|record| {
matches!(
record.state,
ConnectionState::Pending
| ConnectionState::Connecting
| ConnectionState::Connected
)
})
.cloned()
.collect()
}
pub async fn add_scope(&self, id: &str, scope: &str) -> Vec<String> {
let mut store = self.store.write().await;
let Some(peer_id) = resolve_peer_id(&store, id) else {
return Vec::new();
};
let scopes = store.scopes_by_peer.entry(peer_id).or_default();
scopes.insert(scope.to_string());
scopes.iter().cloned().collect()
}
pub async fn release_scope(&self, id: &str, scope: Option<&str>) -> Vec<String> {
let mut store = self.store.write().await;
let Some(peer_id) = resolve_peer_id(&store, id) else {
return Vec::new();
};
let Some(scopes) = store.scopes_by_peer.get_mut(&peer_id) else {
return Vec::new();
};
match scope {
Some(scope) => {
scopes.remove(scope);
}
None => {
scopes.clear();
}
}
if scopes.is_empty() {
store.scopes_by_peer.remove(&peer_id);
return Vec::new();
}
scopes.iter().cloned().collect()
}
pub async fn get_scopes(&self, id: &str) -> Vec<String> {
let store = self.store.read().await;
let Some(peer_id) = resolve_peer_id(&store, id) else {
return Vec::new();
};
store
.scopes_by_peer
.get(&peer_id)
.map(|scopes| scopes.iter().cloned().collect())
.unwrap_or_default()
}
pub async fn are_same_peer(&self, left: &str, right: &str) -> bool {
let store = self.store.read().await;
let left_peer = resolve_peer_id(&store, left);
let right_peer = resolve_peer_id(&store, right);
left_peer.is_some() && left_peer == right_peer
}
pub async fn report_transport_status(
&self,
connection_id: &str,
active_transport: String,
parallel_transport: Option<String>,
) -> Option<PeerSnapshot> {
let mut store = self.store.write().await;
let record = store.by_id.get_mut(connection_id)?;
let transport_changed = record.active_transport != active_transport
|| record.parallel_transport != parallel_transport;
record.active_transport = active_transport;
record.parallel_transport = parallel_transport;
let now = unix_ms_now();
record.updated_at_ms = now;
if transport_changed {
record.transport_generation = record.transport_generation.saturating_add(1);
record.last_transport_change_at_ms = now;
}
let peer_id = canonical_peer_id(record);
build_peer_snapshot(&store, &peer_id)
}
pub async fn set_health(&self, id: &str, health: ConnectionHealth) -> Option<PeerSnapshot> {
let mut store = self.store.write().await;
let peer_id = resolve_peer_id(&store, id)?;
store.health_by_peer.insert(peer_id.clone(), health);
build_peer_snapshot(&store, &peer_id)
}
pub async fn report_transport_health(
&self,
id: &str,
transport: &str,
health: ConnectionHealth,
) -> Option<PeerSnapshot> {
let mut store = self.store.write().await;
let peer_id = resolve_peer_id(&store, id)?;
store
.transport_health_by_peer
.entry(peer_id.clone())
.or_default()
.insert(transport.to_string(), health);
let aggregated =
aggregate_peer_health(store.transport_health_by_peer.get(&peer_id).unwrap());
store.health_by_peer.insert(peer_id.clone(), aggregated);
build_peer_snapshot(&store, &peer_id)
}
pub async fn clear_transport_health(&self, id: &str) {
let mut store = self.store.write().await;
if let Some(peer_id) = resolve_peer_id(&store, id) {
store.transport_health_by_peer.remove(&peer_id);
}
}
pub async fn peer_snapshot(&self, id: &str) -> Option<PeerSnapshot> {
let store = self.store.read().await;
let peer_id = resolve_peer_id(&store, id)?;
build_peer_snapshot(&store, &peer_id)
}
pub async fn list_peer_snapshots(&self) -> Vec<PeerSnapshot> {
let store = self.store.read().await;
let peer_ids = collect_peer_ids(&store);
peer_ids
.iter()
.filter_map(|peer_id| build_peer_snapshot(&store, peer_id))
.collect()
}
pub async fn best_connection_for_peer(&self, id: &str) -> Option<ConnectionRecord> {
let store = self.store.read().await;
let peer_id = resolve_peer_id(&store, id)?;
let snapshot = build_peer_snapshot(&store, &peer_id);
let active_transport_stable_id = snapshot
.as_ref()
.and_then(|value| value.active_transport_stable_id);
let active_transport_generation = snapshot
.as_ref()
.map(|value| value.active_transport_generation)
.unwrap_or_default();
let active_node_id = snapshot.as_ref().and_then(|value| value.node_id.as_deref());
store
.by_id
.values()
.filter(|record| canonical_peer_id(record) == peer_id)
.filter(|record| matches!(record.state, ConnectionState::Connected))
.cloned()
.max_by(|left, right| {
let left_active = left.transport_stable_id == active_transport_stable_id
&& left.transport_generation == active_transport_generation;
let right_active = right.transport_stable_id == active_transport_stable_id
&& right.transport_generation == active_transport_generation;
let left_node_match = active_node_id == left.node_id.as_deref();
let right_node_match = active_node_id == right.node_id.as_deref();
left_active
.cmp(&right_active)
.then_with(|| left_node_match.cmp(&right_node_match))
.then_with(|| left.transport_generation.cmp(&right.transport_generation))
.then_with(|| left.updated_at_ms.cmp(&right.updated_at_ms))
})
}
async fn upsert(
&self,
connection_id: String,
node_id: Option<String>,
device_id_hint: Option<String>,
endpoint_id: Option<String>,
state: ConnectionState,
status_reason: Option<String>,
) -> ConnectionRecord {
let mut store = self.store.write().await;
if let Some(mut existing) = store.by_id.get(&connection_id).cloned() {
let previous_peer_id = canonical_peer_id(&existing);
remove_index(
&mut store.by_node_id,
existing.node_id.as_deref(),
&connection_id,
);
remove_index(
&mut store.by_device_id,
existing.device_id.as_deref(),
&connection_id,
);
remove_index(
&mut store.by_device_hint_id,
existing.device_id_hint.as_deref(),
&connection_id,
);
remove_index(
&mut store.by_endpoint_id,
existing.endpoint_id.as_deref(),
&connection_id,
);
if node_id.is_some() {
existing.node_id = node_id;
}
if let Some(device_id_hint) = device_id_hint {
if existing.device_id.is_none()
|| existing.device_id.as_deref() == Some(device_id_hint.as_str())
{
existing.device_id_hint = Some(device_id_hint);
}
}
if endpoint_id.is_some() {
existing.endpoint_id = endpoint_id;
}
let merged_state = merge_upsert_state(&existing.state, &state);
let merged_reason =
merge_upsert_reason(&existing, &merged_state, &state, status_reason);
existing.state = merged_state;
existing.status_reason = merged_reason;
existing.updated_at_ms = unix_ms_now();
add_indexes(&mut store, &existing);
let next_peer_id = canonical_peer_id(&existing);
store
.health_by_peer
.entry(next_peer_id.clone())
.or_insert_with(|| default_health_for_state(&existing.state));
migrate_peer_state(&mut store, &previous_peer_id, &next_peer_id);
store.by_id.insert(connection_id, existing.clone());
return existing;
}
let mut record =
ConnectionRecord::new(connection_id.clone(), node_id, device_id_hint, endpoint_id);
if let Some(history) = store.history_by_connection_id.get(&connection_id).cloned() {
record.transport_generation = history.transport_generation;
record.transition_count = history.transition_count;
record.connecting_transition_count = history.connecting_transition_count;
record.replacement_count = history.replacement_count;
record.retire_count = history.retire_count;
record.last_disconnect_reason = history.last_disconnect_reason;
record.last_reconnect_reason = history.last_reconnect_reason;
}
record.state = state;
record.status_reason = status_reason;
add_indexes(&mut store, &record);
store
.health_by_peer
.entry(canonical_peer_id(&record))
.or_insert_with(|| default_health_for_state(&record.state));
store.by_id.insert(connection_id, record.clone());
record
}
async fn transition_state(
&self,
connection_id: &str,
next_state: ConnectionState,
status_reason: Option<String>,
) -> Option<ConnectionRecord> {
let mut store = self.store.write().await;
let mut record = store.by_id.get(connection_id)?.clone();
let peer_id = canonical_peer_id(&record);
let previous_state = record.state.clone();
record.state = next_state.clone();
record.status_reason = status_reason;
record.updated_at_ms = unix_ms_now();
if previous_state != next_state {
record.transition_count = record.transition_count.saturating_add(1);
if matches!(next_state, ConnectionState::Connecting) {
record.connecting_transition_count =
record.connecting_transition_count.saturating_add(1);
}
}
if matches!(
next_state,
ConnectionState::Closing | ConnectionState::Closed | ConnectionState::Failed
) {
if !matches!(
previous_state,
ConnectionState::Closing | ConnectionState::Closed | ConnectionState::Failed
) {
record.retire_count = record.retire_count.saturating_add(1);
}
record.last_disconnect_reason = record.status_reason.clone();
} else if matches!(next_state, ConnectionState::Connected) {
record.last_reconnect_reason = record.status_reason.clone();
}
match &next_state {
ConnectionState::Connected => {
store
.health_by_peer
.entry(peer_id)
.or_insert(ConnectionHealth::Unknown);
}
ConnectionState::Pending | ConnectionState::Connecting => {
store
.health_by_peer
.insert(peer_id, ConnectionHealth::Unknown);
}
ConnectionState::Closing | ConnectionState::Closed | ConnectionState::Failed => {
store
.health_by_peer
.insert(peer_id, ConnectionHealth::Stale);
}
}
store
.by_id
.insert(connection_id.to_string(), record.clone());
Some(record)
}
async fn lookup_indexed(&self, key: &str, index_type: IndexType) -> Vec<ConnectionRecord> {
let store = self.store.read().await;
let ids: Vec<String> = match index_type {
IndexType::Node => store
.by_node_id
.get(key)
.map(|values| values.iter().cloned().collect())
.unwrap_or_default(),
IndexType::Device => {
let mut ids = HashSet::new();
if let Some(values) = store.by_device_id.get(key) {
ids.extend(values.iter().cloned());
}
if let Some(values) = store.by_device_hint_id.get(key) {
ids.extend(values.iter().cloned());
}
ids.into_iter().collect()
}
IndexType::Endpoint => store
.by_endpoint_id
.get(key)
.map(|values| values.iter().cloned().collect())
.unwrap_or_default(),
};
ids.into_iter()
.filter_map(|id| store.by_id.get(&id).cloned())
.collect()
}
}
enum IndexType {
Node,
Device,
Endpoint,
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn indexes_stay_consistent_across_state_transitions() {
let manager = ConnectionManager::new();
let record = manager
.upsert_pending(
"conn-1".to_string(),
Some("node-A".to_string()),
Some("device-A".to_string()),
Some("endpoint-A".to_string()),
)
.await;
assert_eq!(record.state, ConnectionState::Pending);
assert_eq!(manager.get_by_node_id("node-A").await.len(), 1);
assert_eq!(manager.get_by_device_id("device-A").await.len(), 1);
assert_eq!(manager.get_by_endpoint_id("endpoint-A").await.len(), 1);
manager.set_connecting("conn-1").await;
manager
.set_connected("conn-1", Some("endpoint-B".to_string()))
.await;
let updated = manager
.get_by_connection_id("conn-1")
.await
.expect("updated record");
assert_eq!(updated.state, ConnectionState::Connected);
assert_eq!(updated.endpoint_id.as_deref(), Some("endpoint-B"));
assert_eq!(manager.get_by_endpoint_id("endpoint-A").await.len(), 0);
assert_eq!(manager.get_by_endpoint_id("endpoint-B").await.len(), 1);
manager
.set_closed("conn-1", Some("test-close".to_string()))
.await;
let closed = manager
.get_by_connection_id("conn-1")
.await
.expect("closed record");
assert_eq!(closed.state, ConnectionState::Closed);
assert_eq!(manager.list_active().await.len(), 0);
manager.remove("conn-1").await;
assert!(manager.get_by_connection_id("conn-1").await.is_none());
assert_eq!(manager.get_by_node_id("node-A").await.len(), 0);
assert_eq!(manager.get_by_device_id("device-A").await.len(), 0);
assert_eq!(manager.get_by_endpoint_id("endpoint-B").await.len(), 0);
}
#[tokio::test]
async fn set_device_id_reindexes_without_changing_state() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-2".to_string(),
Some("node-A".to_string()),
Some("device-A".to_string()),
Some("endpoint-A".to_string()),
)
.await;
manager
.set_connected("conn-2", Some("endpoint-A".to_string()))
.await;
let updated = manager
.set_device_id("conn-2", "device-B".to_string())
.await
.expect("updated record");
assert_eq!(updated.state, ConnectionState::Connected);
assert_eq!(updated.device_id.as_deref(), Some("device-B"));
assert_eq!(updated.device_id_hint.as_deref(), Some("device-A"));
assert_eq!(manager.get_by_device_id("device-A").await.len(), 1);
assert_eq!(manager.get_by_device_id("device-B").await.len(), 1);
assert!(manager.are_same_peer("device-A", "device-B").await);
}
#[tokio::test]
async fn test_pending_to_connected_to_closed_lifecycle() {
let manager = ConnectionManager::new();
let record = manager
.upsert_pending(
"lc-1".to_string(),
Some("node-X".to_string()),
Some("dev-X".to_string()),
None,
)
.await;
assert_eq!(record.state, ConnectionState::Pending);
assert_eq!(manager.list_active().await.len(), 1);
let connecting = manager.set_connecting("lc-1").await.unwrap();
assert_eq!(connecting.state, ConnectionState::Connecting);
assert_eq!(manager.list_active().await.len(), 1);
let connected = manager
.set_connected("lc-1", Some("ep-X".to_string()))
.await
.unwrap();
assert_eq!(connected.state, ConnectionState::Connected);
assert_eq!(connected.endpoint_id.as_deref(), Some("ep-X"));
assert_eq!(manager.list_active().await.len(), 1);
let closing = manager
.set_closing("lc-1", Some("graceful".into()))
.await
.unwrap();
assert_eq!(closing.state, ConnectionState::Closing);
assert_eq!(manager.list_active().await.len(), 0);
let closed = manager
.set_closed("lc-1", Some("done".into()))
.await
.unwrap();
assert_eq!(closed.state, ConnectionState::Closed);
assert_eq!(manager.list_active().await.len(), 0);
assert!(manager.get_by_connection_id("lc-1").await.is_some());
assert_eq!(manager.get_by_node_id("node-X").await.len(), 1);
}
#[tokio::test]
async fn test_upsert_preserves_existing_fields() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"up-1".to_string(),
Some("node-1".to_string()),
Some("dev-1".to_string()),
Some("ep-1".to_string()),
)
.await;
let updated = manager
.upsert(
"up-1".to_string(),
None,
None,
None,
ConnectionState::Connected,
None,
)
.await;
assert_eq!(updated.state, ConnectionState::Connected);
assert_eq!(updated.node_id.as_deref(), Some("node-1"));
assert_eq!(updated.device_id, None);
assert_eq!(updated.device_id_hint.as_deref(), Some("dev-1"));
assert_eq!(updated.endpoint_id.as_deref(), Some("ep-1"));
}
#[tokio::test]
async fn test_set_failed_retains_indexes() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"fail-1".to_string(),
Some("node-F".to_string()),
Some("dev-F".to_string()),
None,
)
.await;
manager.set_failed("fail-1", Some("timeout".into())).await;
let failed = manager
.get_by_connection_id("fail-1")
.await
.expect("record exists");
assert_eq!(failed.state, ConnectionState::Failed);
assert_eq!(failed.status_reason.as_deref(), Some("timeout"));
assert_eq!(manager.get_by_node_id("node-F").await.len(), 1);
assert_eq!(manager.get_by_device_id("dev-F").await.len(), 1);
assert_eq!(manager.list_active().await.len(), 0);
}
#[tokio::test]
async fn test_list_active_excludes_terminal_states() {
let manager = ConnectionManager::new();
manager.upsert_pending("a-1".into(), None, None, None).await;
manager.upsert_pending("a-2".into(), None, None, None).await;
manager.upsert_pending("a-3".into(), None, None, None).await;
manager.set_connecting("a-1").await;
manager.set_connected("a-2", Some("ep".into())).await;
manager.set_failed("a-3", None).await;
let active = manager.list_active().await;
assert_eq!(active.len(), 2);
let ids: Vec<_> = active.iter().map(|r| r.connection_id.as_str()).collect();
assert!(ids.contains(&"a-1"));
assert!(ids.contains(&"a-2"));
assert!(!ids.contains(&"a-3"));
}
#[tokio::test]
async fn test_concurrent_connections_per_node() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"multi-1".into(),
Some("shared-node".into()),
Some("dev-A".into()),
None,
)
.await;
manager
.upsert_pending(
"multi-2".into(),
Some("shared-node".into()),
Some("dev-A".into()),
None,
)
.await;
let by_node = manager.get_by_node_id("shared-node").await;
assert_eq!(by_node.len(), 2);
let by_device = manager.get_by_device_id("dev-A").await;
assert_eq!(by_device.len(), 2);
manager.remove("multi-1").await;
assert_eq!(manager.get_by_node_id("shared-node").await.len(), 1);
assert_eq!(manager.get_by_device_id("dev-A").await.len(), 1);
}
#[tokio::test]
async fn test_connected_with_transport_tracks_generation_changes() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-1".into(),
Some("node-1".into()),
Some("device-1".into()),
Some("endpoint-1".into()),
)
.await;
let first = manager
.set_connected_with_transport(
"conn-1",
Some("endpoint-1".into()),
Some(10),
Some("incoming".into()),
)
.await
.expect("record exists");
assert_eq!(first.transport_generation, 1);
assert_eq!(first.transport_stable_id, Some(10));
assert_eq!(first.transport_source.as_deref(), Some("incoming"));
let unchanged = manager
.set_connected_with_transport(
"conn-1",
Some("endpoint-1".into()),
Some(10),
Some("incoming-refresh".into()),
)
.await
.expect("record exists");
assert_eq!(unchanged.transport_generation, 1);
assert_eq!(unchanged.transport_stable_id, Some(10));
assert_eq!(
unchanged.transport_source.as_deref(),
Some("incoming-refresh")
);
let replaced = manager
.set_connected_with_transport(
"conn-1",
Some("endpoint-1".into()),
Some(11),
Some("replacement".into()),
)
.await
.expect("record exists");
assert_eq!(replaced.transport_generation, 2);
assert_eq!(replaced.transport_stable_id, Some(11));
assert_eq!(replaced.transport_source.as_deref(), Some("replacement"));
assert!(manager.current_transport_matches("conn-1", Some(11)).await);
assert!(!manager.current_transport_matches("conn-1", Some(10)).await);
}
#[tokio::test]
async fn report_transport_status_updates_peer_snapshot_without_replacing_transport() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-status".into(),
Some("node-status".into()),
Some("device-status".into()),
Some("endpoint-status".into()),
)
.await;
manager
.set_connected_with_transport(
"conn-status",
Some("endpoint-status".into()),
Some(42),
Some("incoming".into()),
)
.await;
let before = manager
.get_by_connection_id("conn-status")
.await
.expect("record exists before status report")
.last_transport_change_at_ms;
tokio::time::sleep(std::time::Duration::from_millis(1)).await;
let snapshot = manager
.report_transport_status(
"conn-status",
"webrtc".to_string(),
Some("iroh".to_string()),
)
.await
.expect("peer snapshot");
assert_eq!(snapshot.active_transport, "webrtc");
assert_eq!(snapshot.parallel_transport.as_deref(), Some("iroh"));
assert_eq!(snapshot.active_transport_stable_id, Some(42));
assert_eq!(snapshot.active_transport_generation, 2);
let record = manager
.get_by_connection_id("conn-status")
.await
.expect("record exists after status report");
assert_eq!(record.transport_stable_id, Some(42));
assert_eq!(record.replacement_count, 1);
let after = manager
.get_by_connection_id("conn-status")
.await
.expect("record exists after status report")
.last_transport_change_at_ms;
assert!(after >= before);
}
#[tokio::test]
async fn report_transport_status_advances_generation_for_iroh_path_migration() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-path".into(),
Some("node-path".into()),
Some("device-path".into()),
Some("endpoint-path".into()),
)
.await;
manager
.set_connected_with_transport(
"conn-path",
Some("endpoint-path".into()),
Some(7),
Some("incoming".into()),
)
.await;
let relay = manager
.report_transport_status("conn-path", "iroh-relay".to_string(), None)
.await
.expect("relay snapshot");
assert_eq!(relay.active_transport, "iroh-relay");
assert_eq!(relay.active_transport_generation, 2);
let direct = manager
.report_transport_status("conn-path", "iroh-quic".to_string(), None)
.await
.expect("direct snapshot");
assert_eq!(direct.active_transport, "iroh-quic");
assert_eq!(direct.active_transport_generation, 3);
let record = manager
.get_by_connection_id("conn-path")
.await
.expect("record exists after path migration");
assert_eq!(record.transport_stable_id, Some(7));
assert_eq!(record.replacement_count, 1);
}
#[tokio::test]
async fn test_mark_transport_replaced_advances_generation_without_state_reset() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-2".into(),
Some("node-2".into()),
Some("device-2".into()),
Some("endpoint-2".into()),
)
.await;
manager
.set_connected_with_transport(
"conn-2",
Some("endpoint-2".into()),
Some(20),
Some("outgoing".into()),
)
.await;
let replaced = manager
.mark_transport_replaced(
"conn-2",
Some(21),
Some("incoming-replacement".into()),
Some("replacement-in-progress".into()),
)
.await
.expect("record exists");
assert_eq!(replaced.transport_generation, 2);
assert_eq!(replaced.transport_stable_id, Some(21));
assert_eq!(
replaced.transport_source.as_deref(),
Some("incoming-replacement")
);
assert_eq!(
replaced.status_reason.as_deref(),
Some("replacement-in-progress")
);
assert_eq!(replaced.state, ConnectionState::Connected);
let peer = manager
.peer_snapshot("device-2")
.await
.expect("peer snapshot");
assert_eq!(peer.health, ConnectionHealth::Unknown);
}
#[tokio::test]
async fn mark_transport_replaced_without_new_binding_preserves_generation() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-2b".into(),
Some("node-2b".into()),
Some("device-2b".into()),
Some("endpoint-2b".into()),
)
.await;
manager
.set_connected_with_transport(
"conn-2b",
Some("endpoint-2b".into()),
Some(30),
Some("outgoing".into()),
)
.await
.expect("connected record");
let replaced = manager
.mark_transport_replaced(
"conn-2b",
None,
Some("outgoing-closed".into()),
Some("replacement-in-progress".into()),
)
.await
.expect("record exists");
assert_eq!(replaced.transport_generation, 1);
assert_eq!(replaced.transport_stable_id, None);
assert_eq!(
replaced.status_reason.as_deref(),
Some("replacement-in-progress")
);
assert_eq!(replaced.replacement_count, 1);
assert_eq!(replaced.state, ConnectionState::Connected);
}
#[tokio::test]
async fn upsert_pending_does_not_downgrade_connected_replacement_record() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-stable".into(),
Some("node-1".into()),
Some("device-1".into()),
Some("endpoint-1".into()),
)
.await;
manager
.set_connected_with_transport(
"conn-stable",
Some("endpoint-1".into()),
Some(101),
Some("outgoing".into()),
)
.await
.expect("connected record");
manager
.mark_transport_replaced(
"conn-stable",
Some(102),
Some("incoming-replacement".into()),
Some("replacement-in-progress".into()),
)
.await
.expect("replacement record");
let updated = manager
.upsert_pending(
"conn-stable".into(),
Some("node-1".into()),
Some("device-1".into()),
Some("endpoint-1".into()),
)
.await;
assert_eq!(updated.state, ConnectionState::Connected);
assert_eq!(
updated.status_reason.as_deref(),
Some("replacement-in-progress")
);
assert_eq!(updated.transport_generation, 2);
assert_eq!(updated.transport_stable_id, Some(102));
}
#[tokio::test]
async fn peer_snapshot_uses_live_connected_record_for_transport_metadata() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-live".into(),
Some("node-live".into()),
Some("device-1".into()),
Some("endpoint-live".into()),
)
.await;
manager
.set_connected_with_transport(
"conn-live",
Some("endpoint-live".into()),
Some(101),
Some("incoming-live".into()),
)
.await
.expect("live connected record");
manager
.report_transport_status("conn-live", "webrtc".into(), Some("iroh-relay".into()))
.await
.expect("live peer snapshot");
manager
.upsert_pending(
"conn-pending-replacement".into(),
Some("node-live".into()),
Some("device-1".into()),
Some("endpoint-replacement".into()),
)
.await;
manager
.set_connecting("conn-pending-replacement")
.await
.expect("pending replacement record");
let snapshot = manager
.peer_snapshot("device-1")
.await
.expect("peer snapshot");
assert_eq!(snapshot.status, ConnectionState::Connected);
assert_eq!(
snapshot.connection_ids.first().map(String::as_str),
Some("conn-live")
);
assert_eq!(snapshot.active_transport_stable_id, Some(101));
assert_eq!(snapshot.active_transport_generation, 2);
assert_eq!(snapshot.active_transport, "webrtc");
assert_eq!(snapshot.parallel_transport.as_deref(), Some("iroh-relay"));
}
#[tokio::test]
async fn transition_counters_increment_without_idle_churn() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-counters".into(),
Some("node-1".into()),
Some("device-1".into()),
Some("endpoint-1".into()),
)
.await;
manager.set_connecting("conn-counters").await;
manager
.set_connected_with_transport(
"conn-counters",
Some("endpoint-1".into()),
Some(201),
Some("initial".into()),
)
.await
.expect("connected record");
manager
.mark_transport_replaced(
"conn-counters",
Some(202),
Some("replacement".into()),
Some("replacement-in-progress".into()),
)
.await
.expect("replacement record");
let record = manager
.get_by_connection_id("conn-counters")
.await
.expect("record");
assert_eq!(record.transition_count, 2);
assert_eq!(record.connecting_transition_count, 1);
assert_eq!(record.replacement_count, 2);
assert_eq!(record.last_reconnect_reason.as_deref(), Some("replacement"));
let peer = manager.peer_snapshot("device-1").await.expect("peer");
assert_eq!(peer.transition_count, 2);
assert_eq!(peer.connecting_transition_count, 1);
assert_eq!(peer.replacement_count, 2);
assert_eq!(peer.retire_count, 0);
}
#[tokio::test]
async fn remove_and_recreate_preserves_monotonic_generation_history() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-history".into(),
Some("node-1".into()),
Some("device-1".into()),
Some("endpoint-1".into()),
)
.await;
manager.set_connecting("conn-history").await;
manager
.set_connected_with_transport(
"conn-history",
Some("endpoint-1".into()),
Some(301),
Some("initial".into()),
)
.await
.expect("connected record");
manager
.mark_transport_replaced(
"conn-history",
Some(302),
Some("replacement".into()),
Some("replacement-in-progress".into()),
)
.await
.expect("replacement record");
manager
.remove("conn-history")
.await
.expect("removed record");
let recreated = manager
.upsert_pending(
"conn-history".into(),
Some("node-1".into()),
Some("device-1".into()),
Some("endpoint-1".into()),
)
.await;
assert_eq!(recreated.transport_generation, 2);
assert_eq!(recreated.transition_count, 2);
assert_eq!(recreated.connecting_transition_count, 1);
assert_eq!(recreated.replacement_count, 2);
let reconnect = manager
.set_connected_with_transport(
"conn-history",
Some("endpoint-1".into()),
Some(303),
Some("recreated".into()),
)
.await
.expect("reconnected record");
assert_eq!(reconnect.transport_generation, 3);
assert_eq!(reconnect.transition_count, 3);
assert_eq!(reconnect.replacement_count, 3);
}
#[tokio::test]
async fn peer_snapshot_uses_connection_id_until_device_hint_exists() {
let manager = ConnectionManager::new();
manager
.upsert_pending("conn-provisional".into(), Some("node-1".into()), None, None)
.await;
let provisional = manager
.peer_snapshot("conn-provisional")
.await
.expect("provisional snapshot");
assert_eq!(provisional.peer_id, "conn-provisional");
assert_eq!(provisional.device_id, None);
assert_eq!(provisional.device_id_hint, None);
assert_eq!(provisional.node_id.as_deref(), Some("node-1"));
manager
.upsert_pending(
"conn-provisional".into(),
Some("node-1".into()),
Some("device-hint-1".into()),
None,
)
.await;
let hinted = manager
.peer_snapshot("device-hint-1")
.await
.expect("hinted snapshot");
assert_eq!(hinted.peer_id, "device-hint-1");
assert_eq!(hinted.device_id, None);
assert_eq!(hinted.device_id_hint.as_deref(), Some("device-hint-1"));
assert_eq!(hinted.node_id.as_deref(), Some("node-1"));
assert!(
manager
.are_same_peer("device-hint-1", "conn-provisional")
.await
);
assert!(manager.are_same_peer("device-hint-1", "node-1").await);
}
#[tokio::test]
async fn authoritative_device_binding_rekeys_peer_and_preserves_state() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-authoritative".into(),
Some("node-2".into()),
Some("device-hint-2".into()),
None,
)
.await;
manager.set_connecting("conn-authoritative").await;
manager.add_scope("device-hint-2", "persistent").await;
manager
.set_health("device-hint-2", ConnectionHealth::Suspect)
.await;
let updated = manager
.set_device_id("conn-authoritative", "device-2".into())
.await
.expect("authoritative record");
assert_eq!(updated.device_id.as_deref(), Some("device-2"));
assert_eq!(updated.device_id_hint.as_deref(), Some("device-hint-2"));
let authoritative = manager
.peer_snapshot("device-2")
.await
.expect("authoritative snapshot");
assert_eq!(authoritative.peer_id, "device-2");
assert_eq!(authoritative.device_id.as_deref(), Some("device-2"));
assert_eq!(
authoritative.device_id_hint.as_deref(),
Some("device-hint-2")
);
assert_eq!(authoritative.status, ConnectionState::Connecting);
assert_eq!(authoritative.health, ConnectionHealth::Suspect);
assert_eq!(authoritative.scopes, vec!["persistent".to_string()]);
assert!(
manager
.are_same_peer("device-2", "conn-authoritative")
.await
);
assert!(manager.are_same_peer("device-2", "node-2").await);
assert!(manager.are_same_peer("device-hint-2", "device-2").await);
let hinted = manager
.peer_snapshot("device-hint-2")
.await
.expect("hinted snapshot");
assert_eq!(hinted.peer_id, "device-2");
assert_eq!(hinted.device_id.as_deref(), Some("device-2"));
assert_eq!(hinted.device_id_hint.as_deref(), Some("device-hint-2"));
}
#[tokio::test]
async fn multiple_transports_collapse_into_one_device_peer() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-a".into(),
Some("node-a".into()),
Some("device-3".into()),
None,
)
.await;
manager
.upsert_pending(
"conn-b".into(),
Some("node-b".into()),
Some("device-3".into()),
None,
)
.await;
manager.set_connecting("conn-a").await;
manager
.set_connected_with_transport("conn-b", None, Some(30), Some("incoming".into()))
.await;
let snapshots = manager.list_peer_snapshots().await;
assert_eq!(snapshots.len(), 1);
assert_eq!(snapshots[0].peer_id, "device-3");
assert_eq!(snapshots[0].device_id, None);
assert_eq!(snapshots[0].device_id_hint.as_deref(), Some("device-3"));
assert_eq!(snapshots[0].status, ConnectionState::Connected);
assert_eq!(snapshots[0].connection_ids.len(), 2);
}
#[tokio::test]
async fn best_connection_for_peer_prefers_latest_active_transport() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-old".into(),
Some("node-1".into()),
Some("device-1".into()),
Some("endpoint-1".into()),
)
.await;
manager
.set_device_id("conn-old", "device-1".into())
.await
.expect("old device binding");
manager
.set_connected_with_transport(
"conn-old",
Some("endpoint-1".into()),
Some(10),
Some("incoming".into()),
)
.await
.expect("old record");
manager
.upsert_pending(
"conn-new".into(),
Some("node-1".into()),
Some("device-1".into()),
Some("endpoint-1".into()),
)
.await;
manager
.set_device_id("conn-new", "device-1".into())
.await
.expect("new device binding");
manager
.set_connected_with_transport(
"conn-new",
Some("endpoint-1".into()),
Some(11),
Some("replacement".into()),
)
.await
.expect("new record");
let best = manager
.best_connection_for_peer("device-1")
.await
.expect("best record");
assert_eq!(best.connection_id, "conn-new");
assert_eq!(best.transport_stable_id, Some(11));
}
#[tokio::test]
async fn best_connection_for_peer_returns_none_for_unknown_peer() {
let manager = ConnectionManager::new();
assert!(
manager
.best_connection_for_peer("nonexistent-peer")
.await
.is_none(),
"unknown peer should return None"
);
}
#[tokio::test]
async fn peer_snapshot_surfaces_failed_over_older_connected_same_device() {
let manager = ConnectionManager::new();
manager
.upsert_pending(
"conn-stale".into(),
Some("node-peer".into()),
Some("device-peer".into()),
Some("node-peer".into()),
)
.await;
manager.set_connecting("conn-stale").await;
manager
.set_connected("conn-stale", Some("node-peer".into()))
.await;
manager
.set_device_id("conn-stale", "device-peer-uuid".into())
.await
.expect("set device");
manager
.upsert_pending(
"conn-dial".into(),
Some("node-peer".into()),
Some("device-peer".into()),
Some("node-peer".into()),
)
.await;
manager.set_connecting("conn-dial").await;
manager
.set_device_id("conn-dial", "device-peer-uuid".into())
.await
.expect("set device on dial row");
manager
.set_failed("conn-dial", Some("ensure_connected_addr-timeout".into()))
.await
.expect("mark dial failed");
let snapshot = manager
.peer_snapshot("device-peer-uuid")
.await
.expect("merged peer");
assert_eq!(snapshot.status, ConnectionState::Failed);
}
}