/build/source/nativelink-scheduler/src/awaited_action_db/mod.rs

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

use core::cmp;
use core::ops::Bound;
use core::time::Duration;
use std::sync::Arc;

pub use awaited_action::{AwaitedAction, AwaitedActionSortKey, PersistedSortKey};
use futures::{Future, Stream};
use nativelink_error::{Error, ResultExt, make_input_err};
use nativelink_metric::MetricsComponent;
use nativelink_util::action_messages::{ActionInfo, ActionStage, OperationId};
use nativelink_util::platform_properties::PlatformProperties;
use serde::{Deserialize, Serialize};

use crate::worker_registry::SharedWorkerRegistry;

mod awaited_action;

/// Duration to wait before sending client keep alive messages.
pub const CLIENT_KEEPALIVE_DURATION: Duration = Duration::from_secs(10);

/// A simple enum to represent the state of an `AwaitedAction`.
#[derive(Debug, Clone, Copy)]
pub enum SortedAwaitedActionState {
    CacheCheck,
    Queued,
    Executing,
    Completed,
}

impl TryFrom<&ActionStage> for SortedAwaitedActionState {
    type Error = Error;
    fn try_from(value: &ActionStage) -> Result<Self, Error> {
        match value {
            ActionStage::CacheCheck => Ok(Self::CacheCheck),
            ActionStage::Executing => Ok(Self::Executing),
            ActionStage::Completed(_) => Ok(Self::Completed),
            ActionStage::Queued => Ok(Self::Queued),
            _ => Err(make_input_err!("Invalid State")),
        }
    }
}

impl TryFrom<ActionStage> for SortedAwaitedActionState {
    type Error = Error;
    fn try_from(value: ActionStage) -> Result<Self, Error> {
        Self::try_from(&value)
    }
}

/// A struct pointing to an `AwaitedAction` that can be sorted.
#[derive(Debug, Clone, Serialize, Deserialize, MetricsComponent)]
pub struct SortedAwaitedAction {
    #[metric(help = "The sort key of the AwaitedAction")]
    pub sort_key: AwaitedActionSortKey,
    #[metric(help = "The operation id")]
    pub operation_id: OperationId,
}

impl PartialEq for SortedAwaitedAction {
    fn eq(&self, other: &Self) -> bool {
        self.sort_key == other.sort_key && self.operation_id == other.operation_id
    }
}

impl Eq for SortedAwaitedAction {}

impl PartialOrd for SortedAwaitedAction {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for SortedAwaitedAction {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.sort_key
            .cmp(&other.sort_key)
            .then_with(|| self.operation_id60.3k.cmp60.3k(&other.operation_id60.3k))
    }
}

impl core::fmt::Display for SortedAwaitedAction {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        core::fmt::write(
            f,
            format_args!("{}-{}", self.sort_key.as_u128(), self.operation_id),
        )
    }
}

impl From<&AwaitedAction> for SortedAwaitedAction {
    fn from(value: &AwaitedAction) -> Self {
        Self {
            operation_id: value.operation_id().clone(),
            sort_key: value.sort_key(),
        }
    }
}

impl From<AwaitedAction> for SortedAwaitedAction {
    fn from(value: AwaitedAction) -> Self {
        Self::from(&value)
    }
}

impl TryInto<Vec<u8>> for SortedAwaitedAction {
    type Error = Error;
    fn try_into(self) -> Result<Vec<u8>, Self::Error> {
        serde_json::to_vec(&self)
            .map_err(|e| make_input_err!("{}", e.to_string()))
            .err_tip(|| "In SortedAwaitedAction::TryInto::<Vec<u8>>")
    }
}

impl TryFrom<&[u8]> for SortedAwaitedAction {
    type Error = Error;
    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        serde_json::from_slice(value)
            .map_err(|e| make_input_err!("{}", e.to_string()))
            .err_tip(|| "In AwaitedAction::TryFrom::&[u8]")
    }
}

/// Subscriber that can be used to monitor when `AwaitedActions` change.
pub trait AwaitedActionSubscriber: Send + Sync + Sized + 'static {
    /// Wait for `AwaitedAction` to change.
    fn changed(&mut self) -> impl Future<Output = Result<AwaitedAction, Error>> + Send;

    /// Get the current awaited action.
    fn borrow(&self) -> impl Future<Output = Result<AwaitedAction, Error>> + Send;
}

/// A trait that defines the interface for an `AwaitedActionDb`.
pub trait AwaitedActionDb: Send + Sync + MetricsComponent + Unpin + 'static {
    type Subscriber: AwaitedActionSubscriber;

    /// Get the `AwaitedAction` by the client operation id.
    fn get_awaited_action_by_id(
        &self,
        client_operation_id: &OperationId,
    ) -> impl Future<Output = Result<Option<Self::Subscriber>, Error>> + Send;

    /// Get all `AwaitedActions`. This call should be avoided as much as possible.
    fn get_all_awaited_actions(
        &self,
    ) -> impl Future<
        Output = Result<impl Stream<Item = Result<Self::Subscriber, Error>> + Send, Error>,
    > + Send;

    /// Get the `AwaitedAction` by the operation id.
    fn get_by_operation_id(
        &self,
        operation_id: &OperationId,
    ) -> impl Future<Output = Result<Option<Self::Subscriber>, Error>> + Send;

    /// Get a range of `AwaitedActions` of a specific state in sorted order.
    fn get_range_of_actions(
        &self,
        state: SortedAwaitedActionState,
        start: Bound<SortedAwaitedAction>,
        end: Bound<SortedAwaitedAction>,
        desc: bool,
    ) -> impl Future<
        Output = Result<impl Stream<Item = Result<Self::Subscriber, Error>> + Send, Error>,
    > + Send;

    /// Process a change changed `AwaitedAction` and notify any listeners.
    fn update_awaited_action(
        &self,
        new_awaited_action: AwaitedAction,
    ) -> impl Future<Output = Result<(), Error>> + Send;

    /// Add (or join) an action to the `AwaitedActionDb` and subscribe
    /// to changes.
    fn add_action(
        &self,
        client_operation_id: OperationId,
        action_info: Arc<ActionInfo>,
        no_event_action_timeout: Duration,
    ) -> impl Future<Output = Result<Self::Subscriber, Error>> + Send;

    /// Lets an implementation judge worker liveness before deciding an
    /// executing action was abandoned. Only shared-state implementations
    /// need it, so the default ignores the registry.
    fn set_worker_registry(&mut self, _worker_registry: SharedWorkerRegistry) {}

    /// Tells other schedulers on the same state what the workers connected
    /// here can run, and returns, per peer scheduler, what the workers
    /// connected to it can run. `local` holds the properties each distinct
    /// kind of worker registered with, and stays visible to peers for
    /// `ttl`. Only shared-state implementations need it; the default shares
    /// nothing and sees nothing.
    fn exchange_fleet_capabilities(
        &self,
        _scheduler_id: &str,
        _local: Vec<PlatformProperties>,
        _ttl: Duration,
    ) -> impl Future<Output = Result<Vec<PeerScheduler>, Error>> + Send {
        async { Ok(Vec::new()) }
    }

    /// Whether other schedulers may share this state, and so have workers
    /// of their own that `exchange_fleet_capabilities` reports.
    fn shares_state(&self) -> bool {
        false
    }
}

/// What one peer scheduler's workers can run, as it published them.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PeerScheduler {
    /// The id the peer publishes under, unique to its process.
    pub scheduler_id: String,
    /// The properties each distinct kind of worker connected to the peer
    /// registered with.
    pub workers: Vec<PlatformProperties>,
}

Coverage Report

Created: 2026-11-13 23:29

Line	Count	Source
1		// Copyright 2024 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		use core::cmp;
16		use core::ops::Bound;
17		use core::time::Duration;
18		use std::sync::Arc;
19
20		pub use awaited_action::{AwaitedAction, AwaitedActionSortKey, PersistedSortKey};
21		use futures::{Future, Stream};
22		use nativelink_error::{Error, ResultExt, make_input_err};
23		use nativelink_metric::MetricsComponent;
24		use nativelink_util::action_messages::{ActionInfo, ActionStage, OperationId};
25		use nativelink_util::platform_properties::PlatformProperties;
26		use serde::{Deserialize, Serialize};
27
28		use crate::worker_registry::SharedWorkerRegistry;
29
30		mod awaited_action;
31
32		/// Duration to wait before sending client keep alive messages.
33		pub const CLIENT_KEEPALIVE_DURATION: Duration = Duration::from_secs(10);
34
35		/// A simple enum to represent the state of an `AwaitedAction`.
36		#[derive(Debug, Clone, Copy)]
37		pub enum SortedAwaitedActionState {
38		CacheCheck,
39		Queued,
40		Executing,
41		Completed,
42		}
43
44		impl TryFrom<&ActionStage> for SortedAwaitedActionState {
45		type Error = Error;
46	7.90k	fn try_from(value: &ActionStage) -> Result<Self, Error> {
47	7.90k	match value {
48	0	ActionStage::CacheCheck => Ok(Self::CacheCheck),
49	1.18k	ActionStage::Executing => Ok(Self::Executing),
50	1.55k	ActionStage::Completed(_) => Ok(Self::Completed),
51	5.16k	ActionStage::Queued => Ok(Self::Queued),
52	0	_ => Err(make_input_err!("Invalid State")),
53		}
54	7.90k	}
55		}
56
57		impl TryFrom<ActionStage> for SortedAwaitedActionState {
58		type Error = Error;
59	0	fn try_from(value: ActionStage) -> Result<Self, Error> {
60	0	Self::try_from(&value)
61	0	}
62		}
63
64		/// A struct pointing to an `AwaitedAction` that can be sorted.
65		#[derive(Debug, Clone, Serialize, Deserialize, MetricsComponent)]
66		pub struct SortedAwaitedAction {
67		#[metric(help = "The sort key of the AwaitedAction")]
68		pub sort_key: AwaitedActionSortKey,
69		#[metric(help = "The operation id")]
70		pub operation_id: OperationId,
71		}
72
73		impl PartialEq for SortedAwaitedAction {
74	0	fn eq(&self, other: &Self) -> bool {
75	0	self.sort_key == other.sort_key && self.operation_id == other.operation_id
76	0	}
77		}
78
79		impl Eq for SortedAwaitedAction {}
80
81		impl PartialOrd for SortedAwaitedAction {
82	0	fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
83	0	Some(self.cmp(other))
84	0	}
85		}
86
87		impl Ord for SortedAwaitedAction {
88	404k	fn cmp(&self, other: &Self) -> cmp::Ordering {
89	404k	self.sort_key
90	404k	.cmp(&other.sort_key)
91	404k	.then_with(\|\| self.operation_id60.3k . cmp60.3k ( &other.operation_id60.3k ))
92	404k	}
93		}
94
95		impl core::fmt::Display for SortedAwaitedAction {
96	0	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
97	0	core::fmt::write(
98	0	f,
99	0	format_args!("{}-{}", self.sort_key.as_u128(), self.operation_id),
100		)
101	0	}
102		}
103
104		impl From<&AwaitedAction> for SortedAwaitedAction {
105	7.91k	fn from(value: &AwaitedAction) -> Self {
106	7.91k	Self {
107	7.91k	operation_id: value.operation_id().clone(),
108	7.91k	sort_key: value.sort_key(),
109	7.91k	}
110	7.91k	}
111		}
112
113		impl From<AwaitedAction> for SortedAwaitedAction {
114	0	fn from(value: AwaitedAction) -> Self {
115	0	Self::from(&value)
116	0	}
117		}
118
119		impl TryInto<Vec<u8>> for SortedAwaitedAction {
120		type Error = Error;
121	0	fn try_into(self) -> Result<Vec<u8>, Self::Error> {
122	0	serde_json::to_vec(&self)
123	0	.map_err(\|e\| make_input_err!("{}", e.to_string()))
124	0	.err_tip(\|\| "In SortedAwaitedAction::TryInto::<Vec<u8>>")
125	0	}
126		}
127
128		impl TryFrom<&[u8]> for SortedAwaitedAction {
129		type Error = Error;
130	0	fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
131	0	serde_json::from_slice(value)
132	0	.map_err(\|e\| make_input_err!("{}", e.to_string()))
133	0	.err_tip(\|\| "In AwaitedAction::TryFrom::&[u8]")
134	0	}
135		}
136
137		/// Subscriber that can be used to monitor when `AwaitedActions` change.
138		pub trait AwaitedActionSubscriber: Send + Sync + Sized + 'static {
139		/// Wait for `AwaitedAction` to change.
140		fn changed(&mut self) -> impl Future<Output = Result<AwaitedAction, Error>> + Send;
141
142		/// Get the current awaited action.
143		fn borrow(&self) -> impl Future<Output = Result<AwaitedAction, Error>> + Send;
144		}
145
146		/// A trait that defines the interface for an `AwaitedActionDb`.
147		pub trait AwaitedActionDb: Send + Sync + MetricsComponent + Unpin + 'static {
148		type Subscriber: AwaitedActionSubscriber;
149
150		/// Get the `AwaitedAction` by the client operation id.
151		fn get_awaited_action_by_id(
152		&self,
153		client_operation_id: &OperationId,
154		) -> impl Future<Output = Result<Option<Self::Subscriber>, Error>> + Send;
155
156		/// Get all `AwaitedActions`. This call should be avoided as much as possible.
157		fn get_all_awaited_actions(
158		&self,
159		) -> impl Future<
160		Output = Result<impl Stream<Item = Result<Self::Subscriber, Error>> + Send, Error>,
161		> + Send;
162
163		/// Get the `AwaitedAction` by the operation id.
164		fn get_by_operation_id(
165		&self,
166		operation_id: &OperationId,
167		) -> impl Future<Output = Result<Option<Self::Subscriber>, Error>> + Send;
168
169		/// Get a range of `AwaitedActions` of a specific state in sorted order.
170		fn get_range_of_actions(
171		&self,
172		state: SortedAwaitedActionState,
173		start: Bound<SortedAwaitedAction>,
174		end: Bound<SortedAwaitedAction>,
175		desc: bool,
176		) -> impl Future<
177		Output = Result<impl Stream<Item = Result<Self::Subscriber, Error>> + Send, Error>,
178		> + Send;
179
180		/// Process a change changed `AwaitedAction` and notify any listeners.
181		fn update_awaited_action(
182		&self,
183		new_awaited_action: AwaitedAction,
184		) -> impl Future<Output = Result<(), Error>> + Send;
185
186		/// Add (or join) an action to the `AwaitedActionDb` and subscribe
187		/// to changes.
188		fn add_action(
189		&self,
190		client_operation_id: OperationId,
191		action_info: Arc<ActionInfo>,
192		no_event_action_timeout: Duration,
193		) -> impl Future<Output = Result<Self::Subscriber, Error>> + Send;
194
195		/// Lets an implementation judge worker liveness before deciding an
196		/// executing action was abandoned. Only shared-state implementations
197		/// need it, so the default ignores the registry.
198	271	fn set_worker_registry(&mut self, _worker_registry: SharedWorkerRegistry) {}
199
200		/// Tells other schedulers on the same state what the workers connected
201		/// here can run, and returns, per peer scheduler, what the workers
202		/// connected to it can run. `local` holds the properties each distinct
203		/// kind of worker registered with, and stays visible to peers for
204		/// `ttl`. Only shared-state implementations need it; the default shares
205		/// nothing and sees nothing.
206	0	fn exchange_fleet_capabilities(
207	0	&self,
208	0	_scheduler_id: &str,
209	0	_local: Vec<PlatformProperties>,
210	0	_ttl: Duration,
211	0	) -> impl Future<Output = Result<Vec<PeerScheduler>, Error>> + Send {
212	0	async { Ok(Vec::new()) }
213	0	}
214
215		/// Whether other schedulers may share this state, and so have workers
216		/// of their own that `exchange_fleet_capabilities` reports.
217	271	fn shares_state(&self) -> bool {
218	271	false
219	271	}
220		}
221
222		/// What one peer scheduler's workers can run, as it published them.
223		#[derive(Debug, Clone, PartialEq, Eq)]
224		pub struct PeerScheduler {
225		/// The id the peer publishes under, unique to its process.
226		pub scheduler_id: String,
227		/// The properties each distinct kind of worker connected to the peer
228		/// registered with.
229		pub workers: Vec<PlatformProperties>,
230		}