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

Source
// Copyright 2024 The NativeLink Authors. 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.

//! Worker capability index for fast worker matching.
//!
//! This module provides an index that accelerates worker matching by property.
//! Instead of iterating all workers for each action, we maintain an inverted index
//! that maps property values to sets of workers that have those values.
//!
//! ## Complexity Analysis
//!
//! Without index: O(W × P) where W = workers, P = properties per action
//! With index: O(P × log(W)) for exact properties + O(W' × P') for minimum properties
//!   where W' = filtered workers, P' = minimum property count (typically small)
//!
//! For typical workloads (few minimum properties), this reduces matching from
//! O(n × m) to approximately O(log n).

use std::collections::{HashMap, HashSet};

use nativelink_util::action_messages::WorkerId;
use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue};
use tracing::info;

/// A property key-value pair used for indexing.
#[derive(Clone, Hash, Eq, PartialEq, Debug)]
struct PropertyKey {
    name: String,
    value: PlatformPropertyValue,
}

/// Index structure for fast worker capability lookup.
///
/// Maintains an inverted index from property values to worker IDs.
/// Only indexes `Exact` and `Priority` properties since `Minimum` properties
/// are dynamic and require runtime comparison.
#[derive(Debug, Default)]
pub struct WorkerCapabilityIndex {
    /// Maps `(property_name, property_value)` -> Set of worker IDs with that property.
    /// Only contains `Exact` and `Priority` properties.
    exact_index: HashMap<PropertyKey, HashSet<WorkerId>>,

    /// Maps `property_name` -> Set of worker IDs that have this property (any value).
    /// Used for fast "has property" checks for `Priority` and `Minimum` properties.
    property_presence: HashMap<String, HashSet<WorkerId>>,

    /// Set of all indexed worker IDs.
    all_workers: HashSet<WorkerId>,
}

impl WorkerCapabilityIndex {
    /// Creates a new empty capability index.
    pub fn new() -> Self {
        Self::default()
    }

    /// Adds a worker to the index with their platform properties.
    pub fn add_worker(&mut self, worker_id: &WorkerId, properties: &PlatformProperties) {
        self.all_workers.insert(worker_id.clone());

        for (name32, value32) in &properties.properties {
            // Track property presence
            self.property_presence
                .entry(name.clone())
                .or_default()
                .insert(worker_id.clone());

            match value {
                PlatformPropertyValue::Exact(_)
                | PlatformPropertyValue::Priority(_)
                | PlatformPropertyValue::Unknown(_) => {
                    // Index exact-match properties
                    let key = PropertyKey {
                        name: name.clone(),
                        value: value.clone(),
                    };
                    self.exact_index
                        .entry(key)
                        .or_default()
                        .insert(worker_id.clone());
                }
                PlatformPropertyValue::Minimum(_) | PlatformPropertyValue::Ignore(_) => {
                    // Minimum properties are tracked via `property_presence` only.
                    // Their actual values are checked at runtime since they're dynamic.

                    // Ignore properties we just drop
                }
            }
        }
    }

    /// Removes a worker from the index.
    pub fn remove_worker(&mut self, worker_id: &WorkerId) {
        self.all_workers.remove(worker_id);

        // Remove from exact index
        self.exact_index.retain(|_, workers| {1
            workers.remove(worker_id);
            !workers.is_empty()
        });

        // Remove from presence index
        self.property_presence.retain(|_, workers| {3
            workers.remove(worker_id);
            !workers.is_empty()
        });
    }

    /// Finds workers that can satisfy the given action properties.
    ///
    /// Returns a set of worker IDs that match all required properties.
    /// The caller should apply additional filtering (e.g., worker availability).
    ///
    /// IMPORTANT: This method returns candidates based on STATIC properties only.
    /// - Exact and Unknown properties are fully matched; Unknown properties
    ///   additionally match workers that do not declare the key
    /// - Priority properties just require the key to exist
    /// - Minimum properties return workers that HAVE the property (presence check only)
    ///
    /// The caller MUST still verify Minimum property values at runtime because
    /// worker resources change dynamically as jobs are assigned/completed.
    pub fn find_matching_workers(
        &self,
        action_properties: &PlatformProperties,
        full_worker_logging: bool,
    ) -> HashSet<WorkerId> {
        if self.all_workers.is_empty() {
            if full_worker_logging {
                info!("No workers available to match!");
            }
            return HashSet::new();
        }

        if action_properties.properties.is_empty() {
            // No properties required, all workers match
            return self.all_workers.clone();
        }

        let mut candidates: Option<HashSet<WorkerId>> = None;

        for (name, value) in &action_properties.properties32 {
            match value {
                PlatformPropertyValue::Exact(_) | PlatformPropertyValue::Unknown(_) => {
                    // Look up workers with exact match
                    let key = PropertyKey {
                        name: name.clone(),
                        value: value.clone(),
                    };

                    let mut matching = self.exact_index.get(&key).cloned().unwrap_or_default();

                    // Unknown properties only restrict workers that declare the
                    // key, so workers without the key remain candidates.
                    if matches!6(value, PlatformPropertyValue::Unknown(_)) {
                        match self.property_presence.get(name) {
                            Some(with_key) => {
                                matching.extend(self.all_workers.difference(with_key).cloned());
                            }
                            None => matching.extend(self.all_workers.iter().cloned()),
                        }
                    }

                    let internal_candidates = match candidates {
                        Some(existing) => existing.intersection(&matching).cloned().collect(),
                        None => matching,
                    };

                    // Early exit if no candidates
                    if internal_candidates.is_empty() {
                        if full_worker_logging {
                            let values: Vec<_> = self
                                .exact_index
                                .iter()
                                .filter(|pk| &pk.0.name == name)
                                .map(|pk| pk.0.value.clone())
                                .collect();
                            info!(
                                "No candidate workers due to a lack of matching '{name}' = {value:?}. Workers have: {values:?}"
                            );
                        }
                        return HashSet::new();
                    }
                    candidates = Some(internal_candidates);
                }
                PlatformPropertyValue::Priority(_) | PlatformPropertyValue::Minimum(_) => {
                    // Priority: just requires the key to exist
                    // Minimum: worker must have the property (value checked at runtime by caller)
                    // We only check presence here because Minimum values are DYNAMIC -
                    // they change as jobs are assigned to workers.
                    let workers_with_property = self
                        .property_presence
                        .get(name)
                        .cloned()
                        .unwrap_or_default();

                    let internal_candidates = match candidates {
                        Some(existing) => existing
                            .intersection(&workers_with_property)
                            .cloned()
                            .collect(),
                        None => workers_with_property,
                    };

                    if internal_candidates.is_empty() {
                        if full_worker_logging {
                            info!(
                                "No candidate workers due to a lack of key '{name}'. Job asked for {value:?}"
                            );
                        }
                        return HashSet::new();
                    }
                    candidates = Some(internal_candidates);
                }
                PlatformPropertyValue::Ignore(_) => {}
            }
        }

        candidates.unwrap_or_else(|| self.all_workers1.clone1())
    }

    /// Returns the number of indexed workers.
    pub fn worker_count(&self) -> usize {
        self.all_workers.len()
    }

    /// Returns true if the index is empty.
    pub fn is_empty(&self) -> bool {
        self.all_workers.is_empty()
    }
}

Coverage Report

Created: 2026-09-18 20:40

Line	Count	Source
1		// Copyright 2024 The NativeLink Authors. 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		//! Worker capability index for fast worker matching.
16		//!
17		//! This module provides an index that accelerates worker matching by property.
18		//! Instead of iterating all workers for each action, we maintain an inverted index
19		//! that maps property values to sets of workers that have those values.
20		//!
21		//! ## Complexity Analysis
22		//!
23		//! Without index: O(W × P) where W = workers, P = properties per action
24		//! With index: O(P × log(W)) for exact properties + O(W' × P') for minimum properties
25		//! where W' = filtered workers, P' = minimum property count (typically small)
26		//!
27		//! For typical workloads (few minimum properties), this reduces matching from
28		//! O(n × m) to approximately O(log n).
29
30		use std::collections::{HashMap, HashSet};
31
32		use nativelink_util::action_messages::WorkerId;
33		use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue};
34		use tracing::info;
35
36		/// A property key-value pair used for indexing.
37		#[derive(Clone, Hash, Eq, PartialEq, Debug)]
38		struct PropertyKey {
39		name: String,
40		value: PlatformPropertyValue,
41		}
42
43		/// Index structure for fast worker capability lookup.
44		///
45		/// Maintains an inverted index from property values to worker IDs.
46		/// Only indexes `Exact` and `Priority` properties since `Minimum` properties
47		/// are dynamic and require runtime comparison.
48		#[derive(Debug, Default)]
49		pub struct WorkerCapabilityIndex {
50		/// Maps `(property_name, property_value)` -> Set of worker IDs with that property.
51		/// Only contains `Exact` and `Priority` properties.
52		exact_index: HashMap<PropertyKey, HashSet<WorkerId>>,
53
54		/// Maps `property_name` -> Set of worker IDs that have this property (any value).
55		/// Used for fast "has property" checks for `Priority` and `Minimum` properties.
56		property_presence: HashMap<String, HashSet<WorkerId>>,
57
58		/// Set of all indexed worker IDs.
59		all_workers: HashSet<WorkerId>,
60		}
61
62		impl WorkerCapabilityIndex {
63		/// Creates a new empty capability index.
64	52	pub fn new() -> Self {
65	52	Self::default()
66	52	}
67
68		/// Adds a worker to the index with their platform properties.
69	67	pub fn add_worker(&mut self, worker_id: &WorkerId, properties: &PlatformProperties) {
70	67	self.all_workers.insert(worker_id.clone());
71
72	67	for ( name32 , value32 ) in &properties.properties {
73		// Track property presence
74	32	self.property_presence
75	32	.entry(name.clone())
76	32	.or_default()
77	32	.insert(worker_id.clone());
78
79	32	match value {
80		PlatformPropertyValue::Exact(_)
81		\| PlatformPropertyValue::Priority(_)
82	18	\| PlatformPropertyValue::Unknown(_) => {
83	18	// Index exact-match properties
84	18	let key = PropertyKey {
85	18	name: name.clone(),
86	18	value: value.clone(),
87	18	};
88	18	self.exact_index
89	18	.entry(key)
90	18	.or_default()
91	18	.insert(worker_id.clone());
92	18	}
93	14	PlatformPropertyValue::Minimum(_) \| PlatformPropertyValue::Ignore(_) => {
94	14	// Minimum properties are tracked via `property_presence` only.
95	14	// Their actual values are checked at runtime since they're dynamic.
96	14
97	14	// Ignore properties we just drop
98	14	}
99		}
100		}
101	67	}
102
103		/// Removes a worker from the index.
104	13	pub fn remove_worker(&mut self, worker_id: &WorkerId) {
105	13	self.all_workers.remove(worker_id);
106
107		// Remove from exact index
108	13	self.exact_index.retain(\|_, workers\| {1
109	1	workers.remove(worker_id);
110	1	!workers.is_empty()
111	1	});
112
113		// Remove from presence index
114	13	self.property_presence.retain(\|_, workers\| {3
115	3	workers.remove(worker_id);
116	3	!workers.is_empty()
117	3	});
118	13	}
119
120		/// Finds workers that can satisfy the given action properties.
121		///
122		/// Returns a set of worker IDs that match all required properties.
123		/// The caller should apply additional filtering (e.g., worker availability).
124		///
125		/// IMPORTANT: This method returns candidates based on STATIC properties only.
126		/// - Exact and Unknown properties are fully matched; Unknown properties
127		/// additionally match workers that do not declare the key
128		/// - Priority properties just require the key to exist
129		/// - Minimum properties return workers that HAVE the property (presence check only)
130		///
131		/// The caller MUST still verify Minimum property values at runtime because
132		/// worker resources change dynamically as jobs are assigned/completed.
133	72	pub fn find_matching_workers(
134	72	&self,
135	72	action_properties: &PlatformProperties,
136	72	full_worker_logging: bool,
137	72	) -> HashSet<WorkerId> {
138	72	if self.all_workers.is_empty() {
139	2	if full_worker_logging {
140	2	info!("No workers available to match!");
141	0	}
142	2	return HashSet::new();
143	70	}
144
145	70	if action_properties.properties.is_empty() {
146		// No properties required, all workers match
147	38	return self.all_workers.clone();
148	32	}
149
150	32	let mut candidates: Option<HashSet<WorkerId>> = None;
151
152	34	for (name, value) in &action_properties.properties32 {
153	34	match value {
154		PlatformPropertyValue::Exact(_) \| PlatformPropertyValue::Unknown(_) => {
155		// Look up workers with exact match
156	8	let key = PropertyKey {
157	8	name: name.clone(),
158	8	value: value.clone(),
159	8	};
160
161	8	let mut matching = self.exact_index.get(&key).cloned().unwrap_or_default();
162
163		// Unknown properties only restrict workers that declare the
164		// key, so workers without the key remain candidates.
165	8	if matches!6 (value, PlatformPropertyValue::Unknown(_)) {
166	2	match self.property_presence.get(name) {
167	1	Some(with_key) => {
168	1	matching.extend(self.all_workers.difference(with_key).cloned());
169	1	}
170	1	None => matching.extend(self.all_workers.iter().cloned()),
171		}
172	6	}
173
174	8	let internal_candidates = match candidates {
175	0	Some(existing) => existing.intersection(&matching).cloned().collect(),
176	8	None => matching,
177		};
178
179		// Early exit if no candidates
180	8	if internal_candidates.is_empty() {
181	2	if full_worker_logging {
182	1	let values: Vec<_> = self
183	1	.exact_index
184	1	.iter()
185	1	.filter(\|pk\| &pk.0.name == name)
186	1	.map(\|pk\| pk.0.value.clone())
187	1	.collect();
188	1	info!(
189		"No candidate workers due to a lack of matching '{name}' = {value:?}. Workers have: {values:?}"
190		);
191	1	}
192	2	return HashSet::new();
193	6	}
194	6	candidates = Some(internal_candidates);
195		}
196		PlatformPropertyValue::Priority(_) \| PlatformPropertyValue::Minimum(_) => {
197		// Priority: just requires the key to exist
198		// Minimum: worker must have the property (value checked at runtime by caller)
199		// We only check presence here because Minimum values are DYNAMIC -
200		// they change as jobs are assigned to workers.
201	25	let workers_with_property = self
202	25	.property_presence
203	25	.get(name)
204	25	.cloned()
205	25	.unwrap_or_default();
206
207	25	let internal_candidates = match candidates {
208	2	Some(existing) => existing
209	2	.intersection(&workers_with_property)
210	2	.cloned()
211	2	.collect(),
212	23	None => workers_with_property,
213		};
214
215	25	if internal_candidates.is_empty() {
216	1	if full_worker_logging {
217	1	info!(
218		"No candidate workers due to a lack of key '{name}'. Job asked for {value:?}"
219		);
220	0	}
221	1	return HashSet::new();
222	24	}
223	24	candidates = Some(internal_candidates);
224		}
225	1	PlatformPropertyValue::Ignore(_) => {}
226		}
227		}
228
229	29	candidates.unwrap_or_else(\|\| self.all_workers1 . clone1 ())
230	72	}
231
232		/// Returns the number of indexed workers.
233	2	pub fn worker_count(&self) -> usize {
234	2	self.all_workers.len()
235	2	}
236
237		/// Returns true if the index is empty.
238	0	pub fn is_empty(&self) -> bool {
239	0	self.all_workers.is_empty()
240	0	}
241		}