/build/source/nativelink-scheduler/src/awaited_action_db/mod.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.

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

pub use awaited_action::{AwaitedAction, AwaitedActionSortKey};
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 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_id593.cmp593(&other.operation_id593))
    }
}

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_u64(), 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) {}
}

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		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};
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 serde::{Deserialize, Serialize};
26
27		use crate::worker_registry::SharedWorkerRegistry;
28
29		mod awaited_action;
30
31		/// Duration to wait before sending client keep alive messages.
32		pub const CLIENT_KEEPALIVE_DURATION: Duration = Duration::from_secs(10);
33
34		/// A simple enum to represent the state of an `AwaitedAction`.
35		#[derive(Debug, Clone, Copy)]
36		pub enum SortedAwaitedActionState {
37		CacheCheck,
38		Queued,
39		Executing,
40		Completed,
41		}
42
43		impl TryFrom<&ActionStage> for SortedAwaitedActionState {
44		type Error = Error;
45	19	fn try_from(value: &ActionStage) -> Result<Self, Error> {
46	19	match value {
47	0	ActionStage::CacheCheck => Ok(Self::CacheCheck),
48	7	ActionStage::Executing => Ok(Self::Executing),
49	1	ActionStage::Completed(_) => Ok(Self::Completed),
50	11	ActionStage::Queued => Ok(Self::Queued),
51	0	_ => Err(make_input_err!("Invalid State")),
52		}
53	19	}
54		}
55
56		impl TryFrom<ActionStage> for SortedAwaitedActionState {
57		type Error = Error;
58	0	fn try_from(value: ActionStage) -> Result<Self, Error> {
59	0	Self::try_from(&value)
60	0	}
61		}
62
63		/// A struct pointing to an `AwaitedAction` that can be sorted.
64		#[derive(Debug, Clone, Serialize, Deserialize, MetricsComponent)]
65		pub struct SortedAwaitedAction {
66		#[metric(help = "The sort key of the AwaitedAction")]
67		pub sort_key: AwaitedActionSortKey,
68		#[metric(help = "The operation id")]
69		pub operation_id: OperationId,
70		}
71
72		impl PartialEq for SortedAwaitedAction {
73	0	fn eq(&self, other: &Self) -> bool {
74	0	self.sort_key == other.sort_key && self.operation_id == other.operation_id
75	0	}
76		}
77
78		impl Eq for SortedAwaitedAction {}
79
80		impl PartialOrd for SortedAwaitedAction {
81	0	fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
82	0	Some(self.cmp(other))
83	0	}
84		}
85
86		impl Ord for SortedAwaitedAction {
87	601	fn cmp(&self, other: &Self) -> cmp::Ordering {
88	601	self.sort_key
89	601	.cmp(&other.sort_key)
90	601	.then_with(\|\| self.operation_id593 . cmp593 ( &other.operation_id593 ))
91	601	}
92		}
93
94		impl core::fmt::Display for SortedAwaitedAction {
95	0	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
96	0	core::fmt::write(
97	0	f,
98	0	format_args!("{}-{}", self.sort_key.as_u64(), self.operation_id),
99		)
100	0	}
101		}
102
103		impl From<&AwaitedAction> for SortedAwaitedAction {
104	19	fn from(value: &AwaitedAction) -> Self {
105	19	Self {
106	19	operation_id: value.operation_id().clone(),
107	19	sort_key: value.sort_key(),
108	19	}
109	19	}
110		}
111
112		impl From<AwaitedAction> for SortedAwaitedAction {
113	0	fn from(value: AwaitedAction) -> Self {
114	0	Self::from(&value)
115	0	}
116		}
117
118		impl TryInto<Vec<u8>> for SortedAwaitedAction {
119		type Error = Error;
120	0	fn try_into(self) -> Result<Vec<u8>, Self::Error> {
121	0	serde_json::to_vec(&self)
122	0	.map_err(\|e\| make_input_err!("{}", e.to_string()))
123	0	.err_tip(\|\| "In SortedAwaitedAction::TryInto::<Vec<u8>>")
124	0	}
125		}
126
127		impl TryFrom<&[u8]> for SortedAwaitedAction {
128		type Error = Error;
129	0	fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
130	0	serde_json::from_slice(value)
131	0	.map_err(\|e\| make_input_err!("{}", e.to_string()))
132	0	.err_tip(\|\| "In AwaitedAction::TryFrom::&[u8]")
133	0	}
134		}
135
136		/// Subscriber that can be used to monitor when `AwaitedActions` change.
137		pub trait AwaitedActionSubscriber: Send + Sync + Sized + 'static {
138		/// Wait for `AwaitedAction` to change.
139		fn changed(&mut self) -> impl Future<Output = Result<AwaitedAction, Error>> + Send;
140
141		/// Get the current awaited action.
142		fn borrow(&self) -> impl Future<Output = Result<AwaitedAction, Error>> + Send;
143		}
144
145		/// A trait that defines the interface for an `AwaitedActionDb`.
146		pub trait AwaitedActionDb: Send + Sync + MetricsComponent + Unpin + 'static {
147		type Subscriber: AwaitedActionSubscriber;
148
149		/// Get the `AwaitedAction` by the client operation id.
150		fn get_awaited_action_by_id(
151		&self,
152		client_operation_id: &OperationId,
153		) -> impl Future<Output = Result<Option<Self::Subscriber>, Error>> + Send;
154
155		/// Get all `AwaitedActions`. This call should be avoided as much as possible.
156		fn get_all_awaited_actions(
157		&self,
158		) -> impl Future<
159		Output = Result<impl Stream<Item = Result<Self::Subscriber, Error>> + Send, Error>,
160		> + Send;
161
162		/// Get the `AwaitedAction` by the operation id.
163		fn get_by_operation_id(
164		&self,
165		operation_id: &OperationId,
166		) -> impl Future<Output = Result<Option<Self::Subscriber>, Error>> + Send;
167
168		/// Get a range of `AwaitedActions` of a specific state in sorted order.
169		fn get_range_of_actions(
170		&self,
171		state: SortedAwaitedActionState,
172		start: Bound<SortedAwaitedAction>,
173		end: Bound<SortedAwaitedAction>,
174		desc: bool,
175		) -> impl Future<
176		Output = Result<impl Stream<Item = Result<Self::Subscriber, Error>> + Send, Error>,
177		> + Send;
178
179		/// Process a change changed `AwaitedAction` and notify any listeners.
180		fn update_awaited_action(
181		&self,
182		new_awaited_action: AwaitedAction,
183		) -> impl Future<Output = Result<(), Error>> + Send;
184
185		/// Add (or join) an action to the `AwaitedActionDb` and subscribe
186		/// to changes.
187		fn add_action(
188		&self,
189		client_operation_id: OperationId,
190		action_info: Arc<ActionInfo>,
191		no_event_action_timeout: Duration,
192		) -> impl Future<Output = Result<Self::Subscriber, Error>> + Send;
193
194		/// Lets an implementation judge worker liveness before deciding an
195		/// executing action was abandoned. Only shared-state implementations
196		/// need it, so the default ignores the registry.
197	32	fn set_worker_registry(&mut self, _worker_registry: SharedWorkerRegistry) {}
198		}

Coverage Report

Created: 2026-09-18 20:40