/build/source/nativelink-scheduler/src/admin.rs

Source
// Copyright 2026 Trace Machina, Inc. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Read-only views for the admin API: what the scheduler has queued and what
//! it has connected. A provisioner sizes pods from the first and drains
//! through the second, instead of reading the scheduler's store behind its
//! back.

use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

use futures::StreamExt;
use nativelink_error::{Error, ResultExt, make_input_err};
use nativelink_util::action_messages::ActionStage;
use nativelink_util::metrics::record_awaited_action_orphan;
use nativelink_util::operation_state_manager::{
    ClientStateManager, OperationFilter, OperationStageFlags, OrderDirection,
};
use serde::{Deserialize, Serialize};
use tracing::warn;

use crate::resource_sketch::{Family, ResourceSketch, SAMPLER_GENERATION, SketchKey};
use crate::simple_scheduler_state_manager::is_lost_record;
use crate::worker_scheduler::WorkerSummary;

/// A queued action as the admin API lists it: enough to size a pod for it
/// and to tell how long it has waited.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct QueuedDemand {
    pub operation_id: String,
    /// Milliseconds since the epoch when the action was queued.
    pub queued_since_ms: u64,
    pub priority: i32,
    /// The action's properties as the scheduler matches them, reservations
    /// included.
    pub platform_properties: HashMap<String, String>,
    /// Milliseconds since the epoch when a client last checked in on the
    /// operation. A queue entry no client has touched for longer than the
    /// client timeout is one the scheduler will not dispatch.
    pub client_last_seen_ms: Option<u64>,
}

/// Every queued action of the scheduler, highest priority first and then
/// oldest, as the store yields them.
pub async fn queued_demand(scheduler: &dyn ClientStateManager) -> Result<Vec<QueuedDemand>, Error> {
    let mut stream = scheduler
        .filter_operations(OperationFilter {
            stages: OperationStageFlags::Queued,
            // The Redis backend serves the queue in one direction only, the
            // matcher's: highest priority first, then oldest.
            order_by_priority_direction: Some(OrderDirection::Desc),
            ..Default::default()
        })
        .await
        .err_tip(|| "Listing queued operations for the admin API")?0;
    let mut out = Vec::new();
    while let Some(result11) = stream.next().await {
        // On the Redis backend the reads below return the record as the
        // listing read it; on the memory one they are live, and the record
        // can have moved on or gone in between. One that is gone or cannot
        // be decoded is left out and counted, one no longer queued is no
        // longer demand, and a store failure is the listing's error.
        let Some((state8, _)) = still_listed(result.as_state().await)?1 else {
            continue;
        };
        if !matches!1(state.stage, ActionStage::Queued) {
            continue;
        }
        let Some((action_info, _)) = still_listed(result.as_action_info().await)?0 else {
            continue;
        };
        let Some(client_last_seen) = still_listed(result.client_last_seen().await)?0 else {
            continue;
        };
        let client_last_seen_ms = client_last_seen.map(epoch_ms);
        out.push(QueuedDemand {
            operation_id: state.client_operation_id.to_string(),
            queued_since_ms: epoch_ms(action_info.insert_timestamp),
            priority: action_info.priority,
            platform_properties: action_info.platform_properties.clone(),
            client_last_seen_ms,
        });
    }
    Ok(out)
}

/// A read of a listed record: the value, `None` for a record that is gone
/// or cannot be decoded (skipped and counted), or the store's error.
fn still_listed<T>(read: Result<T, Error>) -> Result<Option<T>, Error> {
    match read {
        Ok(value) => Ok(Some(value)),
        Err(err2) if is_lost_record(&err)2 => {
            warn!(
                ?err,
                "Queued operation listed but its record cannot be read; left out of the demand"
            );
            record_awaited_action_orphan("admin");
            Ok(None)
        }
        Err(err) => Err(err).err_tip(|| "Reading a queued operation for the admin API"),
    }
}

fn epoch_ms(at: SystemTime) -> u64 {
    u64::try_from(
        at.duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis(),
    )
    .unwrap_or(u64::MAX)
}

/// The demand listing as the JSON body the admin API serves.
pub async fn queued_demand_json(scheduler: &dyn ClientStateManager) -> Result<String, Error> {
    let demand = queued_demand(scheduler).await?0;
    serde_json::to_string(&demand).map_err(|e| {0
        Error::new(
            nativelink_error::Code::Internal,
            format!("Serializing the demand listing: {e}"),
        )
    })
}

/// The worker listing as the JSON body the admin API serves.
pub fn workers_json(workers: &[WorkerSummary]) -> Result<String, Error> {
    serde_json::to_string(workers).map_err(|e| {0
        Error::new(
            nativelink_error::Code::Internal,
            format!("Serializing the worker listing: {e}"),
        )
    })
}

/// Counts per sketch family.
#[derive(Clone, Copy, Debug, Serialize)]
pub struct FamilyCounts {
    pub target: u64,
    pub mnemonic: u64,
    pub command: u64,
}

impl FamilyCounts {
    fn from_slice<T: Copy + Into<u64>>(values: [T; 3]) -> Self {
        Self {
            target: values[0].into(),
            mnemonic: values[1].into(),
            command: values[2].into(),
        }
    }
}

/// One hint of the listing, in the hints-file shape plus what the sketch
/// knows about it.
#[derive(Debug, Serialize)]
pub struct HintListing {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub target_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub action_mnemonic: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub command_digest: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub class: Option<String>,
    #[serde(skip_serializing_if = "is_zero")]
    pub cpu_count: u64,
    #[serde(skip_serializing_if = "is_zero")]
    pub memory_kb: u64,
    #[serde(skip_serializing_if = "is_zero")]
    pub disk_kb: u64,
    pub samples: u32,
    pub last_seen_s: u64,
    /// `measured`, or `seed` for a hint from the file no attempt has
    /// replaced yet.
    pub source: &'static str,
}

#[expect(
    clippy::trivially_copy_pass_by_ref,
    reason = "serde's skip_serializing_if takes a reference"
)]
const fn is_zero(value: &u64) -> bool {
    *value == 0
}

/// The hints listing: a header with the sketch's size and the shaped hints,
/// most recently used first.
#[derive(Debug, Serialize)]
pub struct HintsListing {
    pub sampler_generation: u32,
    pub window_days: u64,
    pub headroom_percent: u64,
    pub entries: FamilyCounts,
    pub capacity: FamilyCounts,
    pub evictions: FamilyCounts,
    pub samples: FamilyCounts,
    pub probation: u64,
    pub hints: Vec<HintListing>,
}

/// The hints listing as the JSON body the admin API serves. Its `hints`
/// are a valid `hints_file`, so a new fleet can be seeded from an old one.
/// `family` narrows to `target`, `mnemonic` or `command`; `limit` caps the
/// rows, 100,000 by default.
pub fn hints_json(
    sketch: &ResourceSketch,
    family: Option<&str>,
    limit: Option<usize>,
) -> Result<String, Error> {
    let family = match family {
        None | Some("") => None,
        Some("target") => Some(Family::Target),
        Some("mnemonic") => Some(Family::Mnemonic),
        Some("command") => Some(Family::Command),
        Some(other) => {
            return Err(make_input_err!(
                "family must be target, mnemonic or command, not {other:?}"
            ));
        }
    };
    let stats = sketch.stats();
    let hints = sketch
        .hints(family, limit.unwrap_or(100_000))
        .into_iter()
        .map(|row| {
            let (target_id, action_mnemonic, command_digest) = match row.key {
                SketchKey::Target { target, mnemonic } => (
                    Some(target.to_string()),
                    mnemonic.map(|m| m.to_string()),
                    None,
                ),
                SketchKey::Mnemonic { mnemonic } => (None, Some(mnemonic.to_string()), None),
                SketchKey::Command { digest } => (None, None, Some(digest.to_string())),
            };
            HintListing {
                target_id,
                action_mnemonic,
                command_digest,
                class: row.shaped.class,
                cpu_count: row.shaped.cpu_count,
                memory_kb: row.shaped.memory_kb,
                disk_kb: row.shaped.disk_kb,
                samples: row.samples,
                last_seen_s: row.last_seen_s,
                source: if row.seeded { "seed" } else { "measured" },
            }
        })
        .collect();
    let listing = HintsListing {
        sampler_generation: SAMPLER_GENERATION,
        window_days: sketch.window().as_secs() / 86_400,
        headroom_percent: sketch.policy().headroom_percent,
        entries: FamilyCounts::from_slice(
            stats.entries.map(|n| u64::try_from(n).unwrap_or(u64::MAX)),
        ),
        capacity: FamilyCounts::from_slice(
            stats.capacity.map(|n| u64::try_from(n).unwrap_or(u64::MAX)),
        ),
        evictions: FamilyCounts::from_slice(stats.evictions),
        samples: FamilyCounts::from_slice(stats.samples),
        probation: u64::try_from(stats.probation).unwrap_or(u64::MAX),
        hints,
    };
    serde_json::to_string(&listing).map_err(|e| {
        Error::new(
            nativelink_error::Code::Internal,
            format!("Serializing the hints listing: {e}"),
        )
    })
}

Coverage Report

Created: 2026-11-13 23:29

Line	Count	Source
1		// Copyright 2026 Trace Machina, Inc. All rights reserved.
2		//
3		// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
4		// you may not use this file except in compliance with the License.
5		// You may obtain a copy of the License at
6		//
7		// See LICENSE file for details
8		//
9		// Unless required by applicable law or agreed to in writing, software
10		// distributed under the License is distributed on an "AS IS" BASIS,
11		// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12		// See the License for the specific language governing permissions and
13		// limitations under the License.
14
15		//! Read-only views for the admin API: what the scheduler has queued and what
16		//! it has connected. A provisioner sizes pods from the first and drains
17		//! through the second, instead of reading the scheduler's store behind its
18		//! back.
19
20		use std::collections::HashMap;
21		use std::time::{SystemTime, UNIX_EPOCH};
22
23		use futures::StreamExt;
24		use nativelink_error::{Error, ResultExt, make_input_err};
25		use nativelink_util::action_messages::ActionStage;
26		use nativelink_util::metrics::record_awaited_action_orphan;
27		use nativelink_util::operation_state_manager::{
28		ClientStateManager, OperationFilter, OperationStageFlags, OrderDirection,
29		};
30		use serde::{Deserialize, Serialize};
31		use tracing::warn;
32
33		use crate::resource_sketch::{Family, ResourceSketch, SAMPLER_GENERATION, SketchKey};
34		use crate::simple_scheduler_state_manager::is_lost_record;
35		use crate::worker_scheduler::WorkerSummary;
36
37		/// A queued action as the admin API lists it: enough to size a pod for it
38		/// and to tell how long it has waited.
39		#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40		pub struct QueuedDemand {
41		pub operation_id: String,
42		/// Milliseconds since the epoch when the action was queued.
43		pub queued_since_ms: u64,
44		pub priority: i32,
45		/// The action's properties as the scheduler matches them, reservations
46		/// included.
47		pub platform_properties: HashMap<String, String>,
48		/// Milliseconds since the epoch when a client last checked in on the
49		/// operation. A queue entry no client has touched for longer than the
50		/// client timeout is one the scheduler will not dispatch.
51		pub client_last_seen_ms: Option<u64>,
52		}
53
54		/// Every queued action of the scheduler, highest priority first and then
55		/// oldest, as the store yields them.
56	6	pub async fn queued_demand(scheduler: &dyn ClientStateManager) -> Result<Vec<QueuedDemand>, Error> {
57	6	let mut stream = scheduler
58	6	.filter_operations(OperationFilter {
59	6	stages: OperationStageFlags::Queued,
60	6	// The Redis backend serves the queue in one direction only, the
61	6	// matcher's: highest priority first, then oldest.
62	6	order_by_priority_direction: Some(OrderDirection::Desc),
63	6	..Default::default()
64	6	})
65	6	.await
66	6	.err_tip(\|\| "Listing queued operations for the admin API") ?0 ;
67	6	let mut out = Vec::new();
68	16	while let Some( result11 ) = stream.next().await {
69		// On the Redis backend the reads below return the record as the
70		// listing read it; on the memory one they are live, and the record
71		// can have moved on or gone in between. One that is gone or cannot
72		// be decoded is left out and counted, one no longer queued is no
73		// longer demand, and a store failure is the listing's error.
74	11	let Some(( state8 , _)) = still_listed(result.as_state().await) ?1 else {
75	2	continue;
76		};
77	8	if ! matches!1 (state.stage, ActionStage::Queued) {
78	1	continue;
79	7	}
80	7	let Some((action_info, _)) = still_listed(result.as_action_info().await) ?0 else {
81	0	continue;
82		};
83	7	let Some(client_last_seen) = still_listed(result.client_last_seen().await) ?0 else {
84	0	continue;
85		};
86	7	let client_last_seen_ms = client_last_seen.map(epoch_ms);
87	7	out.push(QueuedDemand {
88	7	operation_id: state.client_operation_id.to_string(),
89	7	queued_since_ms: epoch_ms(action_info.insert_timestamp),
90	7	priority: action_info.priority,
91	7	platform_properties: action_info.platform_properties.clone(),
92	7	client_last_seen_ms,
93	7	});
94		}
95	5	Ok(out)
96	6	}
97
98		/// A read of a listed record: the value, `None` for a record that is gone
99		/// or cannot be decoded (skipped and counted), or the store's error.
100	25	fn still_listed<T>(read: Result<T, Error>) -> Result<Option<T>, Error> {
101	3	match read {
102	22	Ok(value) => Ok(Some(value)),
103	3	Err( err2 ) if is_lost_record(&err )2 => {
104	2	warn!(
105		?err,
106		"Queued operation listed but its record cannot be read; left out of the demand"
107		);
108	2	record_awaited_action_orphan("admin");
109	2	Ok(None)
110		}
111	1	Err(err) => Err(err).err_tip(\|\| "Reading a queued operation for the admin API"),
112		}
113	25	}
114
115	11	fn epoch_ms(at: SystemTime) -> u64 {
116	11	u64::try_from(
117	11	at.duration_since(UNIX_EPOCH)
118	11	.unwrap_or_default()
119	11	.as_millis(),
120		)
121	11	.unwrap_or(u64::MAX)
122	11	}
123
124		/// The demand listing as the JSON body the admin API serves.
125	1	pub async fn queued_demand_json(scheduler: &dyn ClientStateManager) -> Result<String, Error> {
126	1	let demand = queued_demand(scheduler).await ?0 ;
127	1	serde_json::to_string(&demand).map_err(\|e\| {0
128	0	Error::new(
129	0	nativelink_error::Code::Internal,
130	0	format!("Serializing the demand listing: {e}"),
131		)
132	0	})
133	1	}
134
135		/// The worker listing as the JSON body the admin API serves.
136	1	pub fn workers_json(workers: &[WorkerSummary]) -> Result<String, Error> {
137	1	serde_json::to_string(workers).map_err(\|e\| {0
138	0	Error::new(
139	0	nativelink_error::Code::Internal,
140	0	format!("Serializing the worker listing: {e}"),
141		)
142	0	})
143	1	}
144
145		/// Counts per sketch family.
146		#[derive(Clone, Copy, Debug, Serialize)]
147		pub struct FamilyCounts {
148		pub target: u64,
149		pub mnemonic: u64,
150		pub command: u64,
151		}
152
153		impl FamilyCounts {
154	0	fn from_slice<T: Copy + Into<u64>>(values: [T; 3]) -> Self {
155	0	Self {
156	0	target: values[0].into(),
157	0	mnemonic: values[1].into(),
158	0	command: values[2].into(),
159	0	}
160	0	}
161		}
162
163		/// One hint of the listing, in the hints-file shape plus what the sketch
164		/// knows about it.
165		#[derive(Debug, Serialize)]
166		pub struct HintListing {
167		#[serde(skip_serializing_if = "Option::is_none")]
168		pub target_id: Option<String>,
169		#[serde(skip_serializing_if = "Option::is_none")]
170		pub action_mnemonic: Option<String>,
171		#[serde(skip_serializing_if = "Option::is_none")]
172		pub command_digest: Option<String>,
173		#[serde(skip_serializing_if = "Option::is_none")]
174		pub class: Option<String>,
175		#[serde(skip_serializing_if = "is_zero")]
176		pub cpu_count: u64,
177		#[serde(skip_serializing_if = "is_zero")]
178		pub memory_kb: u64,
179		#[serde(skip_serializing_if = "is_zero")]
180		pub disk_kb: u64,
181		pub samples: u32,
182		pub last_seen_s: u64,
183		/// `measured`, or `seed` for a hint from the file no attempt has
184		/// replaced yet.
185		pub source: &'static str,
186		}
187
188		#[expect(
189		clippy::trivially_copy_pass_by_ref,
190		reason = "serde's skip_serializing_if takes a reference"
191		)]
192	0	const fn is_zero(value: &u64) -> bool {
193	0	*value == 0
194	0	}
195
196		/// The hints listing: a header with the sketch's size and the shaped hints,
197		/// most recently used first.
198		#[derive(Debug, Serialize)]
199		pub struct HintsListing {
200		pub sampler_generation: u32,
201		pub window_days: u64,
202		pub headroom_percent: u64,
203		pub entries: FamilyCounts,
204		pub capacity: FamilyCounts,
205		pub evictions: FamilyCounts,
206		pub samples: FamilyCounts,
207		pub probation: u64,
208		pub hints: Vec<HintListing>,
209		}
210
211		/// The hints listing as the JSON body the admin API serves. Its `hints`
212		/// are a valid `hints_file`, so a new fleet can be seeded from an old one.
213		/// `family` narrows to `target`, `mnemonic` or `command`; `limit` caps the
214		/// rows, 100,000 by default.
215	0	pub fn hints_json(
216	0	sketch: &ResourceSketch,
217	0	family: Option<&str>,
218	0	limit: Option<usize>,
219	0	) -> Result<String, Error> {
220	0	let family = match family {
221	0	None \| Some("") => None,
222	0	Some("target") => Some(Family::Target),
223	0	Some("mnemonic") => Some(Family::Mnemonic),
224	0	Some("command") => Some(Family::Command),
225	0	Some(other) => {
226	0	return Err(make_input_err!(
227	0	"family must be target, mnemonic or command, not {other:?}"
228	0	));
229		}
230		};
231	0	let stats = sketch.stats();
232	0	let hints = sketch
233	0	.hints(family, limit.unwrap_or(100_000))
234	0	.into_iter()
235	0	.map(\|row\| {
236	0	let (target_id, action_mnemonic, command_digest) = match row.key {
237	0	SketchKey::Target { target, mnemonic } => (
238	0	Some(target.to_string()),
239	0	mnemonic.map(\|m\| m.to_string()),
240	0	None,
241		),
242	0	SketchKey::Mnemonic { mnemonic } => (None, Some(mnemonic.to_string()), None),
243	0	SketchKey::Command { digest } => (None, None, Some(digest.to_string())),
244		};
245		HintListing {
246	0	target_id,
247	0	action_mnemonic,
248	0	command_digest,
249	0	class: row.shaped.class,
250	0	cpu_count: row.shaped.cpu_count,
251	0	memory_kb: row.shaped.memory_kb,
252	0	disk_kb: row.shaped.disk_kb,
253	0	samples: row.samples,
254	0	last_seen_s: row.last_seen_s,
255	0	source: if row.seeded { "seed" } else { "measured" },
256		}
257	0	})
258	0	.collect();
259	0	let listing = HintsListing {
260		sampler_generation: SAMPLER_GENERATION,
261	0	window_days: sketch.window().as_secs() / 86_400,
262	0	headroom_percent: sketch.policy().headroom_percent,
263	0	entries: FamilyCounts::from_slice(
264	0	stats.entries.map(\|n\| u64::try_from(n).unwrap_or(u64::MAX)),
265		),
266	0	capacity: FamilyCounts::from_slice(
267	0	stats.capacity.map(\|n\| u64::try_from(n).unwrap_or(u64::MAX)),
268		),
269	0	evictions: FamilyCounts::from_slice(stats.evictions),
270	0	samples: FamilyCounts::from_slice(stats.samples),
271	0	probation: u64::try_from(stats.probation).unwrap_or(u64::MAX),
272	0	hints,
273		};
274	0	serde_json::to_string(&listing).map_err(\|e\| {
275	0	Error::new(
276	0	nativelink_error::Code::Internal,
277	0	format!("Serializing the hints listing: {e}"),
278		)
279	0	})
280	0	}