refactor(state): share consistent device reads and row assembly

This commit is contained in:
lda committed 2026-10-11 19:36:37 +07:00 Verified
1 parent c172ce6603
commit a04df6f573
3 files changed
+473 -223

No files matched your search

@@ -5,6 +5,7 @@ use super::Store;
use super::helpers::core::*; use super::helpers::core::*;
use crate::state::types::*; use crate::state::types::*;
mod read;
mod snapshot; mod snapshot;
impl Store { impl Store {
@@ -95,242 +96,23 @@ impl Store {
Ok(incoming_keys.len()) Ok(incoming_keys.len())
} }
/// List all agent device rows with their child MAC/IP/hostname/fact rows. /// List all agent devices and their child rows from one consistent read view.
#[tracing::instrument(name = "db_list_fleet_devices", skip_all)] #[tracing::instrument(name = "db_list_fleet_devices", skip_all)]
pub async fn list_agent_device_rows(&self) -> Result<Vec<AgentDeviceWithChildren>> { pub async fn list_agent_device_rows(&self) -> Result<Vec<AgentDeviceWithChildren>> {
let devices = sqlx::query_as!( read::load(&self.pool, None).await
AgentDeviceRow,
r#"SELECT agent_id, device_key, presence, display_name,
first_seen_unix, last_seen_unix
FROM agent_devices
ORDER BY agent_id, device_key"#,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent devices")?;
let macs = sqlx::query_as!(
AgentDeviceMacRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!"
FROM agent_device_macs"#,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent device macs")?;
let ips = sqlx::query_as!(
AgentDeviceIpRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!"
FROM agent_device_ips"#,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent device ips")?;
let hostnames = sqlx::query_as!(
AgentDeviceHostnameRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!"
FROM agent_device_hostnames"#,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent device hostnames")?;
let endpoints = sqlx::query_as::<_, AgentDeviceEndpointRow>(
r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip,
hostname, interface, presence, first_seen_unix, last_seen_unix
FROM agent_device_endpoints"#,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent device endpoints")?;
let facts = sqlx::query_as!(
AgentDeviceFactRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!"
FROM agent_device_facts"#,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent device facts")?;
Ok(assemble_device_rows(
devices, macs, ips, hostnames, endpoints, facts,
))
} }
/// List agent device rows for a single agent. /// List one agent's devices using the same assembly path as fleet reads.
#[allow(dead_code)] #[allow(dead_code)]
#[tracing::instrument(name = "db_list_agent_devices", skip_all)] #[tracing::instrument(name = "db_list_agent_devices", skip_all)]
pub async fn list_agent_device_rows_for_agent( pub async fn list_agent_device_rows_for_agent(
&self, &self,
agent_id: &str, agent_id: &str,
) -> Result<Vec<AgentDeviceWithChildren>> { ) -> Result<Vec<AgentDeviceWithChildren>> {
let devices = sqlx::query_as!( read::load(&self.pool, Some(agent_id)).await
AgentDeviceRow,
r#"SELECT agent_id, device_key, presence, display_name,
first_seen_unix, last_seen_unix
FROM agent_devices
WHERE agent_id = ?1
ORDER BY device_key"#,
agent_id,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent devices for agent")?;
let macs = sqlx::query_as!(
AgentDeviceMacRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!"
FROM agent_device_macs
WHERE agent_id = ?1"#,
agent_id,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent device macs for agent")?;
let ips = sqlx::query_as!(
AgentDeviceIpRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!"
FROM agent_device_ips
WHERE agent_id = ?1"#,
agent_id,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent device ips for agent")?;
let hostnames = sqlx::query_as!(
AgentDeviceHostnameRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!"
FROM agent_device_hostnames
WHERE agent_id = ?1"#,
agent_id,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent device hostnames for agent")?;
let endpoints = sqlx::query_as::<_, AgentDeviceEndpointRow>(
r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip,
hostname, interface, presence, first_seen_unix, last_seen_unix
FROM agent_device_endpoints
WHERE agent_id = ?1"#,
)
.bind(agent_id)
.fetch_all(&self.pool)
.await
.context("failed listing agent device endpoints for agent")?;
let facts = sqlx::query_as!(
AgentDeviceFactRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!"
FROM agent_device_facts
WHERE agent_id = ?1"#,
agent_id,
)
.fetch_all(&self.pool)
.await
.context("failed listing agent device facts for agent")?;
Ok(assemble_device_rows(
devices, macs, ips, hostnames, endpoints, facts,
))
} }
} }
fn assemble_device_rows(
devices: Vec<AgentDeviceRow>,
macs: Vec<AgentDeviceMacRow>,
ips: Vec<AgentDeviceIpRow>,
hostnames: Vec<AgentDeviceHostnameRow>,
endpoints: Vec<AgentDeviceEndpointRow>,
facts: Vec<AgentDeviceFactRow>,
) -> Vec<AgentDeviceWithChildren> {
use std::collections::BTreeMap;
let mut mac_map: BTreeMap<(&str, &str), Vec<&AgentDeviceMacRow>> = BTreeMap::new();
for row in &macs {
mac_map
.entry((row.agent_id.as_str(), row.device_key.as_str()))
.or_default()
.push(row);
}
let mut ip_map: BTreeMap<(&str, &str), Vec<&AgentDeviceIpRow>> = BTreeMap::new();
for row in &ips {
ip_map
.entry((row.agent_id.as_str(), row.device_key.as_str()))
.or_default()
.push(row);
}
let mut hostname_map: BTreeMap<(&str, &str), Vec<String>> = BTreeMap::new();
for row in &hostnames {
hostname_map
.entry((row.agent_id.as_str(), row.device_key.as_str()))
.or_default()
.push(row.hostname.clone());
}
let mut endpoint_map: BTreeMap<(&str, &str), Vec<DeviceEndpoint>> = BTreeMap::new();
for row in &endpoints {
if let Some(endpoint) = row.to_endpoint() {
endpoint_map
.entry((row.agent_id.as_str(), row.device_key.as_str()))
.or_default()
.push(endpoint);
}
}
let mut fact_map: BTreeMap<(&str, &str), Vec<String>> = BTreeMap::new();
for row in &facts {
fact_map
.entry((row.agent_id.as_str(), row.device_key.as_str()))
.or_default()
.push(row.fact_json.clone());
}
devices
.into_iter()
.map(|device| {
let key = (device.agent_id.as_str(), device.device_key.as_str());
let macs: Vec<macaddr::MacAddr> = mac_map
.get(&key)
.into_iter()
.flatten()
.filter_map(|row| match macaddr::MacAddr::try_from(*row) {
Ok(mac) => Some(mac),
Err(e) => {
::tracing::warn!(error = %e, agent_id = %device.agent_id, device_key = %device.device_key, raw_mac = %row.mac, "failed to parse mac from agent_device_macs row");
None
}
})
.collect();
let ips: Vec<std::net::IpAddr> = ip_map
.get(&key)
.into_iter()
.flatten()
.filter_map(|row| match std::net::IpAddr::try_from(*row) {
Ok(ip) => Some(ip),
Err(e) => {
::tracing::warn!(error = %e, agent_id = %device.agent_id, device_key = %device.device_key, raw_ip = %row.ip, "failed to parse ip from agent_device_ips row");
None
}
})
.collect();
let hostnames = hostname_map.get(&key).cloned().unwrap_or_default();
let endpoints = endpoint_map.get(&key).cloned().unwrap_or_default();
let facts = fact_map.get(&key).cloned().unwrap_or_default();
AgentDeviceWithChildren {
macs,
ips,
hostnames,
endpoints,
facts,
device,
}
})
.collect()
}
pub fn device_key_from_id(device_id: &DeviceId) -> String { pub fn device_key_from_id(device_id: &DeviceId) -> String {
match device_id { match device_id {
DeviceId::Mac(mac) => format!("mac:{}", mac.to_string().to_ascii_lowercase()), DeviceId::Mac(mac) => format!("mac:{}", mac.to_string().to_ascii_lowercase()),
@@ -0,0 +1,270 @@
//! Fleet and single-agent reads share one sequence and a single SQLite read view.
use crate::state::types::*;
use anyhow::{Context, Result};
use sqlx::{SqliteConnection, SqlitePool};
mod assembly;
use assembly::DeviceRows;
/// A read transaction prevents child rows from belonging to a newer snapshot
/// than their parents. Keep scoped SQL explicit so agent reads use indexed filters.
pub(super) async fn load(
pool: &SqlitePool,
agent_id: Option<&str>,
) -> Result<Vec<AgentDeviceWithChildren>> {
let mut tx = pool
.begin()
.await
.context("failed starting device read transaction")?;
let rows = DeviceRows {
devices: load_devices(&mut tx, agent_id).await?,
macs: load_macs(&mut tx, agent_id).await?,
ips: load_ips(&mut tx, agent_id).await?,
hostnames: load_hostnames(&mut tx, agent_id).await?,
endpoints: load_endpoints(&mut tx, agent_id).await?,
facts: load_facts(&mut tx, agent_id).await?,
};
tx.commit()
.await
.context("failed completing device read transaction")?;
Ok(rows.into_devices())
}
async fn load_devices(
connection: &mut SqliteConnection,
agent_id: Option<&str>,
) -> Result<Vec<AgentDeviceRow>> {
match agent_id {
Some(agent_id) => sqlx::query_as!(
AgentDeviceRow,
r#"SELECT agent_id, device_key, presence, display_name,
first_seen_unix, last_seen_unix
FROM agent_devices
WHERE agent_id = ?1
ORDER BY device_key"#,
agent_id,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent devices for agent"),
None => sqlx::query_as!(
AgentDeviceRow,
r#"SELECT agent_id, device_key, presence, display_name,
first_seen_unix, last_seen_unix
FROM agent_devices
ORDER BY agent_id, device_key"#,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent devices"),
}
}
async fn load_macs(
connection: &mut SqliteConnection,
agent_id: Option<&str>,
) -> Result<Vec<AgentDeviceMacRow>> {
match agent_id {
Some(agent_id) => sqlx::query_as!(
AgentDeviceMacRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!"
FROM agent_device_macs
WHERE agent_id = ?1"#,
agent_id,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device macs for agent"),
None => sqlx::query_as!(
AgentDeviceMacRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!"
FROM agent_device_macs"#,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device macs"),
}
}
async fn load_ips(
connection: &mut SqliteConnection,
agent_id: Option<&str>,
) -> Result<Vec<AgentDeviceIpRow>> {
match agent_id {
Some(agent_id) => sqlx::query_as!(
AgentDeviceIpRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!"
FROM agent_device_ips
WHERE agent_id = ?1"#,
agent_id,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device ips for agent"),
None => sqlx::query_as!(
AgentDeviceIpRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!"
FROM agent_device_ips"#,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device ips"),
}
}
async fn load_hostnames(
connection: &mut SqliteConnection,
agent_id: Option<&str>,
) -> Result<Vec<AgentDeviceHostnameRow>> {
match agent_id {
Some(agent_id) => sqlx::query_as!(
AgentDeviceHostnameRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!"
FROM agent_device_hostnames
WHERE agent_id = ?1"#,
agent_id,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device hostnames for agent"),
None => sqlx::query_as!(
AgentDeviceHostnameRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!"
FROM agent_device_hostnames"#,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device hostnames"),
}
}
async fn load_endpoints(
connection: &mut SqliteConnection,
agent_id: Option<&str>,
) -> Result<Vec<AgentDeviceEndpointRow>> {
match agent_id {
Some(agent_id) => sqlx::query_as::<_, AgentDeviceEndpointRow>(
r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip,
hostname, interface, presence, first_seen_unix, last_seen_unix
FROM agent_device_endpoints
WHERE agent_id = ?1"#,
)
.bind(agent_id)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device endpoints for agent"),
None => sqlx::query_as::<_, AgentDeviceEndpointRow>(
r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip,
hostname, interface, presence, first_seen_unix, last_seen_unix
FROM agent_device_endpoints"#,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device endpoints"),
}
}
async fn load_facts(
connection: &mut SqliteConnection,
agent_id: Option<&str>,
) -> Result<Vec<AgentDeviceFactRow>> {
match agent_id {
Some(agent_id) => {
sqlx::query_as!(
AgentDeviceFactRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!"
FROM agent_device_facts
WHERE agent_id = ?1"#,
agent_id,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device facts for agent")
}
None => {
sqlx::query_as!(
AgentDeviceFactRow,
r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!"
FROM agent_device_facts"#,
)
.fetch_all(&mut *connection)
.await
.context("failed listing agent device facts")
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::state::store::helpers::test_helpers::TestStore;
use wakey_core::{Device, NeighborEntry, NeighborState};
fn device(ip: &str) -> Device {
Device::from_parts(
vec![NeighborEntry {
ip: ip.parse().unwrap(),
mac: Some("aa:bb:cc:dd:ee:01".parse().unwrap()),
dev: Some("br-lan".into()),
state: NeighborState::Reachable,
}],
vec![],
)
}
#[tokio::test]
async fn scoped_reads_match_fleet_rows_without_leaking_other_agents() {
let ts = TestStore::new().await;
let store = ts.store();
for (agent, ip) in [("a", "192.168.1.10"), ("b", "192.168.1.20")] {
store
.replace_agent_device_snapshot(agent, &[device(ip)])
.await
.unwrap();
}
let fleet = store.list_agent_device_rows().await.unwrap();
assert_eq!(fleet.len(), 2);
for agent in ["a", "b", "missing"] {
let scoped = store.list_agent_device_rows_for_agent(agent).await.unwrap();
let expected = fleet
.iter()
.filter(|row| row.device.agent_id == agent)
.cloned()
.collect::<Vec<_>>();
// Compare every parent/child field, including timestamps and endpoint metadata.
assert_eq!(format!("{scoped:?}"), format!("{expected:?}"));
}
}
#[tokio::test]
async fn child_reads_share_the_parent_snapshot_during_a_concurrent_write() {
let ts = TestStore::new().await;
let store = ts.store();
store
.replace_agent_device_snapshot("a", &[device("192.168.1.10")])
.await
.unwrap();
let mut tx = store.pool.begin().await.unwrap();
let parents = load_devices(&mut tx, Some("a")).await.unwrap();
assert_eq!(parents.len(), 1);
// In WAL mode the writer can commit while the reader still sees the
// original snapshot. No sleeps or timing assumptions are needed.
store
.replace_agent_device_snapshot("a", &[device("192.168.1.20")])
.await
.unwrap();
let ips = load_ips(&mut tx, Some("a")).await.unwrap();
let endpoints = load_endpoints(&mut tx, Some("a")).await.unwrap();
assert_eq!(ips[0].ip, "192.168.1.10");
assert_eq!(endpoints[0].ip.as_deref(), Some("192.168.1.10"));
tx.commit().await.unwrap();
let current = store.list_agent_device_rows_for_agent("a").await.unwrap();
assert_eq!(
current[0].ips,
["192.168.1.20".parse::<std::net::IpAddr>().unwrap()]
);
}
}
@@ -0,0 +1,198 @@
//! Attach decoded children to their parent using the full agent/device identity.
use std::collections::BTreeMap;
use crate::state::types::*;
#[derive(Default)]
pub(super) struct DeviceRows {
pub(super) devices: Vec<AgentDeviceRow>,
pub(super) macs: Vec<AgentDeviceMacRow>,
pub(super) ips: Vec<AgentDeviceIpRow>,
pub(super) hostnames: Vec<AgentDeviceHostnameRow>,
pub(super) endpoints: Vec<AgentDeviceEndpointRow>,
pub(super) facts: Vec<AgentDeviceFactRow>,
}
type DeviceIndex = BTreeMap<(String, String), AgentDeviceWithChildren>;
impl DeviceRows {
pub(super) fn into_devices(self) -> Vec<AgentDeviceWithChildren> {
let mut index = self
.devices
.into_iter()
.map(|device| {
let key = (device.agent_id.clone(), device.device_key.clone());
(
key,
AgentDeviceWithChildren {
device,
macs: Vec::new(),
ips: Vec::new(),
hostnames: Vec::new(),
endpoints: Vec::new(),
facts: Vec::new(),
},
)
})
.collect::<DeviceIndex>();
add_macs(&mut index, self.macs);
add_ips(&mut index, self.ips);
add_hostnames(&mut index, self.hostnames);
add_endpoints(&mut index, self.endpoints);
add_facts(&mut index, self.facts);
// Matches the SQL order: agent_id/device_key for fleet reads, and
// device_key within one agent. Child row order is left unchanged.
index.into_values().collect()
}
}
fn add_macs(index: &mut DeviceIndex, rows: Vec<AgentDeviceMacRow>) {
for row in rows {
let mac = match macaddr::MacAddr::try_from(&row) {
Ok(mac) => mac,
Err(error) => {
tracing::warn!(%error, agent_id = %row.agent_id, device_key = %row.device_key,
raw_mac = %row.mac, "failed to parse mac from agent_device_macs row");
continue;
}
};
if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) {
device.macs.push(mac);
}
}
}
fn add_ips(index: &mut DeviceIndex, rows: Vec<AgentDeviceIpRow>) {
for row in rows {
let ip = match std::net::IpAddr::try_from(&row) {
Ok(ip) => ip,
Err(error) => {
tracing::warn!(%error, agent_id = %row.agent_id, device_key = %row.device_key,
raw_ip = %row.ip, "failed to parse ip from agent_device_ips row");
continue;
}
};
if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) {
device.ips.push(ip);
}
}
}
fn add_hostnames(index: &mut DeviceIndex, rows: Vec<AgentDeviceHostnameRow>) {
for row in rows {
if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) {
device.hostnames.push(row.hostname);
}
}
}
fn add_endpoints(index: &mut DeviceIndex, rows: Vec<AgentDeviceEndpointRow>) {
for row in rows {
if let Some(endpoint) = row.to_endpoint()
&& let Some(device) = index.get_mut(&(row.agent_id, row.device_key))
{
device.endpoints.push(endpoint);
}
}
}
fn add_facts(index: &mut DeviceIndex, rows: Vec<AgentDeviceFactRow>) {
for row in rows {
if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) {
device.facts.push(row.fact_json);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn parent(agent: &str, key: &str) -> AgentDeviceRow {
AgentDeviceRow {
agent_id: agent.into(),
device_key: key.into(),
presence: "unknown".into(),
display_name: None,
first_seen_unix: 1,
last_seen_unix: 2,
}
}
#[test]
fn assembly_keeps_agent_boundaries_order_and_empty_parents() {
let rows = DeviceRows {
devices: vec![
parent("b", "shared"),
parent("a", "shared"),
parent("a", "empty"),
],
hostnames: vec![
AgentDeviceHostnameRow {
agent_id: "b".into(),
device_key: "shared".into(),
hostname: "second".into(),
},
AgentDeviceHostnameRow {
agent_id: "a".into(),
device_key: "shared".into(),
hostname: "first".into(),
},
AgentDeviceHostnameRow {
agent_id: "missing".into(),
device_key: "shared".into(),
hostname: "orphan".into(),
},
],
facts: vec![AgentDeviceFactRow {
agent_id: "a".into(),
device_key: "shared".into(),
fact_json: "{\"action\":\"remove\"}".into(),
}],
..Default::default()
}
.into_devices();
assert_eq!(rows.len(), 3);
assert_eq!(rows[0].device.device_key, "empty");
assert!(rows[0].hostnames.is_empty());
assert_eq!(rows[1].hostnames, ["first"]);
assert_eq!(rows[1].facts, ["{\"action\":\"remove\"}"]);
assert_eq!(rows[2].hostnames, ["second"]);
assert!(rows[2].facts.is_empty());
}
#[test]
fn malformed_children_do_not_hide_valid_children_or_the_parent() {
let rows = DeviceRows {
devices: vec![parent("a", "pc")],
macs: ["not-a-mac", "aa:bb:cc:dd:ee:01"]
.into_iter()
.map(|mac| AgentDeviceMacRow {
agent_id: "a".into(),
device_key: "pc".into(),
mac: mac.into(),
})
.collect(),
ips: ["not-an-ip", "192.168.1.10"]
.into_iter()
.map(|ip| AgentDeviceIpRow {
agent_id: "a".into(),
device_key: "pc".into(),
ip: ip.into(),
})
.collect(),
..Default::default()
}
.into_devices();
assert_eq!(rows.len(), 1);
assert_eq!(
rows[0].macs,
["aa:bb:cc:dd:ee:01".parse::<macaddr::MacAddr>().unwrap()]
);
assert_eq!(
rows[0].ips,
["192.168.1.10".parse::<std::net::IpAddr>().unwrap()]
);
}
}