fix(fleet): separate wake selection and remove capped device lookup

This commit is contained in:
lda committed 2026-10-11 19:56:55 +07:00 Verified
1 parent a04df6f573
commit 08cb41ce73
3 files changed
+407 -99

No files matched your search

+7 -4
View File
@@ -1,6 +1,6 @@
# Fleet View Cache # Fleet View Cache
CC retains one complete, grouped fleet view in memory. Requests still apply CC retains one complete, grouped fleet view in memory. List requests still apply
their own search, visibility, presence, agent filters and result limit. Neither their own search, visibility, presence, agent filters and result limit. Neither
the database nor the response format changes. the database nor the response format changes.
@@ -22,7 +22,9 @@ list/wake requests return HTTP 503 with code `fleet_busy`; clients can retry
later. No stale view or command is used as a fallback. This bounds the number later. No stale view or command is used as a fallback. This bounds the number
of rebuilds per reader, not database execution time or total concurrent requests. of rebuilds per reader, not database execution time or total concurrent requests.
Wake uses the same view, but command dispatch still checks the current session; Wake searches the complete cached view by device key, without list visibility
filters or the list API's 1,000-row limit. It borrows the view and clones only the
selected route. Command dispatch still checks the current session;
a cached route cannot grant access to a disconnected agent. a cached route cannot grant access to a disconnected agent.
## External Writes ## External Writes
@@ -41,5 +43,6 @@ keep their existing timer-driven evaluator rather than being forced into this
lazy cache. lazy cache.
`fleet_load` spans remain on requests. A cache miss adds `fleet_rebuild`, with `fleet_load` spans remain on requests. A cache miss adds `fleet_rebuild`, with
database loads and `fleet_group_devices` nested below it. A cache hit performs database loads and `fleet_group_devices` nested below it. On a cache hit, list
only key checking, cloning the view, and request-specific filtering. requests clone/filter the view; wake requests look up their device in the borrowed
view. Neither path repeats the database reads or grouping.
@@ -6,7 +6,7 @@ use axum::http::StatusCode;
use axum::response::IntoResponse; use axum::response::IntoResponse;
use futures_util::StreamExt; use futures_util::StreamExt;
use tracing::{Instrument, info, warn}; use tracing::{Instrument, info, warn};
use wakey_agent::protocol::{AgentCommand, InventoryRequest, WakeRequest}; use wakey_agent::protocol::{AgentCommand, InventoryRequest};
use crate::api::ApiError; use crate::api::ApiError;
use crate::api::commands::{RelayCommandResponse, relay_agent_command}; use crate::api::commands::{RelayCommandResponse, relay_agent_command};
@@ -14,6 +14,8 @@ use crate::runtime::AppState;
mod build; mod build;
mod types; mod types;
mod wake;
pub use wake::wake_fleet_device;
#[cfg(test)] #[cfg(test)]
mod tests; mod tests;
@@ -24,7 +26,7 @@ use build::{
}; };
use types::{ use types::{
FleetDevice, ListFleetDevicesQuery, RefreshFleetAgentResult, RefreshFleetDevicesRequest, FleetDevice, ListFleetDevicesQuery, RefreshFleetAgentResult, RefreshFleetDevicesRequest,
RefreshFleetDevicesResponse, WakeFleetDeviceRequest, WakeFleetDeviceResponse, RefreshFleetDevicesResponse,
}; };
pub(crate) type FleetCache = pub(crate) type FleetCache =
@@ -249,99 +251,6 @@ async fn refresh_fleet_agent_result(
} }
} }
#[tracing::instrument(name = "fleet_wake", skip_all, fields(explicit_route = req.route_id.is_some(), outcome = tracing::field::Empty))]
pub async fn wake_fleet_device(
State(state): State<AppState>,
Json(req): Json<WakeFleetDeviceRequest>,
) -> Result<impl IntoResponse, ApiError> {
// Early database/relay errors retain an outcome even without a command result.
tracing::Span::current().record("outcome", "error");
let query = ListFleetDevicesQuery {
query: None,
presence: None,
known: None,
agent_id: None,
visibility: Some("all".into()),
limit: Some(1000),
};
let devices = load_fleet_devices(&state, &query).await.map_err(|err| {
warn!(error = %err, "failed loading fleet devices for wake");
fleet_error(err, "load_fleet_devices_failed")
})?;
let device = devices
.into_iter()
.find(|device| device.device_key == req.device_key)
.ok_or_else(|| {
tracing::Span::current().record("outcome", "device_not_found");
ApiError::new(
StatusCode::NOT_FOUND,
"fleet_device_not_found",
"fleet device not found",
)
})?;
let route = match req.route_id.as_deref() {
Some(route_id) => device
.route_candidates
.into_iter()
.find(|route| route.route_id == route_id),
None => device.recommended_route,
}
.ok_or_else(|| {
tracing::Span::current().record("outcome", "route_unavailable");
ApiError::new(
StatusCode::BAD_REQUEST,
"wake_route_unavailable",
"no wakeable connected MAC-backed route is available",
)
})?;
let Some(mac) = route.mac else {
tracing::Span::current().record("outcome", "route_without_mac");
return Err(ApiError::new(
StatusCode::BAD_REQUEST,
"wake_route_unavailable",
"selected route does not include a MAC address",
));
};
if !route.connected {
tracing::Span::current().record("outcome", "route_disconnected");
return Err(ApiError::new(
StatusCode::BAD_REQUEST,
"wake_route_unavailable",
"selected route agent is not connected",
));
}
if !route.wakeable {
tracing::Span::current().record("outcome", "route_not_wakeable");
return Err(ApiError::new(
StatusCode::BAD_REQUEST,
"wake_route_unavailable",
"selected route is not wakeable",
));
}
info!(agent_id = %route.agent_id, has_ip = route.ip.is_some(), explicit_route = req.route_id.is_some(), "fleet wake route selected");
let command = relay_agent_command(
&state,
&route.agent_id,
AgentCommand::Wake(WakeRequest {
query: None,
mac: Some(mac),
ip: route.ip,
}),
req.timeout_ms,
)
.await?;
tracing::Span::current().record("outcome", command.status.as_str());
Ok((
StatusCode::OK,
Json(WakeFleetDeviceResponse { route, command }),
))
}
#[tracing::instrument(name = "fleet_load", skip_all, fields(grouped = tracing::field::Empty, filtered = tracing::field::Empty, returned = tracing::field::Empty))] #[tracing::instrument(name = "fleet_load", skip_all, fields(grouped = tracing::field::Empty, filtered = tracing::field::Empty, returned = tracing::field::Empty))]
async fn load_fleet_devices( async fn load_fleet_devices(
state: &AppState, state: &AppState,
@@ -0,0 +1,396 @@
use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
use macaddr::MacAddr;
use tracing::{Instrument, info, warn};
use wakey_agent::protocol::{AgentCommand, WakeRequest};
use super::types::{FleetDevice, FleetWakeRoute, WakeFleetDeviceRequest, WakeFleetDeviceResponse};
use super::{fleet_error, fleet_view};
use crate::api::{ApiError, commands::relay_agent_command};
use crate::runtime::AppState;
#[tracing::instrument(name = "fleet_wake", skip_all, fields(explicit_route = req.route_id.is_some(), outcome = tracing::field::Empty))]
pub async fn wake_fleet_device(
State(state): State<AppState>,
Json(req): Json<WakeFleetDeviceRequest>,
) -> Result<impl IntoResponse, ApiError> {
// Early database/relay errors retain an outcome even without a command result.
tracing::Span::current().record("outcome", "error");
let selected = load_selected_route(&state, &req.device_key, req.route_id.as_deref()).await?;
info!(agent_id = %selected.route.agent_id, has_ip = selected.route.ip.is_some(),
explicit_route = req.route_id.is_some(), "fleet wake route selected");
let response = dispatch_wake(&state, selected, req.timeout_ms).await?;
tracing::Span::current().record("outcome", response.command.status.as_str());
Ok((StatusCode::OK, Json(response)))
}
/// List visibility and pagination are presentation policy, not wake eligibility.
/// Borrow the complete cached view and clone only the selected route.
async fn load_selected_route(
state: &AppState,
device_key: &str,
route_id: Option<&str>,
) -> Result<SelectedRoute, ApiError> {
let span = tracing::info_span!("fleet_load", grouped = tracing::field::Empty);
let devices = fleet_view(state)
.instrument(span.clone())
.await
.map_err(|err| {
warn!(error = %err, "failed loading fleet devices for wake");
fleet_error(err, "load_fleet_devices_failed")
})?;
span.record("grouped", devices.len());
select_route(&devices, device_key, route_id).map_err(|error| {
tracing::Span::current().record("outcome", error.outcome());
error.into_api_error()
})
}
/// Selection does not rank routes: the fleet builder already chose the recommendation.
/// An explicit route never falls back, including when it is disconnected or IP-only.
fn select_route(
devices: &[FleetDevice],
device_key: &str,
route_id: Option<&str>,
) -> Result<SelectedRoute, RouteError> {
let device = devices
.iter()
.find(|device| device.device_key == device_key)
.ok_or(RouteError::DeviceNotFound)?;
let route = match route_id {
Some(id) => device
.route_candidates
.iter()
.find(|route| route.route_id == id),
None => device.recommended_route.as_ref(),
}
.ok_or(RouteError::Unavailable)?;
let mac = route.mac.ok_or(RouteError::MissingMac)?;
if !route.connected {
return Err(RouteError::Disconnected);
}
if !route.wakeable {
return Err(RouteError::NotWakeable);
}
Ok(SelectedRoute {
route: route.clone(),
mac,
})
}
/// Only the selector constructs this; dispatch cannot accidentally use an IP-only route.
#[derive(Debug)]
struct SelectedRoute {
route: FleetWakeRoute,
mac: MacAddr,
}
impl SelectedRoute {
fn command(&self) -> AgentCommand {
AgentCommand::Wake(WakeRequest {
query: None,
mac: Some(self.mac),
ip: self.route.ip,
})
}
}
async fn dispatch_wake(
state: &AppState,
selected: SelectedRoute,
timeout_ms: Option<u64>,
) -> Result<WakeFleetDeviceResponse, ApiError> {
// The relay checks the current agent connection again; cached connectivity
// is only selection evidence, not a guarantee that the agent is still there.
let command = relay_agent_command(
state,
&selected.route.agent_id,
selected.command(),
timeout_ms,
)
.await?;
Ok(WakeFleetDeviceResponse {
route: selected.route,
command,
})
}
#[derive(Debug, PartialEq, Eq)]
enum RouteError {
DeviceNotFound,
Unavailable,
MissingMac,
Disconnected,
NotWakeable,
}
impl RouteError {
fn outcome(&self) -> &'static str {
match self {
Self::DeviceNotFound => "device_not_found",
Self::Unavailable => "route_unavailable",
Self::MissingMac => "route_without_mac",
Self::Disconnected => "route_disconnected",
Self::NotWakeable => "route_not_wakeable",
}
}
fn into_api_error(self) -> ApiError {
let message = match self {
Self::DeviceNotFound => {
return ApiError::new(
StatusCode::NOT_FOUND,
"fleet_device_not_found",
"fleet device not found",
);
}
Self::Unavailable => "no wakeable connected MAC-backed route is available",
Self::MissingMac => "selected route does not include a MAC address",
Self::Disconnected => "selected route agent is not connected",
Self::NotWakeable => "selected route is not wakeable",
};
ApiError::new(StatusCode::BAD_REQUEST, "wake_route_unavailable", message)
}
}
#[cfg(test)]
mod tests {
use super::super::{load_fleet_devices, types::ListFleetDevicesQuery};
use super::*;
use wakey_core::Presence;
#[tokio::test]
async fn wake_lookup_can_select_a_device_beyond_the_list_limit() {
use std::collections::{BTreeSet, HashMap};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{Mutex, RwLock};
let dir = std::env::temp_dir().join(format!("wakey-wake-limit-{}", uuid::Uuid::new_v4()));
let state = AppState {
store: Arc::new(
crate::state::Store::load_or_init(
&dir.join("state.sqlite3"),
vec![],
Duration::from_secs(60),
)
.await
.unwrap(),
),
sessions: Arc::new(RwLock::new(HashMap::new())),
pending: Arc::new(Mutex::new(HashMap::new())),
public_url: "http://localhost".into(),
command_timeout: Duration::from_secs(30),
enroll_token_ttl: Duration::from_secs(60),
terminals: crate::runtime::terminals::TerminalRegistry::new(),
alerts: crate::api::AlertsHub::new(),
fleet: Arc::new(super::super::FleetCache::default()),
};
// Seed the complete cached view, as the builder would. Exercise both
// production read paths without thousands of unrelated SQL inserts.
state
.fleet
.get_or_try_init(
|| async { (state.store.fleet_revision(), BTreeSet::new()) },
|| async {
Ok((0..1001)
.map(|i| {
let mut row = device(vec![route("valid")], Some(route("valid")));
row.device_key = format!("pc-{i}");
row
})
.collect())
},
)
.await
.unwrap();
let listed = load_fleet_devices(
&state,
&ListFleetDevicesQuery {
visibility: Some("all".into()),
limit: Some(1000),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(listed.len(), 1000);
assert!(listed.iter().all(|row| row.device_key != "pc-1000"));
for route_id in [None, Some("valid")] {
let selected = load_selected_route(&state, "pc-1000", route_id)
.await
.unwrap();
assert_eq!(selected.route.route_id, "valid");
}
drop(state);
std::fs::remove_dir_all(dir).unwrap();
}
fn route(id: &str) -> FleetWakeRoute {
FleetWakeRoute {
route_id: id.into(),
agent_id: "agent-a".into(),
nickname: None,
connected: true,
mac: Some("aa:bb:cc:dd:ee:01".parse().unwrap()),
ip: Some("192.168.1.10".parse().unwrap()),
hostname: None,
interface: None,
source: "neighbor".into(),
presence: Presence::Offline,
last_seen_unix: 1,
wakeable: true,
}
}
fn device(candidates: Vec<FleetWakeRoute>, recommended: Option<FleetWakeRoute>) -> FleetDevice {
FleetDevice {
device_key: "pc".into(),
display_name: "pc".into(),
known_device: None,
pinned: false,
ips: vec![],
macs: vec![],
hostnames: vec![],
agents: vec![],
sources: vec![],
endpoints: vec![],
first_seen_unix: None,
last_seen_unix: None,
presence: Presence::Offline,
route_candidates: candidates,
recommended_route: recommended,
}
}
#[test]
fn selection_uses_recommendation_or_exact_override_and_preserves_address_pair() {
let first = route("first");
let mut recommended = route("recommended");
recommended.agent_id = "agent-b".into();
recommended.ip = Some("192.168.2.20".parse().unwrap());
let devices = [device(vec![first, recommended.clone()], Some(recommended))];
for (override_id, route_id, agent, ip) in [
(None, "recommended", "agent-b", "192.168.2.20"),
(Some("first"), "first", "agent-a", "192.168.1.10"),
] {
let selected = select_route(&devices, "pc", override_id).unwrap();
assert_eq!(selected.route.route_id, route_id);
assert_eq!(selected.route.agent_id, agent);
let AgentCommand::Wake(command) = selected.command() else {
panic!("not a wake command")
};
assert!(command.query.is_none());
assert_eq!(command.mac, selected.route.mac);
assert_eq!(command.ip, Some(ip.parse().unwrap()));
}
}
#[test]
fn missing_device_and_missing_routes_keep_distinct_errors() {
let valid = route("valid");
let devices = [device(vec![valid.clone()], Some(valid))];
assert_eq!(
select_route(&devices, "missing", None).unwrap_err(),
RouteError::DeviceNotFound
);
assert_eq!(
select_route(&devices, "pc", Some("missing")).unwrap_err(),
RouteError::Unavailable
);
let devices = [device(vec![route("valid")], None)];
assert_eq!(
select_route(&devices, "pc", None).unwrap_err(),
RouteError::Unavailable
);
}
#[test]
fn invalid_explicit_routes_do_not_fall_back_and_validation_order_is_preserved() {
for (mac, connected, wakeable, expected) in [
(None, false, false, RouteError::MissingMac),
(
Some("aa:bb:cc:dd:ee:01"),
false,
false,
RouteError::Disconnected,
),
(
Some("aa:bb:cc:dd:ee:01"),
true,
false,
RouteError::NotWakeable,
),
] {
let mut invalid = route("invalid");
invalid.mac = mac.map(|value| value.parse().unwrap());
invalid.connected = connected;
invalid.wakeable = wakeable;
let devices = [device(vec![invalid, route("valid")], Some(route("valid")))];
assert_eq!(
select_route(&devices, "pc", Some("invalid")).unwrap_err(),
expected
);
}
}
#[test]
fn mac_only_offline_route_is_still_wakeable() {
let mut mac_only = route("mac-only");
mac_only.ip = None;
let devices = [device(vec![mac_only.clone()], Some(mac_only))];
let selected = select_route(&devices, "pc", None).unwrap();
let AgentCommand::Wake(command) = selected.command() else {
panic!("not a wake command")
};
assert!(command.mac.is_some());
assert!(command.ip.is_none());
}
#[test]
fn route_errors_preserve_http_codes_messages_and_trace_outcomes() {
for (error, status, code, message, outcome) in [
(
RouteError::DeviceNotFound,
StatusCode::NOT_FOUND,
"fleet_device_not_found",
"fleet device not found",
"device_not_found",
),
(
RouteError::Unavailable,
StatusCode::BAD_REQUEST,
"wake_route_unavailable",
"no wakeable connected MAC-backed route is available",
"route_unavailable",
),
(
RouteError::MissingMac,
StatusCode::BAD_REQUEST,
"wake_route_unavailable",
"selected route does not include a MAC address",
"route_without_mac",
),
(
RouteError::Disconnected,
StatusCode::BAD_REQUEST,
"wake_route_unavailable",
"selected route agent is not connected",
"route_disconnected",
),
(
RouteError::NotWakeable,
StatusCode::BAD_REQUEST,
"wake_route_unavailable",
"selected route is not wakeable",
"route_not_wakeable",
),
] {
assert_eq!(error.outcome(), outcome);
let response = error.into_api_error();
assert_eq!(response.status, status);
assert_eq!(response.code, code);
assert_eq!(response.message, message);
}
}
}