lda/wakey
1
0 Fork 0
Code Issues 6 Pull Requests Actions Packages Projects Releases 20 Wiki Activity

refactor(fleet): separate naming, routes, and snapshot writes

Browse Source
This commit is contained in:
lda committed 2026-10-11 19:24:39 +07:00 Verified
1 parent a4029e20bf
commit c172ce6603
8 files changed
+663 -411
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

wakey-control-plane/src/api/control/fleet/build.rs
+36 -237
View File
@@ -2,10 +2,15 @@ use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::net::IpAddr;
use macaddr::MacAddr;
use wakey_core::{EndpointSource, Presence};
use wakey_core::Presence;
use crate::state::{AgentDeviceWithChildren, DeviceIdentifier, KnownDevice, KnownDeviceSummary};
mod addresses;
mod routes;
use addresses::add_address_evidence;
use routes::{add_endpoint_routes, add_legacy_routes, compare_route_quality};
use super::types::{
FleetDevice, FleetDeviceAgent, FleetDeviceEndpoint, FleetWakeRoute, ListFleetDevicesQuery,
};
@@ -26,7 +31,9 @@ pub(crate) struct AgentRuntimeStatus {
struct FleetAccumulator {
device_key: String,
display_name: Option<String>,
label_ip: Option<IpAddr>,
/// Naming-only evidence, used when no name, MAC, or summary IP is available.
/// Must not contribute to summary IPs or wake routes.
fallback_ip: Option<IpAddr>,
known_device: Option<KnownDeviceSummary>,
pinned: bool,
ips: BTreeSet<IpAddr>,
@@ -46,7 +53,7 @@ impl Default for FleetAccumulator {
Self {
device_key: String::new(),
display_name: None,
label_ip: None,
fallback_ip: None,
known_device: None,
pinned: false,
ips: BTreeSet::new(),
@@ -197,23 +204,22 @@ fn add_agent_device_to_entry(
agent_device: AgentDeviceWithChildren,
context: &FleetBuildContext,
) {
// Retain address evidence for naming only, without promoting stale IPs into
// the active summary or wake-route candidates.
let label_ip = agent_device
.ips
.iter()
.copied()
.chain(
agent_device
.endpoints
.iter()
.filter_map(|endpoint| endpoint.key.ip),
)
.min();
entry.label_ip = entry.label_ip.into_iter().chain(label_ip).min();
let device_presence = agent_device.presence();
let device_offline = device_presence == Presence::Offline;
if let Some(summary) = device_known_device(&agent_device, context)
add_address_evidence(entry, &agent_device);
add_names_and_identity(entry, &agent_device, context);
let status = add_agent_summary(entry, &agent_device, context);
if agent_device.endpoints.is_empty() {
add_legacy_routes(entry, &agent_device, &status);
} else {
add_endpoint_routes(entry, &agent_device, &status);
}
}
fn add_names_and_identity(
entry: &mut FleetAccumulator,
agent_device: &AgentDeviceWithChildren,
context: &FleetBuildContext,
) {
if let Some(summary) = device_known_device(agent_device, context)
&& entry.known_device.is_none()
{
if let Some(ref name) = agent_device.device.display_name {
@@ -222,33 +228,19 @@ fn add_agent_device_to_entry(
entry.pinned = summary.pinned;
entry.known_device = Some(summary);
}
if agent_device.endpoints.is_empty() {
for mac in &agent_device.macs {
entry.macs.insert(*mac);
}
if !device_offline {
for ip in &agent_device.ips {
entry.ips.insert(*ip);
}
}
} else {
for endpoint in &agent_device.endpoints {
if let Some(mac) = endpoint.key.mac {
entry.macs.insert(mac);
}
if endpoint_ip_is_summary_eligible(endpoint.key.source, endpoint.presence)
&& let Some(ip) = endpoint.key.ip
{
entry.ips.insert(ip);
}
}
}
for hostname in &agent_device.hostnames {
if entry.display_name.is_none() {
entry.display_name = Some(hostname.clone());
}
entry.hostnames.insert(hostname.clone());
}
}
fn add_agent_summary(
entry: &mut FleetAccumulator,
agent_device: &AgentDeviceWithChildren,
context: &FleetBuildContext,
) -> AgentRuntimeStatus {
let first_seen = agent_device.device.first_seen();
let last_seen = agent_device.device.last_seen();
entry.first_seen_unix = Some(
@@ -263,7 +255,7 @@ fn add_agent_device_to_entry(
.map(|current| current.max(last_seen))
.unwrap_or(last_seen),
);
entry.presence = std::cmp::max(entry.presence, device_presence);
entry.presence = std::cmp::max(entry.presence, agent_device.presence());
let agent_id = agent_device.device.agent_id.clone();
let status = context
@@ -288,105 +280,7 @@ fn add_agent_device_to_entry(
connected: status.connected,
last_seen_unix: last_seen,
});
if agent_device.endpoints.is_empty() {
entry.sources.insert("device".to_string());
}
for endpoint in &agent_device.endpoints {
let source = endpoint_source_label(endpoint.key.source).to_string();
entry.sources.insert(source.clone());
let endpoint_last_seen = endpoint.last_seen_unix.unwrap_or(last_seen);
let endpoint_first_seen = endpoint.first_seen_unix.or(Some(first_seen));
entry.endpoints.push(FleetDeviceEndpoint {
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
source: source.clone(),
mac: endpoint.key.mac,
ip: endpoint.key.ip,
hostname: endpoint.hostname.clone(),
interface: endpoint.interface.clone(),
presence: endpoint.presence,
first_seen_unix: endpoint_first_seen,
last_seen_unix: Some(endpoint_last_seen),
});
let rid = route_id(
&agent_id,
endpoint.key.mac.as_ref(),
endpoint.key.ip.as_ref(),
);
let wakeable = status.connected && endpoint.key.mac.is_some();
insert_route(
entry,
FleetWakeRoute {
route_id: rid,
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
mac: endpoint.key.mac,
ip: endpoint.key.ip,
hostname: endpoint.hostname.clone(),
interface: endpoint.interface.clone(),
source,
presence: endpoint.presence,
last_seen_unix: endpoint_last_seen,
wakeable,
},
);
}
if agent_device.endpoints.is_empty() && !agent_device.macs.is_empty() {
for mac in &agent_device.macs {
let ip_for_mac = agent_device.ips.first().copied();
let hostname_for_mac = agent_device.hostnames.first().cloned();
let rid = route_id(&agent_id, Some(mac), ip_for_mac.as_ref());
let wakeable = status.connected;
insert_route(
entry,
FleetWakeRoute {
route_id: rid,
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
mac: Some(*mac),
ip: ip_for_mac,
hostname: hostname_for_mac,
interface: None,
source: "device".to_string(),
presence: device_presence,
last_seen_unix: last_seen,
wakeable,
},
);
}
}
if agent_device.endpoints.is_empty()
&& agent_device.macs.is_empty()
&& let Some(ip) = agent_device.ips.first()
{
let hostname = agent_device.hostnames.first().cloned();
let rid = route_id(&agent_id, None, Some(ip));
insert_route(
entry,
FleetWakeRoute {
route_id: rid,
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
mac: None,
ip: Some(*ip),
hostname,
interface: None,
source: "device".to_string(),
presence: device_presence,
last_seen_unix: last_seen,
wakeable: false,
},
);
}
status
}
fn add_identifier_to_entry(entry: &mut FleetAccumulator, identifier: &DeviceIdentifier) {
@@ -439,37 +333,6 @@ fn device_known_device(
})
}
fn endpoint_ip_is_summary_eligible(source: EndpointSource, presence: Presence) -> bool {
match source {
EndpointSource::DhcpLease => true,
EndpointSource::Neighbor => presence != Presence::Offline,
EndpointSource::HookNeighbor | EndpointSource::HookDhcp => false,
}
}
fn insert_route(entry: &mut FleetAccumulator, candidate: FleetWakeRoute) {
match entry.routes.entry(candidate.route_id.clone()) {
std::collections::btree_map::Entry::Vacant(slot) => {
slot.insert(candidate);
}
std::collections::btree_map::Entry::Occupied(mut slot) => {
if compare_route_quality(&candidate, slot.get()).is_gt() {
slot.insert(candidate);
}
}
}
}
fn compare_route_quality(a: &FleetWakeRoute, b: &FleetWakeRoute) -> std::cmp::Ordering {
a.connected
.cmp(&b.connected)
.then_with(|| a.wakeable.cmp(&b.wakeable))
.then_with(|| a.ip.is_some().cmp(&b.ip.is_some()))
.then_with(|| a.presence.cmp(&b.presence))
.then_with(|| a.last_seen_unix.cmp(&b.last_seen_unix))
.then_with(|| source_quality_rank(&a.source).cmp(&source_quality_rank(&b.source)))
}
pub(crate) fn known_device_summary(device: &KnownDevice) -> KnownDeviceSummary {
KnownDeviceSummary {
device_id: device.device_id.clone(),
@@ -480,6 +343,7 @@ pub(crate) fn known_device_summary(device: &KnownDevice) -> KnownDeviceSummary {
impl FleetAccumulator {
fn into_response(self) -> FleetDevice {
let display_name = self.display_label();
let mut route_candidates = self.routes.into_values().collect::<Vec<_>>();
route_candidates
.sort_by(|a, b| compare_route_quality(b, a).then_with(|| a.agent_id.cmp(&b.agent_id)));
@@ -487,43 +351,6 @@ impl FleetAccumulator {
.iter()
.find(|route| route.wakeable)
.cloned();
let display_name = self
.display_name
.or_else(|| self.hostnames.iter().next().cloned())
.or_else(|| {
// Summary IPs omit offline endpoints; labels must not discard
// their address evidence or imply that unknown means offline.
let address = self
.ips
.iter()
.next()
.copied()
.or(self.label_ip)
.or_else(|| {
self.endpoints
.iter()
.filter_map(|endpoint| endpoint.ip)
.min()
})
.map(|ip| ip.to_string())
.or_else(|| self.macs.iter().next().map(|mac| mac.to_string()))
.or_else(|| {
self.endpoints
.iter()
.find_map(|endpoint| endpoint.mac)
.map(|mac| mac.to_string())
});
address.map(|address| {
let prefix = if self.presence == Presence::Offline {
"Offline device"
} else {
"Device"
};
format!("{prefix}: {address}")
})
})
.unwrap_or_else(|| "(unknown device)".to_string());
FleetDevice {
device_key: self.device_key,
display_name,
@@ -544,34 +371,6 @@ impl FleetAccumulator {
}
}
fn endpoint_source_label(source: EndpointSource) -> &'static str {
match source {
EndpointSource::Neighbor => "neighbor",
EndpointSource::DhcpLease => "dhcp_lease",
EndpointSource::HookNeighbor => "hook_neighbor",
EndpointSource::HookDhcp => "hook_dhcp",
}
}
fn source_quality_rank(source: &str) -> u8 {
match source {
"neighbor" => 4,
"dhcp_lease" => 3,
"hook_neighbor" => 2,
"hook_dhcp" => 1,
_ => 0,
}
}
fn route_id(agent_id: &str, mac: Option<&MacAddr>, ip: Option<&IpAddr>) -> String {
format!(
"{}|{}|{}",
agent_id,
mac.map(|m| m.to_string()).unwrap_or_default(),
ip.map(|i| i.to_string()).unwrap_or_default()
)
}
fn normalize_filter(value: Option<&str>) -> Option<String> {
value
.map(|value| value.trim().to_ascii_lowercase())
wakey-control-plane/src/api/control/fleet/build/addresses.rs
+81
View File
@@ -0,0 +1,81 @@
//! Current-address summaries and naming-only evidence have different eligibility.
use super::FleetAccumulator;
use crate::state::AgentDeviceWithChildren;
use wakey_core::{EndpointSource, Presence};
/// Keep naming evidence distinct from addresses eligible for the summary.
pub(super) fn add_address_evidence(
entry: &mut FleetAccumulator,
agent_device: &AgentDeviceWithChildren,
) {
// Retain address evidence for naming only, without promoting stale IPs into
// the active summary or wake-route candidates.
let fallback_ip = agent_device
.ips
.iter()
.copied()
.chain(
agent_device
.endpoints
.iter()
.filter_map(|endpoint| endpoint.key.ip),
)
.min();
entry.fallback_ip = entry.fallback_ip.into_iter().chain(fallback_ip).min();
if agent_device.endpoints.is_empty() {
for mac in &agent_device.macs {
entry.macs.insert(*mac);
}
if agent_device.presence() != Presence::Offline {
for ip in &agent_device.ips {
entry.ips.insert(*ip);
}
}
} else {
for endpoint in &agent_device.endpoints {
if let Some(mac) = endpoint.key.mac {
entry.macs.insert(mac);
}
if endpoint_ip_is_summary_eligible(endpoint.key.source, endpoint.presence)
&& let Some(ip) = endpoint.key.ip
{
entry.ips.insert(ip);
}
}
}
}
fn endpoint_ip_is_summary_eligible(source: EndpointSource, presence: Presence) -> bool {
match source {
EndpointSource::DhcpLease => true,
EndpointSource::Neighbor => presence != Presence::Offline,
EndpointSource::HookNeighbor | EndpointSource::HookDhcp => false,
}
}
impl FleetAccumulator {
/// Names and MACs outrank address evidence; presence belongs in the badge.
/// Summary IPs outrank naming-only fallback IPs, which may be inactive.
pub(super) fn display_label(&self) -> String {
self.display_name
.clone()
.or_else(|| self.hostnames.iter().next().cloned())
.or_else(|| {
self.macs
.iter()
.next()
.map(|mac| mac.to_string())
.or_else(|| {
self.ips
.iter()
.next()
.copied()
.or(self.fallback_ip)
.map(|ip| ip.to_string())
})
.map(|address| format!("Device: {address}"))
})
.unwrap_or_else(|| "(unknown device)".to_string())
}
}
wakey-control-plane/src/api/control/fleet/build/routes.rs
+173
View File
@@ -0,0 +1,173 @@
//! Build and rank agent-specific wake routes without promoting IPs into summaries.
use super::{AgentRuntimeStatus, FleetAccumulator, FleetDeviceEndpoint, FleetWakeRoute};
use crate::state::AgentDeviceWithChildren;
use macaddr::MacAddr;
use std::net::IpAddr;
use wakey_core::EndpointSource;
/// Endpoint evidence is preserved even when its IP is excluded from the summary.
pub(super) fn add_endpoint_routes(
entry: &mut FleetAccumulator,
agent_device: &AgentDeviceWithChildren,
status: &AgentRuntimeStatus,
) {
let agent_id = &agent_device.device.agent_id;
let first_seen = agent_device.device.first_seen();
let last_seen = agent_device.device.last_seen();
for endpoint in &agent_device.endpoints {
let source = endpoint_source_label(endpoint.key.source).to_string();
entry.sources.insert(source.clone());
let endpoint_last_seen = endpoint.last_seen_unix.unwrap_or(last_seen);
let endpoint_first_seen = endpoint.first_seen_unix.or(Some(first_seen));
entry.endpoints.push(FleetDeviceEndpoint {
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
source: source.clone(),
mac: endpoint.key.mac,
ip: endpoint.key.ip,
hostname: endpoint.hostname.clone(),
interface: endpoint.interface.clone(),
presence: endpoint.presence,
first_seen_unix: endpoint_first_seen,
last_seen_unix: Some(endpoint_last_seen),
});
let rid = route_id(
agent_id,
endpoint.key.mac.as_ref(),
endpoint.key.ip.as_ref(),
);
let wakeable = status.connected && endpoint.key.mac.is_some();
insert_route(
entry,
FleetWakeRoute {
route_id: rid,
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
mac: endpoint.key.mac,
ip: endpoint.key.ip,
hostname: endpoint.hostname.clone(),
interface: endpoint.interface.clone(),
source,
presence: endpoint.presence,
last_seen_unix: endpoint_last_seen,
wakeable,
},
);
}
}
/// Compatibility for stored rows without authoritative endpoint evidence.
pub(super) fn add_legacy_routes(
entry: &mut FleetAccumulator,
agent_device: &AgentDeviceWithChildren,
status: &AgentRuntimeStatus,
) {
entry.sources.insert("device".to_string());
let agent_id = &agent_device.device.agent_id;
let device_presence = agent_device.presence();
let last_seen = agent_device.device.last_seen();
if !agent_device.macs.is_empty() {
for mac in &agent_device.macs {
let ip_for_mac = agent_device.ips.first().copied();
let hostname_for_mac = agent_device.hostnames.first().cloned();
let rid = route_id(agent_id, Some(mac), ip_for_mac.as_ref());
let wakeable = status.connected;
insert_route(
entry,
FleetWakeRoute {
route_id: rid,
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
mac: Some(*mac),
ip: ip_for_mac,
hostname: hostname_for_mac,
interface: None,
source: "device".to_string(),
presence: device_presence,
last_seen_unix: last_seen,
wakeable,
},
);
}
}
if agent_device.macs.is_empty()
&& let Some(ip) = agent_device.ips.first()
{
let hostname = agent_device.hostnames.first().cloned();
let rid = route_id(agent_id, None, Some(ip));
insert_route(
entry,
FleetWakeRoute {
route_id: rid,
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
mac: None,
ip: Some(*ip),
hostname,
interface: None,
source: "device".to_string(),
presence: device_presence,
last_seen_unix: last_seen,
wakeable: false,
},
);
}
}
fn insert_route(entry: &mut FleetAccumulator, candidate: FleetWakeRoute) {
match entry.routes.entry(candidate.route_id.clone()) {
std::collections::btree_map::Entry::Vacant(slot) => {
slot.insert(candidate);
}
std::collections::btree_map::Entry::Occupied(mut slot) => {
if compare_route_quality(&candidate, slot.get()).is_gt() {
slot.insert(candidate);
}
}
}
}
pub(super) fn compare_route_quality(a: &FleetWakeRoute, b: &FleetWakeRoute) -> std::cmp::Ordering {
a.connected
.cmp(&b.connected)
.then_with(|| a.wakeable.cmp(&b.wakeable))
.then_with(|| a.ip.is_some().cmp(&b.ip.is_some()))
.then_with(|| a.presence.cmp(&b.presence))
.then_with(|| a.last_seen_unix.cmp(&b.last_seen_unix))
.then_with(|| source_quality_rank(&a.source).cmp(&source_quality_rank(&b.source)))
}
fn endpoint_source_label(source: EndpointSource) -> &'static str {
match source {
EndpointSource::Neighbor => "neighbor",
EndpointSource::DhcpLease => "dhcp_lease",
EndpointSource::HookNeighbor => "hook_neighbor",
EndpointSource::HookDhcp => "hook_dhcp",
}
}
fn source_quality_rank(source: &str) -> u8 {
match source {
"neighbor" => 4,
"dhcp_lease" => 3,
"hook_neighbor" => 2,
"hook_dhcp" => 1,
_ => 0,
}
}
fn route_id(agent_id: &str, mac: Option<&MacAddr>, ip: Option<&IpAddr>) -> String {
format!(
"{}|{}|{}",
agent_id,
mac.map(|m| m.to_string()).unwrap_or_default(),
ip.map(|i| i.to_string()).unwrap_or_default()
)
}
wakey-control-plane/src/api/control/fleet/tests.rs
+74 -5
View File
@@ -728,7 +728,7 @@ fn unknown_ip_only_offline_device_is_hidden_by_default() {
);
assert_eq!(devices.len(), 1);
assert_eq!(devices[0].display_name, "Offline device: 192.168.1.2");
assert_eq!(devices[0].display_name, "Device: 192.168.1.2");
filter_fleet_devices(&mut devices, &ListFleetDevicesQuery::default());
assert!(devices.is_empty());
}
@@ -762,7 +762,7 @@ fn endpoint_presence_overrides_conflicting_cached_summary() {
assert_eq!(devices[0].endpoints[0].presence, evidence);
assert_eq!(devices[0].agents[0].connected, connected);
if evidence == Presence::Offline {
assert_eq!(devices[0].display_name, "Offline device: 192.168.1.2");
assert_eq!(devices[0].display_name, "Device: 192.168.1.2");
assert!(devices[0].ips.is_empty());
}
}
@@ -776,7 +776,7 @@ fn unnamed_endpoint_labels_preserve_presence_and_inactive_ips() {
Presence::Offline,
None,
Some("192.168.1.2"),
"Offline device: 192.168.1.2",
"Device: 192.168.1.2",
),
(
Presence::Unknown,
@@ -788,7 +788,7 @@ fn unnamed_endpoint_labels_preserve_presence_and_inactive_ips() {
Presence::Offline,
Some("aa:bb:cc:dd:ee:01"),
None,
"Offline device: AA:BB:CC:DD:EE:01",
"Device: AA:BB:CC:DD:EE:01",
),
] {
let mut row = offline_ip_only_unknown("agent-a", "test", "192.168.1.2", 20);
@@ -806,6 +806,75 @@ fn unnamed_endpoint_labels_preserve_presence_and_inactive_ips() {
}
}
#[test]
fn unnamed_mac_device_does_not_use_inactive_ip_or_presence_as_its_name() {
for presence in [Presence::Offline, Presence::Unknown, Presence::Online] {
let mut row =
offline_ip_only_unknown("agent-a", "mac:22:07:53:BF:FC:12", "192.168.100.125", 20);
row.macs = vec!["22:07:53:bf:fc:12".parse().unwrap()];
row.hostnames.clear();
row.endpoints = [EndpointSource::Neighbor, EndpointSource::HookNeighbor]
.into_iter()
.map(|source| {
let mut endpoint = endpoint_with_presence(
source,
Some("22:07:53:bf:fc:12"),
Some("192.168.100.125"),
presence,
20,
);
endpoint.hostname = None;
endpoint
})
.collect();
let devices = build_fleet_devices(Vec::new(), vec![row], &context(&["agent-a"]));
assert_eq!(devices[0].display_name, "Device: 22:07:53:BF:FC:12");
assert_eq!(devices[0].presence, presence);
assert_eq!(devices[0].endpoints.len(), 2);
if presence == Presence::Offline {
assert!(devices[0].ips.is_empty());
}
}
}
#[test]
fn ip_only_label_prefers_summary_ip_and_uses_inactive_evidence_only_as_fallback() {
for has_current_ip in [true, false] {
let mut row = offline_ip_only_unknown("agent-a", "test", "192.168.1.2", 20);
let mut inactive = endpoint_with_presence(
EndpointSource::Neighbor,
None,
Some("192.168.1.2"),
Presence::Offline,
20,
);
inactive.hostname = None;
row.endpoints = vec![inactive];
if has_current_ip {
let mut current = endpoint_with_presence(
EndpointSource::Neighbor,
None,
Some("192.168.1.10"),
Presence::Online,
20,
);
current.hostname = None;
row.endpoints.push(current);
}
let devices = build_fleet_devices(Vec::new(), vec![row], &context(&["agent-a"]));
if has_current_ip {
assert_eq!(devices[0].display_name, "Device: 192.168.1.10");
assert_eq!(
devices[0].ips,
vec!["192.168.1.10".parse::<IpAddr>().unwrap()]
);
} else {
assert_eq!(devices[0].display_name, "Device: 192.168.1.2");
assert!(devices[0].ips.is_empty());
}
}
}
#[test]
fn visibility_all_keeps_unknown_ip_only_offline_device() {
let mut devices = build_fleet_devices(
@@ -827,7 +896,7 @@ fn visibility_all_keeps_unknown_ip_only_offline_device() {
},
);
assert_eq!(devices.len(), 1);
assert_eq!(devices[0].display_name, "Offline device: 192.168.1.2");
assert_eq!(devices[0].display_name, "Device: 192.168.1.2");
}
#[test]
wakey-control-plane/src/state/store/agent_devices.rs
+53 -168
View File
@@ -5,6 +5,8 @@ use super::Store;
use super::helpers::core::*;
use crate::state::types::*;
mod snapshot;
impl Store {
/// Expire agent-reported devices, preserving all operator-saved identities.
/// `last_seen_unix` on the parent is CC snapshot inclusion time, not proof
@@ -52,7 +54,6 @@ impl Store {
let mut incoming_keys = std::collections::HashSet::with_capacity(devices.len());
let snapshot_time = now_unix();
let snapshot_time_i64 = i64::try_from(snapshot_time).context("snapshot time overflow")?;
let existing_keys: Vec<String> = sqlx::query_scalar!(
"SELECT device_key FROM agent_devices WHERE agent_id = ?1",
@@ -69,173 +70,7 @@ impl Store {
let device_key = device_key_from_id(device_id);
incoming_keys.insert(device_key.clone());
let endpoints = endpoints_for_storage(device);
let presence = presence_to_str(endpoint_presence(&endpoints, device.presence));
let display_name: Option<String> = None;
sqlx::query!(
"INSERT INTO agent_devices (agent_id, device_key, presence, display_name, first_seen_unix, last_seen_unix)
VALUES (?1, ?2, ?3, ?4, ?5, ?6)
ON CONFLICT (agent_id, device_key) DO UPDATE SET
presence = excluded.presence,
display_name = excluded.display_name,
last_seen_unix = excluded.last_seen_unix",
agent_id,
device_key,
presence,
display_name,
snapshot_time_i64,
snapshot_time_i64
)
.execute(&mut *tx)
.await
.context("failed upserting agent device")?;
// Replace child rows.
sqlx::query!(
"DELETE FROM agent_device_macs WHERE agent_id = ?1 AND device_key = ?2",
agent_id,
device_key
)
.execute(&mut *tx)
.await
.context("failed deleting device macs")?;
if !device.macs.is_empty() {
let mut builder = sqlx::QueryBuilder::new(
"INSERT INTO agent_device_macs (agent_id, device_key, mac) ",
);
builder.push_values(&device.macs, |mut b, mac| {
b.push_bind(agent_id)
.push_bind(device_key.clone())
.push_bind(mac.to_string().to_ascii_lowercase());
});
builder
.build()
.execute(&mut *tx)
.await
.context("failed inserting device macs")?;
}
sqlx::query!(
"DELETE FROM agent_device_ips WHERE agent_id = ?1 AND device_key = ?2",
agent_id,
device_key
)
.execute(&mut *tx)
.await
.context("failed deleting device ips")?;
if !device.ips.is_empty() {
let mut builder = sqlx::QueryBuilder::new(
"INSERT INTO agent_device_ips (agent_id, device_key, ip) ",
);
builder.push_values(&device.ips, |mut b, ip| {
b.push_bind(agent_id)
.push_bind(device_key.clone())
.push_bind(ip.to_string());
});
builder
.build()
.execute(&mut *tx)
.await
.context("failed inserting device ips")?;
}
sqlx::query(
"DELETE FROM agent_device_endpoints WHERE agent_id = ?1 AND device_key = ?2",
)
.bind(agent_id)
.bind(&device_key)
.execute(&mut *tx)
.await
.context("failed deleting device endpoints")?;
if !endpoints.is_empty() {
let mut builder = sqlx::QueryBuilder::new(
"INSERT INTO agent_device_endpoints \
(agent_id, device_key, endpoint_key, source, mac, ip, hostname, interface, presence, first_seen_unix, last_seen_unix) ",
);
builder.push_values(endpoints, |mut b, endpoint| {
let first_seen = endpoint.first_seen_unix.unwrap_or(snapshot_time);
let last_seen = endpoint.last_seen_unix.unwrap_or(snapshot_time);
b.push_bind(agent_id)
.push_bind(device_key.clone())
.push_bind(endpoint_storage_key(&endpoint))
.push_bind(endpoint_source_to_str(endpoint.key.source))
.push_bind(
endpoint
.key
.mac
.map(|mac| mac.to_string().to_ascii_lowercase()),
)
.push_bind(endpoint.key.ip.map(|ip| ip.to_string()))
.push_bind(endpoint.hostname.clone())
.push_bind(endpoint.interface.clone())
.push_bind(presence_to_str(endpoint.presence))
.push_bind(i64::try_from(first_seen).unwrap_or(i64::MAX))
.push_bind(i64::try_from(last_seen).unwrap_or(i64::MAX));
});
builder
.build()
.execute(&mut *tx)
.await
.context("failed inserting device endpoints")?;
}
sqlx::query!(
"DELETE FROM agent_device_hostnames WHERE agent_id = ?1 AND device_key = ?2",
agent_id,
device_key
)
.execute(&mut *tx)
.await
.context("failed deleting device hostnames")?;
if !device.names.is_empty() {
let mut builder = sqlx::QueryBuilder::new(
"INSERT INTO agent_device_hostnames (agent_id, device_key, hostname) ",
);
builder.push_values(&device.names, |mut b, hostname| {
b.push_bind(agent_id)
.push_bind(device_key.clone())
.push_bind(hostname);
});
builder
.build()
.execute(&mut *tx)
.await
.context("failed inserting device hostnames")?;
}
sqlx::query!(
"DELETE FROM agent_device_facts WHERE agent_id = ?1 AND device_key = ?2",
agent_id,
device_key
)
.execute(&mut *tx)
.await
.context("failed deleting device facts")?;
if !device.observations.is_empty() {
let mut facts_json = Vec::with_capacity(device.observations.len());
for obs in &device.observations {
facts_json.push(serde_json::to_string(obs).context("failed serializing fact")?);
}
let mut builder = sqlx::QueryBuilder::new(
"INSERT INTO agent_device_facts (agent_id, device_key, fact_json) ",
);
builder.push_values(facts_json, |mut b, fact| {
b.push_bind(agent_id)
.push_bind(device_key.clone())
.push_bind(fact);
});
builder
.build()
.execute(&mut *tx)
.await
.context("failed inserting device facts")?;
}
snapshot::write_device(&mut tx, agent_id, &device_key, device, snapshot_time).await?;
}
for old_key in existing_keys {
@@ -784,6 +619,56 @@ mod tests {
assert_eq!(rows[0].endpoints[0].presence, Presence::Unknown);
}
#[tokio::test]
async fn child_write_failure_rolls_back_the_whole_snapshot_without_invalidating_fleet() {
let ts = TestStore::new().await;
let store = ts.store();
let original = sample_device("aa:bb:cc:dd:ee:01", "192.168.1.10", "original");
store
.replace_agent_device_snapshot("agent-a", &[original])
.await
.unwrap();
let before = store
.list_agent_device_rows_for_agent("agent-a")
.await
.unwrap()
.remove(0);
let revision = store.fleet_revision();
let mut updated = sample_device("aa:bb:cc:dd:ee:01", "192.168.1.20", "updated");
updated.neighbors[0].state = NeighborState::Failed;
let mut invalid = sample_device("aa:bb:cc:dd:ee:02", "192.168.1.30", "duplicate");
// The second device fails late, after the first device and several child
// tables were written. Every helper must share the outer transaction.
invalid.names.push("duplicate".into());
let error = store
.replace_agent_device_snapshot("agent-a", &[updated, invalid])
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("failed inserting device hostnames")
);
assert_eq!(store.fleet_revision(), revision);
let mut rows = store
.list_agent_device_rows_for_agent("agent-a")
.await
.unwrap();
assert_eq!(rows.len(), 1, "the failed new device must not survive");
let after = rows.remove(0);
assert_eq!(after.device.device_key, before.device.device_key);
assert_eq!(after.device.presence, before.device.presence);
assert_eq!(after.device.first_seen_unix, before.device.first_seen_unix);
assert_eq!(after.device.last_seen_unix, before.device.last_seen_unix);
assert_eq!(after.macs, before.macs);
assert_eq!(after.ips, before.ips);
assert_eq!(after.hostnames, before.hostnames);
assert_eq!(after.endpoints, before.endpoints);
assert_eq!(after.facts, before.facts);
}
#[tokio::test]
async fn snapshot_inserts_devices_and_children() {
let ts = TestStore::new().await;
wakey-control-plane/src/state/store/agent_devices/snapshot.rs
+237
View File
@@ -0,0 +1,237 @@
//! Writes one snapshot device using the caller's transaction. No helper commits.
use anyhow::{Context, Result};
use sqlx::SqliteConnection;
use wakey_core::{Device, DeviceEndpoint};
use super::{
endpoint_presence, endpoint_source_to_str, endpoint_storage_key, endpoints_for_storage,
presence_to_str,
};
pub(super) async fn write_device(
connection: &mut SqliteConnection,
agent_id: &str,
device_key: &str,
device: &Device,
snapshot_time: u64,
) -> Result<()> {
let endpoints = endpoints_for_storage(device);
let presence = presence_to_str(endpoint_presence(&endpoints, device.presence));
upsert_parent(connection, agent_id, device_key, presence, snapshot_time).await?;
replace_macs(connection, agent_id, device_key, device).await?;
replace_ips(connection, agent_id, device_key, device).await?;
replace_endpoints(connection, agent_id, device_key, &endpoints, snapshot_time).await?;
replace_hostnames(connection, agent_id, device_key, device).await?;
replace_facts(connection, agent_id, device_key, device).await
}
/// Updates snapshot inclusion time, while preserving the parent's first-seen time.
async fn upsert_parent(
connection: &mut SqliteConnection,
agent_id: &str,
device_key: &str,
presence: &str,
snapshot_time: u64,
) -> Result<()> {
let snapshot_time_i64 = i64::try_from(snapshot_time).context("snapshot time overflow")?;
let display_name: Option<String> = None;
sqlx::query!(
"INSERT INTO agent_devices (agent_id, device_key, presence, display_name, first_seen_unix, last_seen_unix)
VALUES (?1, ?2, ?3, ?4, ?5, ?6)
ON CONFLICT (agent_id, device_key) DO UPDATE SET
presence = excluded.presence,
display_name = excluded.display_name,
last_seen_unix = excluded.last_seen_unix",
agent_id,
device_key,
presence,
display_name,
snapshot_time_i64,
snapshot_time_i64
)
.execute(&mut *connection)
.await
.context("failed upserting agent device")?;
Ok(())
}
async fn replace_macs(
connection: &mut SqliteConnection,
agent_id: &str,
device_key: &str,
device: &Device,
) -> Result<()> {
sqlx::query!(
"DELETE FROM agent_device_macs WHERE agent_id = ?1 AND device_key = ?2",
agent_id,
device_key
)
.execute(&mut *connection)
.await
.context("failed deleting device macs")?;
if !device.macs.is_empty() {
let mut builder =
sqlx::QueryBuilder::new("INSERT INTO agent_device_macs (agent_id, device_key, mac) ");
builder.push_values(&device.macs, |mut b, mac| {
b.push_bind(agent_id)
.push_bind(device_key)
.push_bind(mac.to_string().to_ascii_lowercase());
});
builder
.build()
.execute(&mut *connection)
.await
.context("failed inserting device macs")?;
}
Ok(())
}
async fn replace_ips(
connection: &mut SqliteConnection,
agent_id: &str,
device_key: &str,
device: &Device,
) -> Result<()> {
sqlx::query!(
"DELETE FROM agent_device_ips WHERE agent_id = ?1 AND device_key = ?2",
agent_id,
device_key
)
.execute(&mut *connection)
.await
.context("failed deleting device ips")?;
if !device.ips.is_empty() {
let mut builder =
sqlx::QueryBuilder::new("INSERT INTO agent_device_ips (agent_id, device_key, ip) ");
builder.push_values(&device.ips, |mut b, ip| {
b.push_bind(agent_id)
.push_bind(device_key)
.push_bind(ip.to_string());
});
builder
.build()
.execute(&mut *connection)
.await
.context("failed inserting device ips")?;
}
Ok(())
}
async fn replace_endpoints(
connection: &mut SqliteConnection,
agent_id: &str,
device_key: &str,
endpoints: &[DeviceEndpoint],
snapshot_time: u64,
) -> Result<()> {
sqlx::query("DELETE FROM agent_device_endpoints WHERE agent_id = ?1 AND device_key = ?2")
.bind(agent_id)
.bind(device_key)
.execute(&mut *connection)
.await
.context("failed deleting device endpoints")?;
if !endpoints.is_empty() {
let mut builder = sqlx::QueryBuilder::new(
"INSERT INTO agent_device_endpoints \
(agent_id, device_key, endpoint_key, source, mac, ip, hostname, interface, presence, first_seen_unix, last_seen_unix) ",
);
builder.push_values(endpoints, |mut b, endpoint| {
let first_seen = endpoint.first_seen_unix.unwrap_or(snapshot_time);
let last_seen = endpoint.last_seen_unix.unwrap_or(snapshot_time);
b.push_bind(agent_id)
.push_bind(device_key)
.push_bind(endpoint_storage_key(endpoint))
.push_bind(endpoint_source_to_str(endpoint.key.source))
.push_bind(
endpoint
.key
.mac
.map(|mac| mac.to_string().to_ascii_lowercase()),
)
.push_bind(endpoint.key.ip.map(|ip| ip.to_string()))
.push_bind(endpoint.hostname.clone())
.push_bind(endpoint.interface.clone())
.push_bind(presence_to_str(endpoint.presence))
.push_bind(i64::try_from(first_seen).unwrap_or(i64::MAX))
.push_bind(i64::try_from(last_seen).unwrap_or(i64::MAX));
});
builder
.build()
.execute(&mut *connection)
.await
.context("failed inserting device endpoints")?;
}
Ok(())
}
async fn replace_hostnames(
connection: &mut SqliteConnection,
agent_id: &str,
device_key: &str,
device: &Device,
) -> Result<()> {
sqlx::query!(
"DELETE FROM agent_device_hostnames WHERE agent_id = ?1 AND device_key = ?2",
agent_id,
device_key
)
.execute(&mut *connection)
.await
.context("failed deleting device hostnames")?;
if !device.names.is_empty() {
let mut builder = sqlx::QueryBuilder::new(
"INSERT INTO agent_device_hostnames (agent_id, device_key, hostname) ",
);
builder.push_values(&device.names, |mut b, hostname| {
b.push_bind(agent_id)
.push_bind(device_key)
.push_bind(hostname);
});
builder
.build()
.execute(&mut *connection)
.await
.context("failed inserting device hostnames")?;
}
Ok(())
}
async fn replace_facts(
connection: &mut SqliteConnection,
agent_id: &str,
device_key: &str,
device: &Device,
) -> Result<()> {
sqlx::query!(
"DELETE FROM agent_device_facts WHERE agent_id = ?1 AND device_key = ?2",
agent_id,
device_key
)
.execute(&mut *connection)
.await
.context("failed deleting device facts")?;
if !device.observations.is_empty() {
let mut facts_json = Vec::with_capacity(device.observations.len());
for obs in &device.observations {
facts_json.push(serde_json::to_string(obs).context("failed serializing fact")?);
}
let mut builder = sqlx::QueryBuilder::new(
"INSERT INTO agent_device_facts (agent_id, device_key, fact_json) ",
);
builder.push_values(facts_json, |mut b, fact| {
b.push_bind(agent_id).push_bind(device_key).push_bind(fact);
});
builder
.build()
.execute(&mut *connection)
.await
.context("failed inserting device facts")?;
}
Ok(())
}