/build/source/nativelink-scheduler/src/simple_scheduler_state_manager.rs
Line | Count | Source |
1 | | // Copyright 2024 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Apache License, Version 2.0 (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 | | // http://www.apache.org/licenses/LICENSE-2.0 |
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 std::ops::Bound; |
16 | | use std::string::ToString; |
17 | | use std::sync::{Arc, Weak}; |
18 | | use std::time::{Duration, SystemTime}; |
19 | | |
20 | | use async_lock::Mutex; |
21 | | use async_trait::async_trait; |
22 | | use futures::{stream, StreamExt, TryStreamExt}; |
23 | | use nativelink_error::{make_err, Code, Error, ResultExt}; |
24 | | use nativelink_metric::MetricsComponent; |
25 | | use nativelink_util::action_messages::{ |
26 | | ActionInfo, ActionResult, ActionStage, ActionState, ActionUniqueQualifier, ExecutionMetadata, |
27 | | OperationId, WorkerId, |
28 | | }; |
29 | | use nativelink_util::instant_wrapper::InstantWrapper; |
30 | | use nativelink_util::known_platform_property_provider::KnownPlatformPropertyProvider; |
31 | | use nativelink_util::operation_state_manager::{ |
32 | | ActionStateResult, ActionStateResultStream, ClientStateManager, MatchingEngineStateManager, |
33 | | OperationFilter, OperationStageFlags, OrderDirection, UpdateOperationType, WorkerStateManager, |
34 | | }; |
35 | | use tracing::{event, Level}; |
36 | | |
37 | | use super::awaited_action_db::{ |
38 | | AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, SortedAwaitedActionState, |
39 | | }; |
40 | | |
41 | | /// Maximum number of times an update to the database |
42 | | /// can fail before giving up. |
43 | | const MAX_UPDATE_RETRIES: usize = 5; |
44 | | |
45 | | /// Simple struct that implements the `ActionStateResult` trait and always returns an error. |
46 | | struct ErrorActionStateResult(Error); |
47 | | |
48 | | #[async_trait] |
49 | | impl ActionStateResult for ErrorActionStateResult { |
50 | 0 | async fn as_state(&self) -> Result<Arc<ActionState>, Error> { |
51 | 0 | Err(self.0.clone()) |
52 | 0 | } |
53 | | |
54 | 0 | async fn changed(&mut self) -> Result<Arc<ActionState>, Error> { |
55 | 0 | Err(self.0.clone()) |
56 | 0 | } |
57 | | |
58 | 0 | async fn as_action_info(&self) -> Result<Arc<ActionInfo>, Error> { |
59 | 0 | Err(self.0.clone()) |
60 | 0 | } |
61 | | } |
62 | | |
63 | | struct ClientActionStateResult<U, T, I, NowFn> |
64 | | where |
65 | | U: AwaitedActionSubscriber, |
66 | | T: AwaitedActionDb, |
67 | | I: InstantWrapper, |
68 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
69 | | { |
70 | | inner: MatchingEngineActionStateResult<U, T, I, NowFn>, |
71 | | } |
72 | | |
73 | | impl<U, T, I, NowFn> ClientActionStateResult<U, T, I, NowFn> |
74 | | where |
75 | | U: AwaitedActionSubscriber, |
76 | | T: AwaitedActionDb, |
77 | | I: InstantWrapper, |
78 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
79 | | { |
80 | 28 | fn new( |
81 | 28 | sub: U, |
82 | 28 | simple_scheduler_state_manager: Weak<SimpleSchedulerStateManager<T, I, NowFn>>, |
83 | 28 | no_event_action_timeout: Duration, |
84 | 28 | now_fn: NowFn, |
85 | 28 | ) -> Self { |
86 | 28 | Self { |
87 | 28 | inner: MatchingEngineActionStateResult::new( |
88 | 28 | sub, |
89 | 28 | simple_scheduler_state_manager, |
90 | 28 | no_event_action_timeout, |
91 | 28 | now_fn, |
92 | 28 | ), |
93 | 28 | } |
94 | 28 | } |
95 | | } |
96 | | |
97 | | #[async_trait] |
98 | | impl<U, T, I, NowFn> ActionStateResult for ClientActionStateResult<U, T, I, NowFn> |
99 | | where |
100 | | U: AwaitedActionSubscriber, |
101 | | T: AwaitedActionDb, |
102 | | I: InstantWrapper, |
103 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
104 | | { |
105 | 3 | async fn as_state(&self) -> Result<Arc<ActionState>, Error> { |
106 | 3 | self.inner.as_state().await |
107 | 6 | } |
108 | | |
109 | 44 | async fn changed(&mut self) -> Result<Arc<ActionState>, Error> { |
110 | 44 | self.inner.changed().await |
111 | 86 | } |
112 | | |
113 | 0 | async fn as_action_info(&self) -> Result<Arc<ActionInfo>, Error> { |
114 | 0 | self.inner.as_action_info().await |
115 | 0 | } |
116 | | } |
117 | | |
118 | | struct MatchingEngineActionStateResult<U, T, I, NowFn> |
119 | | where |
120 | | U: AwaitedActionSubscriber, |
121 | | T: AwaitedActionDb, |
122 | | I: InstantWrapper, |
123 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
124 | | { |
125 | | awaited_action_sub: U, |
126 | | simple_scheduler_state_manager: Weak<SimpleSchedulerStateManager<T, I, NowFn>>, |
127 | | no_event_action_timeout: Duration, |
128 | | now_fn: NowFn, |
129 | | } |
130 | | impl<U, T, I, NowFn> MatchingEngineActionStateResult<U, T, I, NowFn> |
131 | | where |
132 | | U: AwaitedActionSubscriber, |
133 | | T: AwaitedActionDb, |
134 | | I: InstantWrapper, |
135 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
136 | | { |
137 | 70 | fn new( |
138 | 70 | awaited_action_sub: U, |
139 | 70 | simple_scheduler_state_manager: Weak<SimpleSchedulerStateManager<T, I, NowFn>>, |
140 | 70 | no_event_action_timeout: Duration, |
141 | 70 | now_fn: NowFn, |
142 | 70 | ) -> Self { |
143 | 70 | Self { |
144 | 70 | awaited_action_sub, |
145 | 70 | simple_scheduler_state_manager, |
146 | 70 | no_event_action_timeout, |
147 | 70 | now_fn, |
148 | 70 | } |
149 | 70 | } |
150 | | } |
151 | | |
152 | | #[async_trait] |
153 | | impl<U, T, I, NowFn> ActionStateResult for MatchingEngineActionStateResult<U, T, I, NowFn> |
154 | | where |
155 | | U: AwaitedActionSubscriber, |
156 | | T: AwaitedActionDb, |
157 | | I: InstantWrapper, |
158 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
159 | | { |
160 | 31 | async fn as_state(&self) -> Result<Arc<ActionState>, Error> { |
161 | 31 | Ok(self |
162 | 31 | .awaited_action_sub |
163 | 31 | .borrow() |
164 | 31 | .await |
165 | 31 | .err_tip(|| "In MatchingEngineActionStateResult::as_state"0 )?0 |
166 | 31 | .state() |
167 | 31 | .clone()) |
168 | 62 | } |
169 | | |
170 | 44 | async fn changed(&mut self) -> Result<Arc<ActionState>, Error> { |
171 | 44 | let mut timeout_attempts = 0; |
172 | | loop { |
173 | 54 | tokio::select! { |
174 | 54 | awaited_action_result42 = self.awaited_action_sub.changed() => { |
175 | 42 | return awaited_action_result |
176 | 42 | .err_tip(|| "In MatchingEngineActionStateResult::changed"0 ) |
177 | 42 | .map(|v| v.state().clone()); |
178 | | } |
179 | 54 | () = (self.now_fn)().sleep(self.no_event_action_timeout) => { |
180 | 10 | // Timeout happened, do additional checks below. |
181 | 10 | } |
182 | | } |
183 | | |
184 | 10 | let awaited_action = self |
185 | 10 | .awaited_action_sub |
186 | 10 | .borrow() |
187 | 10 | .await |
188 | 10 | .err_tip(|| "In MatchingEngineActionStateResult::changed"0 )?0 ; |
189 | | |
190 | 10 | if matches!1 (awaited_action.state().stage, ActionStage::Queued) { |
191 | | // Actions in queued state do not get periodically updated, |
192 | | // so we don't need to timeout them. |
193 | 9 | continue; |
194 | 1 | } |
195 | | |
196 | 1 | let simple_scheduler_state_manager = self |
197 | 1 | .simple_scheduler_state_manager |
198 | 1 | .upgrade() |
199 | 1 | .err_tip(|| format!("Failed to upgrade weak reference to SimpleSchedulerStateManager in MatchingEngineActionStateResult::changed at attempt: {timeout_attempts}")0 )?0 ; |
200 | | |
201 | 1 | event!( |
202 | 1 | Level::WARN, |
203 | | ?awaited_action, |
204 | 1 | "OperationId {} / {} timed out after {} seconds issuing a retry", |
205 | 1 | awaited_action.operation_id(), |
206 | 1 | awaited_action.state().client_operation_id, |
207 | 1 | self.no_event_action_timeout.as_secs_f32(), |
208 | | ); |
209 | | |
210 | 1 | simple_scheduler_state_manager |
211 | 1 | .timeout_operation_id(awaited_action.operation_id()) |
212 | 1 | .await |
213 | 1 | .err_tip(|| "In MatchingEngineActionStateResult::changed"0 )?0 ; |
214 | | |
215 | 1 | if timeout_attempts >= MAX_UPDATE_RETRIES { Branch (215:16): [True: 0, False: 0]
Branch (215:16): [True: 0, False: 0]
Branch (215:16): [Folded - Ignored]
Branch (215:16): [True: 0, False: 1]
Branch (215:16): [True: 0, False: 0]
|
216 | 0 | return Err(make_err!( |
217 | 0 | Code::Internal, |
218 | 0 | "Failed to update action after {} retries with no error set in MatchingEngineActionStateResult::changed - {} {:?}", |
219 | 0 | MAX_UPDATE_RETRIES, |
220 | 0 | awaited_action.operation_id(), |
221 | 0 | awaited_action.state().stage, |
222 | 0 | )); |
223 | 1 | } |
224 | 1 | timeout_attempts += 1; |
225 | | } |
226 | 86 | } |
227 | | |
228 | 42 | async fn as_action_info(&self) -> Result<Arc<ActionInfo>, Error> { |
229 | 42 | Ok(self |
230 | 42 | .awaited_action_sub |
231 | 42 | .borrow() |
232 | 42 | .await |
233 | 42 | .err_tip(|| "In MatchingEngineActionStateResult::as_action_info"0 )?0 |
234 | 42 | .action_info() |
235 | 42 | .clone()) |
236 | 84 | } |
237 | | } |
238 | | |
239 | | /// `SimpleSchedulerStateManager` is responsible for maintaining the state of the scheduler. |
240 | | /// Scheduler state includes the actions that are queued, active, and recently completed. |
241 | | /// It also includes the workers that are available to execute actions based on allocation |
242 | | /// strategy. |
243 | | #[derive(MetricsComponent)] |
244 | | pub struct SimpleSchedulerStateManager<T, I, NowFn> |
245 | | where |
246 | | T: AwaitedActionDb, |
247 | | I: InstantWrapper, |
248 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
249 | | { |
250 | | /// Database for storing the state of all actions. |
251 | | #[metric(group = "action_db")] |
252 | | action_db: T, |
253 | | |
254 | | /// Maximum number of times a job can be retried. |
255 | | // TODO(allada) This should be a scheduler decorator instead |
256 | | // of always having it on every SimpleScheduler. |
257 | | #[metric(help = "Maximum number of times a job can be retried")] |
258 | | max_job_retries: usize, |
259 | | |
260 | | /// Duration after which an action is considered to be timed out if |
261 | | /// no event is received. |
262 | | #[metric( |
263 | | help = "Duration after which an action is considered to be timed out if no event is received" |
264 | | )] |
265 | | no_event_action_timeout: Duration, |
266 | | |
267 | | /// Mark operation as timed out if the worker has not updated in this duration. |
268 | | /// This is used to prevent operations from being stuck in the queue forever |
269 | | /// if it is not being processed by any worker. |
270 | | client_action_timeout: Duration, |
271 | | |
272 | | // A lock to ensure only one timeout operation is running at a time |
273 | | // on this service. |
274 | | timeout_operation_mux: Mutex<()>, |
275 | | |
276 | | /// Weak reference to self. |
277 | | // We use a weak reference to reduce the risk of a memory leak from |
278 | | // future changes. If this becomes some kind of perforamnce issue, |
279 | | // we can consider using a strong reference. |
280 | | weak_self: Weak<Self>, |
281 | | |
282 | | /// Function to get the current time. |
283 | | now_fn: NowFn, |
284 | | } |
285 | | |
286 | | impl<T, I, NowFn> SimpleSchedulerStateManager<T, I, NowFn> |
287 | | where |
288 | | T: AwaitedActionDb, |
289 | | I: InstantWrapper, |
290 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
291 | | { |
292 | 21 | pub fn new( |
293 | 21 | max_job_retries: usize, |
294 | 21 | no_event_action_timeout: Duration, |
295 | 21 | client_action_timeout: Duration, |
296 | 21 | action_db: T, |
297 | 21 | now_fn: NowFn, |
298 | 21 | ) -> Arc<Self> { |
299 | 21 | Arc::new_cyclic(|weak_self| Self { |
300 | 21 | action_db, |
301 | 21 | max_job_retries, |
302 | 21 | no_event_action_timeout, |
303 | 21 | client_action_timeout, |
304 | 21 | timeout_operation_mux: Mutex::new(()), |
305 | 21 | weak_self: weak_self.clone(), |
306 | 21 | now_fn, |
307 | 21 | }) |
308 | 21 | } |
309 | | |
310 | 45 | async fn apply_filter_predicate( |
311 | 45 | &self, |
312 | 45 | awaited_action: &AwaitedAction, |
313 | 45 | filter: &OperationFilter, |
314 | 45 | ) -> bool { |
315 | 45 | // Note: The caller must filter `client_operation_id`. |
316 | 45 | |
317 | 45 | if awaited_action.last_client_keepalive_timestamp() + self.client_action_timeout Branch (317:12): [True: 0, False: 0]
Branch (317:12): [True: 0, False: 0]
Branch (317:12): [Folded - Ignored]
Branch (317:12): [True: 0, False: 43]
Branch (317:12): [True: 0, False: 2]
|
318 | 45 | < (self.now_fn)().now() |
319 | | { |
320 | 0 | if !awaited_action.state().stage.is_finished() { Branch (320:16): [True: 0, False: 0]
Branch (320:16): [True: 0, False: 0]
Branch (320:16): [Folded - Ignored]
Branch (320:16): [True: 0, False: 0]
Branch (320:16): [True: 0, False: 0]
|
321 | 0 | let mut state = awaited_action.state().as_ref().clone(); |
322 | 0 | state.stage = ActionStage::Completed(ActionResult { |
323 | 0 | error: Some(make_err!( |
324 | 0 | Code::DeadlineExceeded, |
325 | 0 | "Operation timed out {} seconds of having no more clients listening", |
326 | 0 | self.client_action_timeout.as_secs_f32(), |
327 | 0 | )), |
328 | 0 | ..ActionResult::default() |
329 | 0 | }); |
330 | 0 | let mut new_awaited_action = awaited_action.clone(); |
331 | 0 | new_awaited_action.worker_set_state(Arc::new(state), (self.now_fn)().now()); |
332 | 0 | if let Err(err) = self Branch (332:24): [True: 0, False: 0]
Branch (332:24): [True: 0, False: 0]
Branch (332:24): [Folded - Ignored]
Branch (332:24): [True: 0, False: 0]
Branch (332:24): [True: 0, False: 0]
|
333 | 0 | .action_db |
334 | 0 | .update_awaited_action(new_awaited_action) |
335 | 0 | .await |
336 | | { |
337 | 0 | event!( |
338 | 0 | Level::WARN, |
339 | 0 | "Failed to update action to timed out state after client keepalive timeout. This is ok if multiple schedulers tried to set the state at the same time: {err}", |
340 | | ); |
341 | 0 | } |
342 | 0 | } |
343 | 0 | return false; |
344 | 45 | } |
345 | | |
346 | 45 | if let Some(operation_id0 ) = &filter.operation_id { Branch (346:16): [True: 0, False: 0]
Branch (346:16): [True: 0, False: 0]
Branch (346:16): [Folded - Ignored]
Branch (346:16): [True: 0, False: 43]
Branch (346:16): [True: 0, False: 2]
|
347 | 0 | if operation_id != awaited_action.operation_id() { Branch (347:16): [True: 0, False: 0]
Branch (347:16): [True: 0, False: 0]
Branch (347:16): [Folded - Ignored]
Branch (347:16): [True: 0, False: 0]
Branch (347:16): [True: 0, False: 0]
|
348 | 0 | return false; |
349 | 0 | } |
350 | 45 | } |
351 | | |
352 | 45 | if filter.worker_id.is_some() && filter.worker_id != awaited_action.worker_id()0 { Branch (352:12): [True: 0, False: 0]
Branch (352:42): [True: 0, False: 0]
Branch (352:12): [True: 0, False: 0]
Branch (352:42): [True: 0, False: 0]
Branch (352:12): [Folded - Ignored]
Branch (352:42): [Folded - Ignored]
Branch (352:12): [True: 0, False: 43]
Branch (352:42): [True: 0, False: 0]
Branch (352:12): [True: 0, False: 2]
Branch (352:42): [True: 0, False: 0]
|
353 | 0 | return false; |
354 | 45 | } |
355 | | |
356 | | { |
357 | 45 | if let Some(filter_unique_key0 ) = &filter.unique_key { Branch (357:20): [True: 0, False: 0]
Branch (357:20): [True: 0, False: 0]
Branch (357:20): [Folded - Ignored]
Branch (357:20): [True: 0, False: 43]
Branch (357:20): [True: 0, False: 2]
|
358 | 0 | match &awaited_action.action_info().unique_qualifier { |
359 | 0 | ActionUniqueQualifier::Cachable(unique_key) => { |
360 | 0 | if filter_unique_key != unique_key { Branch (360:28): [True: 0, False: 0]
Branch (360:28): [True: 0, False: 0]
Branch (360:28): [Folded - Ignored]
Branch (360:28): [True: 0, False: 0]
Branch (360:28): [True: 0, False: 0]
|
361 | 0 | return false; |
362 | 0 | } |
363 | | } |
364 | | ActionUniqueQualifier::Uncachable(_) => { |
365 | 0 | return false; |
366 | | } |
367 | | } |
368 | 45 | } |
369 | 45 | if let Some(action_digest0 ) = filter.action_digest { Branch (369:20): [True: 0, False: 0]
Branch (369:20): [True: 0, False: 0]
Branch (369:20): [Folded - Ignored]
Branch (369:20): [True: 0, False: 43]
Branch (369:20): [True: 0, False: 2]
|
370 | 0 | if action_digest != awaited_action.action_info().digest() { Branch (370:20): [True: 0, False: 0]
Branch (370:20): [True: 0, False: 0]
Branch (370:20): [Folded - Ignored]
Branch (370:20): [True: 0, False: 0]
Branch (370:20): [True: 0, False: 0]
|
371 | 0 | return false; |
372 | 0 | } |
373 | 45 | } |
374 | | } |
375 | | |
376 | | { |
377 | 45 | let last_worker_update_timestamp = awaited_action.last_worker_updated_timestamp(); |
378 | 45 | if let Some(worker_update_before0 ) = filter.worker_update_before { Branch (378:20): [True: 0, False: 0]
Branch (378:20): [True: 0, False: 0]
Branch (378:20): [Folded - Ignored]
Branch (378:20): [True: 0, False: 43]
Branch (378:20): [True: 0, False: 2]
|
379 | 0 | if worker_update_before < last_worker_update_timestamp { Branch (379:20): [True: 0, False: 0]
Branch (379:20): [True: 0, False: 0]
Branch (379:20): [Folded - Ignored]
Branch (379:20): [True: 0, False: 0]
Branch (379:20): [True: 0, False: 0]
|
380 | 0 | return false; |
381 | 0 | } |
382 | 45 | } |
383 | 45 | if let Some(completed_before0 ) = filter.completed_before { Branch (383:20): [True: 0, False: 0]
Branch (383:20): [True: 0, False: 0]
Branch (383:20): [Folded - Ignored]
Branch (383:20): [True: 0, False: 43]
Branch (383:20): [True: 0, False: 2]
|
384 | 0 | if awaited_action.state().stage.is_finished() Branch (384:20): [True: 0, False: 0]
Branch (384:20): [True: 0, False: 0]
Branch (384:20): [Folded - Ignored]
Branch (384:20): [True: 0, False: 0]
Branch (384:20): [True: 0, False: 0]
|
385 | 0 | && completed_before < last_worker_update_timestamp Branch (385:24): [True: 0, False: 0]
Branch (385:24): [True: 0, False: 0]
Branch (385:24): [Folded - Ignored]
Branch (385:24): [True: 0, False: 0]
Branch (385:24): [True: 0, False: 0]
|
386 | | { |
387 | 0 | return false; |
388 | 0 | } |
389 | 45 | } |
390 | 45 | if filter.stages != OperationStageFlags::Any { Branch (390:16): [True: 0, False: 0]
Branch (390:16): [True: 0, False: 0]
Branch (390:16): [Folded - Ignored]
Branch (390:16): [True: 40, False: 3]
Branch (390:16): [True: 2, False: 0]
|
391 | 42 | let stage_flag = match awaited_action.state().stage { |
392 | 0 | ActionStage::Unknown => OperationStageFlags::Any, |
393 | 0 | ActionStage::CacheCheck => OperationStageFlags::CacheCheck, |
394 | 42 | ActionStage::Queued => OperationStageFlags::Queued, |
395 | 0 | ActionStage::Executing => OperationStageFlags::Executing, |
396 | | ActionStage::Completed(_) | ActionStage::CompletedFromCache(_) => { |
397 | 0 | OperationStageFlags::Completed |
398 | | } |
399 | | }; |
400 | 42 | if !filter.stages.intersects(stage_flag) { Branch (400:20): [True: 0, False: 0]
Branch (400:20): [True: 0, False: 0]
Branch (400:20): [Folded - Ignored]
Branch (400:20): [True: 0, False: 40]
Branch (400:20): [True: 0, False: 2]
|
401 | 0 | return false; |
402 | 42 | } |
403 | 3 | } |
404 | | } |
405 | | |
406 | 45 | true |
407 | 45 | } |
408 | | |
409 | | /// Let the scheduler know that an operation has timed out from |
410 | | /// the client side (ie: worker has not updated in a while). |
411 | 1 | async fn timeout_operation_id(&self, operation_id: &OperationId) -> Result<(), Error> { |
412 | | // Ensure that only one timeout operation is running at a time. |
413 | | // Failing to do this could result in the same operation being |
414 | | // timed out multiple times at the same time. |
415 | | // Note: We could implement this on a per-operation_id basis, but it is quite |
416 | | // complex to manage the locks. |
417 | 1 | let _lock = self.timeout_operation_mux.lock().await; |
418 | | |
419 | 1 | let awaited_action_subscriber = self |
420 | 1 | .action_db |
421 | 1 | .get_by_operation_id(operation_id) |
422 | 1 | .await |
423 | 1 | .err_tip(|| "In SimpleSchedulerStateManager::timeout_operation_id"0 )?0 |
424 | 1 | .err_tip(|| { |
425 | 0 | format!("Operation id {operation_id} does not exist in SimpleSchedulerStateManager::timeout_operation_id") |
426 | 1 | })?0 ; |
427 | | |
428 | 1 | let awaited_action = awaited_action_subscriber |
429 | 1 | .borrow() |
430 | 1 | .await |
431 | 1 | .err_tip(|| "In SimpleSchedulerStateManager::timeout_operation_id"0 )?0 ; |
432 | | |
433 | | // If the action is not executing, we should not timeout the action. |
434 | 1 | if !matches!0 (awaited_action.state().stage, ActionStage::Executing) { Branch (434:12): [True: 0, False: 0]
Branch (434:12): [True: 0, False: 0]
Branch (434:12): [Folded - Ignored]
Branch (434:12): [True: 0, False: 1]
Branch (434:12): [True: 0, False: 0]
|
435 | 0 | return Ok(()); |
436 | 1 | } |
437 | | |
438 | 1 | let last_worker_updated = awaited_action |
439 | 1 | .last_worker_updated_timestamp() |
440 | 1 | .duration_since(SystemTime::UNIX_EPOCH) |
441 | 1 | .map_err(|e| { |
442 | 0 | make_err!( |
443 | 0 | Code::Internal, |
444 | 0 | "Failed to convert last_worker_updated to duration since epoch {e:?}" |
445 | 0 | ) |
446 | 1 | })?0 ; |
447 | 1 | let worker_should_update_before = last_worker_updated |
448 | 1 | .checked_add(self.no_event_action_timeout) |
449 | 1 | .err_tip(|| "Timestamp too big in SimpleSchedulerStateManager::timeout_operation_id"0 )?0 ; |
450 | 1 | if worker_should_update_before < (self.now_fn)().elapsed() { Branch (450:12): [True: 0, False: 0]
Branch (450:12): [True: 0, False: 0]
Branch (450:12): [Folded - Ignored]
Branch (450:12): [True: 0, False: 1]
Branch (450:12): [True: 0, False: 0]
|
451 | | // The action was updated recently, we should not timeout the action. |
452 | | // This is to prevent timing out actions that have recently been updated |
453 | | // (like multiple clients timeout the same action at the same time). |
454 | 0 | return Ok(()); |
455 | 1 | } |
456 | 1 | |
457 | 1 | self.assign_operation( |
458 | 1 | operation_id, |
459 | 1 | Err(make_err!( |
460 | 1 | Code::DeadlineExceeded, |
461 | 1 | "Operation timed out after {} seconds", |
462 | 1 | self.no_event_action_timeout.as_secs_f32(), |
463 | 1 | )), |
464 | 1 | ) |
465 | 1 | .await |
466 | 1 | } |
467 | | |
468 | 41 | async fn inner_update_operation( |
469 | 41 | &self, |
470 | 41 | operation_id: &OperationId, |
471 | 41 | maybe_worker_id: Option<&WorkerId>, |
472 | 41 | update: UpdateOperationType, |
473 | 41 | ) -> Result<(), Error> { |
474 | 41 | let mut last_err = None; |
475 | 42 | for _ in 0..MAX_UPDATE_RETRIES { |
476 | 42 | let maybe_awaited_action_subscriber = self |
477 | 42 | .action_db |
478 | 42 | .get_by_operation_id(operation_id) |
479 | 42 | .await |
480 | 42 | .err_tip(|| "In SimpleSchedulerStateManager::update_operation"0 )?0 ; |
481 | 42 | let Some(awaited_action_subscriber41 ) = maybe_awaited_action_subscriber else { Branch (481:17): [True: 0, False: 0]
Branch (481:17): [True: 0, False: 0]
Branch (481:17): [Folded - Ignored]
Branch (481:17): [True: 39, False: 0]
Branch (481:17): [True: 2, False: 1]
|
482 | | // No action found. It is ok if the action was not found. It |
483 | | // probably means that the action was dropped, but worker was |
484 | | // still processing it. |
485 | 1 | return Ok(()); |
486 | | }; |
487 | | |
488 | 41 | let mut awaited_action = awaited_action_subscriber |
489 | 41 | .borrow() |
490 | 41 | .await |
491 | 41 | .err_tip(|| "In SimpleSchedulerStateManager::update_operation"0 )?0 ; |
492 | | |
493 | | // Make sure the worker id matches the awaited action worker id. |
494 | | // This might happen if the worker sending the update is not the |
495 | | // worker that was assigned. |
496 | 41 | if awaited_action.worker_id().is_some() Branch (496:16): [True: 0, False: 0]
Branch (496:16): [True: 0, False: 0]
Branch (496:16): [Folded - Ignored]
Branch (496:16): [True: 13, False: 26]
Branch (496:16): [True: 0, False: 2]
|
497 | 13 | && maybe_worker_id.is_some() Branch (497:20): [True: 0, False: 0]
Branch (497:20): [True: 0, False: 0]
Branch (497:20): [Folded - Ignored]
Branch (497:20): [True: 12, False: 1]
Branch (497:20): [True: 0, False: 0]
|
498 | 12 | && maybe_worker_id != awaited_action.worker_id().as_ref() Branch (498:20): [True: 0, False: 0]
Branch (498:20): [True: 0, False: 0]
Branch (498:20): [Folded - Ignored]
Branch (498:20): [True: 0, False: 12]
Branch (498:20): [True: 0, False: 0]
|
499 | | { |
500 | | // If another worker is already assigned to the action, another |
501 | | // worker probably picked up the action. We should not update the |
502 | | // action in this case and abort this operation. |
503 | 0 | let err = make_err!( |
504 | 0 | Code::Aborted, |
505 | 0 | "Worker ids do not match - {:?} != {:?} for {:?}", |
506 | 0 | maybe_worker_id, |
507 | 0 | awaited_action.worker_id(), |
508 | 0 | awaited_action, |
509 | 0 | ); |
510 | 0 | event!( |
511 | 0 | Level::INFO, |
512 | 0 | "Worker ids do not match - {:?} != {:?} for {:?}. This is probably due to another worker picking up the action.", |
513 | 0 | maybe_worker_id, |
514 | 0 | awaited_action.worker_id(), |
515 | | awaited_action, |
516 | | ); |
517 | 0 | return Err(err); |
518 | 41 | } |
519 | 41 | |
520 | 41 | // Make sure we don't update an action that is already completed. |
521 | 41 | if awaited_action.state().stage.is_finished() { Branch (521:16): [True: 0, False: 0]
Branch (521:16): [True: 0, False: 0]
Branch (521:16): [Folded - Ignored]
Branch (521:16): [True: 0, False: 39]
Branch (521:16): [True: 0, False: 2]
|
522 | 0 | return Err(make_err!( |
523 | 0 | Code::Internal, |
524 | 0 | "Action {operation_id:?} is already completed with state {:?} - maybe_worker_id: {:?}", |
525 | 0 | awaited_action.state().stage, |
526 | 0 | maybe_worker_id, |
527 | 0 | )); |
528 | 41 | } |
529 | | |
530 | 41 | let stage = match &update { |
531 | | UpdateOperationType::KeepAlive => { |
532 | 0 | awaited_action.worker_keep_alive((self.now_fn)().now()); |
533 | 0 | return self |
534 | 0 | .action_db |
535 | 0 | .update_awaited_action(awaited_action) |
536 | 0 | .await |
537 | 0 | .err_tip(|| "Failed to send KeepAlive in SimpleSchedulerStateManager::update_operation"); |
538 | | } |
539 | 33 | UpdateOperationType::UpdateWithActionStage(stage) => stage.clone(), |
540 | 8 | UpdateOperationType::UpdateWithError(err) => { |
541 | 8 | // Don't count a backpressure failure as an attempt for an action. |
542 | 8 | let due_to_backpressure = err.code == Code::ResourceExhausted; |
543 | 8 | if !due_to_backpressure { Branch (543:24): [True: 0, False: 0]
Branch (543:24): [True: 0, False: 0]
Branch (543:24): [Folded - Ignored]
Branch (543:24): [True: 8, False: 0]
Branch (543:24): [True: 0, False: 0]
|
544 | 8 | awaited_action.attempts += 1; |
545 | 8 | }0 |
546 | | |
547 | 8 | if awaited_action.attempts > self.max_job_retries { Branch (547:24): [True: 0, False: 0]
Branch (547:24): [True: 0, False: 0]
Branch (547:24): [Folded - Ignored]
Branch (547:24): [True: 1, False: 7]
Branch (547:24): [True: 0, False: 0]
|
548 | 1 | ActionStage::Completed(ActionResult { |
549 | 1 | execution_metadata: ExecutionMetadata { |
550 | 1 | worker: maybe_worker_id.map_or_else(String::default, ToString::to_string), |
551 | 1 | ..ExecutionMetadata::default() |
552 | 1 | }, |
553 | 1 | error: Some(err.clone().merge(make_err!( |
554 | 1 | Code::Internal, |
555 | 1 | "Job cancelled because it attempted to execute too many times {} > {} times {}", |
556 | 1 | awaited_action.attempts, |
557 | 1 | self.max_job_retries, |
558 | 1 | format!("for operation_id: {operation_id}, maybe_worker_id: {maybe_worker_id:?}"), |
559 | 1 | ))), |
560 | 1 | ..ActionResult::default() |
561 | 1 | }) |
562 | | } else { |
563 | 7 | ActionStage::Queued |
564 | | } |
565 | | } |
566 | | }; |
567 | 41 | let now = (self.now_fn)().now(); |
568 | 41 | if matches!34 (stage, ActionStage::Queued) { |
569 | 7 | // If the action is queued, we need to unset the worker id regardless of |
570 | 7 | // which worker sent the update. |
571 | 7 | awaited_action.set_worker_id(None, now); |
572 | 34 | } else { |
573 | 34 | awaited_action.set_worker_id(maybe_worker_id.copied(), now); |
574 | 34 | } |
575 | 41 | awaited_action.worker_set_state( |
576 | 41 | Arc::new(ActionState { |
577 | 41 | stage, |
578 | 41 | // Client id is not known here, it is the responsibility of |
579 | 41 | // the the subscriber impl to replace this with the |
580 | 41 | // correct client id. |
581 | 41 | client_operation_id: operation_id.clone(), |
582 | 41 | action_digest: awaited_action.action_info().digest(), |
583 | 41 | }), |
584 | 41 | now, |
585 | 41 | ); |
586 | | |
587 | 41 | let update_action_result = self |
588 | 41 | .action_db |
589 | 41 | .update_awaited_action(awaited_action) |
590 | 41 | .await |
591 | 41 | .err_tip(|| "In SimpleSchedulerStateManager::update_operation"2 ); |
592 | 41 | if let Err(err2 ) = update_action_result { Branch (592:20): [True: 0, False: 0]
Branch (592:20): [True: 0, False: 0]
Branch (592:20): [Folded - Ignored]
Branch (592:20): [True: 0, False: 39]
Branch (592:20): [True: 2, False: 0]
|
593 | | // We use Aborted to signal that the action was not |
594 | | // updated due to the data being set was not the latest |
595 | | // but can be retried. |
596 | 2 | if err.code == Code::Aborted { Branch (596:20): [True: 0, False: 0]
Branch (596:20): [True: 0, False: 0]
Branch (596:20): [Folded - Ignored]
Branch (596:20): [True: 0, False: 0]
Branch (596:20): [True: 1, False: 1]
|
597 | 1 | last_err = Some(err); |
598 | 1 | continue; |
599 | 1 | } |
600 | 1 | return Err(err); |
601 | 39 | } |
602 | 39 | return Ok(()); |
603 | | } |
604 | 0 | match last_err { |
605 | 0 | Some(err) => Err(err), |
606 | 0 | None => Err(make_err!( |
607 | 0 | Code::Internal, |
608 | 0 | "Failed to update action after {} retries with no error set", |
609 | 0 | MAX_UPDATE_RETRIES, |
610 | 0 | )), |
611 | | } |
612 | 41 | } |
613 | | |
614 | 25 | async fn inner_add_operation( |
615 | 25 | &self, |
616 | 25 | new_client_operation_id: OperationId, |
617 | 25 | action_info: Arc<ActionInfo>, |
618 | 25 | ) -> Result<T::Subscriber, Error> { |
619 | 25 | self.action_db |
620 | 25 | .add_action(new_client_operation_id, action_info) |
621 | 25 | .await |
622 | 25 | .err_tip(|| "In SimpleSchedulerStateManager::add_operation"0 ) |
623 | 25 | } |
624 | | |
625 | 594 | async fn inner_filter_operations<'a, F>( |
626 | 594 | &'a self, |
627 | 594 | filter: OperationFilter, |
628 | 594 | to_action_state_result: F, |
629 | 594 | ) -> Result<ActionStateResultStream<'a>, Error> |
630 | 594 | where |
631 | 594 | F: Fn(T::Subscriber) -> Box<dyn ActionStateResult> + Send + Sync + 'a, |
632 | 594 | { |
633 | 91 | fn sorted_awaited_action_state_for_flags( |
634 | 91 | stage: OperationStageFlags, |
635 | 91 | ) -> Option<SortedAwaitedActionState> { |
636 | 91 | match stage { |
637 | 0 | OperationStageFlags::CacheCheck => Some(SortedAwaitedActionState::CacheCheck), |
638 | 91 | OperationStageFlags::Queued => Some(SortedAwaitedActionState::Queued), |
639 | 0 | OperationStageFlags::Executing => Some(SortedAwaitedActionState::Executing), |
640 | 0 | OperationStageFlags::Completed => Some(SortedAwaitedActionState::Completed), |
641 | 0 | _ => None, |
642 | | } |
643 | 91 | } |
644 | | |
645 | 594 | if let Some(operation_id0 ) = &filter.operation_id { Branch (645:16): [True: 0, False: 0]
Branch (645:16): [True: 0, False: 0]
Branch (645:16): [True: 0, False: 0]
Branch (645:16): [True: 0, False: 0]
Branch (645:16): [Folded - Ignored]
Branch (645:16): [True: 0, False: 503]
Branch (645:16): [True: 0, False: 87]
Branch (645:16): [True: 0, False: 0]
Branch (645:16): [True: 0, False: 4]
|
646 | 0 | let maybe_subscriber = self |
647 | 0 | .action_db |
648 | 0 | .get_by_operation_id(operation_id) |
649 | 0 | .await |
650 | 0 | .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?; |
651 | 0 | let Some(subscriber) = maybe_subscriber else { Branch (651:17): [True: 0, False: 0]
Branch (651:17): [True: 0, False: 0]
Branch (651:17): [True: 0, False: 0]
Branch (651:17): [True: 0, False: 0]
Branch (651:17): [Folded - Ignored]
Branch (651:17): [True: 0, False: 0]
Branch (651:17): [True: 0, False: 0]
Branch (651:17): [True: 0, False: 0]
Branch (651:17): [True: 0, False: 0]
|
652 | 0 | return Ok(Box::pin(stream::empty())); |
653 | | }; |
654 | 0 | let awaited_action = subscriber |
655 | 0 | .borrow() |
656 | 0 | .await |
657 | 0 | .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?; |
658 | 0 | if !self.apply_filter_predicate(&awaited_action, &filter).await { Branch (658:16): [True: 0, False: 0]
Branch (658:16): [True: 0, False: 0]
Branch (658:16): [True: 0, False: 0]
Branch (658:16): [True: 0, False: 0]
Branch (658:16): [Folded - Ignored]
Branch (658:16): [True: 0, False: 0]
Branch (658:16): [True: 0, False: 0]
Branch (658:16): [True: 0, False: 0]
Branch (658:16): [True: 0, False: 0]
|
659 | 0 | return Ok(Box::pin(stream::empty())); |
660 | 0 | } |
661 | 0 | return Ok(Box::pin(stream::once(async move { |
662 | 0 | to_action_state_result(subscriber) |
663 | 0 | }))); |
664 | 594 | } |
665 | 594 | if let Some(client_operation_id503 ) = &filter.client_operation_id { Branch (665:16): [True: 0, False: 0]
Branch (665:16): [True: 0, False: 0]
Branch (665:16): [True: 0, False: 0]
Branch (665:16): [True: 0, False: 0]
Branch (665:16): [Folded - Ignored]
Branch (665:16): [True: 503, False: 0]
Branch (665:16): [True: 0, False: 87]
Branch (665:16): [True: 0, False: 0]
Branch (665:16): [True: 0, False: 4]
|
666 | 503 | let maybe_subscriber = self |
667 | 503 | .action_db |
668 | 503 | .get_awaited_action_by_id(client_operation_id) |
669 | 503 | .await |
670 | 503 | .err_tip(|| "In SimpleSchedulerStateManager::filter_operations"0 )?0 ; |
671 | 503 | let Some(subscriber3 ) = maybe_subscriber else { Branch (671:17): [True: 0, False: 0]
Branch (671:17): [True: 0, False: 0]
Branch (671:17): [True: 0, False: 0]
Branch (671:17): [True: 0, False: 0]
Branch (671:17): [Folded - Ignored]
Branch (671:17): [True: 3, False: 500]
Branch (671:17): [True: 0, False: 0]
Branch (671:17): [True: 0, False: 0]
Branch (671:17): [True: 0, False: 0]
|
672 | 500 | return Ok(Box::pin(stream::empty())); |
673 | | }; |
674 | 3 | let awaited_action = subscriber |
675 | 3 | .borrow() |
676 | 3 | .await |
677 | 3 | .err_tip(|| "In SimpleSchedulerStateManager::filter_operations"0 )?0 ; |
678 | 3 | if !self.apply_filter_predicate(&awaited_action, &filter).await { Branch (678:16): [True: 0, False: 0]
Branch (678:16): [True: 0, False: 0]
Branch (678:16): [True: 0, False: 0]
Branch (678:16): [True: 0, False: 0]
Branch (678:16): [Folded - Ignored]
Branch (678:16): [True: 0, False: 3]
Branch (678:16): [True: 0, False: 0]
Branch (678:16): [True: 0, False: 0]
Branch (678:16): [True: 0, False: 0]
|
679 | 0 | return Ok(Box::pin(stream::empty())); |
680 | 3 | } |
681 | 3 | return Ok(Box::pin(stream::once(async move { |
682 | 3 | to_action_state_result(subscriber) |
683 | 3 | }))); |
684 | 91 | } |
685 | | |
686 | 91 | let Some(sorted_awaited_action_state) = Branch (686:13): [True: 0, False: 0]
Branch (686:13): [True: 0, False: 0]
Branch (686:13): [True: 0, False: 0]
Branch (686:13): [True: 0, False: 0]
Branch (686:13): [Folded - Ignored]
Branch (686:13): [True: 0, False: 0]
Branch (686:13): [True: 87, False: 0]
Branch (686:13): [True: 0, False: 0]
Branch (686:13): [True: 4, False: 0]
|
687 | 91 | sorted_awaited_action_state_for_flags(filter.stages) |
688 | | else { |
689 | 0 | let mut all_items: Vec<_> = self |
690 | 0 | .action_db |
691 | 0 | .get_all_awaited_actions() |
692 | 0 | .await |
693 | 0 | .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")? |
694 | 0 | .and_then(|awaited_action_subscriber| async move { |
695 | 0 | let awaited_action = awaited_action_subscriber |
696 | 0 | .borrow() |
697 | 0 | .await |
698 | 0 | .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?; |
699 | 0 | Ok((awaited_action_subscriber, awaited_action)) |
700 | 0 | }) |
701 | 0 | .try_filter_map(|(subscriber, awaited_action)| { |
702 | 0 | let filter = filter.clone(); |
703 | 0 | async move { |
704 | 0 | if self.apply_filter_predicate(&awaited_action, &filter).await { Branch (704:28): [True: 0, False: 0]
Branch (704:28): [True: 0, False: 0]
Branch (704:28): [True: 0, False: 0]
Branch (704:28): [True: 0, False: 0]
Branch (704:28): [Folded - Ignored]
Branch (704:28): [True: 0, False: 0]
Branch (704:28): [True: 0, False: 0]
Branch (704:28): [True: 0, False: 0]
Branch (704:28): [True: 0, False: 0]
|
705 | 0 | Ok(Some((subscriber, awaited_action.sort_key()))) |
706 | | } else { |
707 | 0 | Ok(None) |
708 | | } |
709 | 0 | } |
710 | 0 | }) |
711 | 0 | .try_collect() |
712 | 0 | .await |
713 | 0 | .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?; |
714 | 0 | match filter.order_by_priority_direction { |
715 | 0 | Some(OrderDirection::Asc) => all_items.sort_unstable_by(|(_, a), (_, b)| a.cmp(b)), |
716 | 0 | Some(OrderDirection::Desc) => all_items.sort_unstable_by(|(_, a), (_, b)| b.cmp(a)), |
717 | 0 | None => {} |
718 | | } |
719 | 0 | return Ok(Box::pin(stream::iter( |
720 | 0 | all_items |
721 | 0 | .into_iter() |
722 | 0 | .map(move |(subscriber, _)| to_action_state_result(subscriber)), |
723 | 0 | ))); |
724 | | }; |
725 | | |
726 | 91 | let desc = matches!0 ( |
727 | 91 | filter.order_by_priority_direction, |
728 | | Some(OrderDirection::Desc) |
729 | | ); |
730 | 91 | let stream = self |
731 | 91 | .action_db |
732 | 91 | .get_range_of_actions( |
733 | 91 | sorted_awaited_action_state, |
734 | 91 | Bound::Unbounded, |
735 | 91 | Bound::Unbounded, |
736 | 91 | desc, |
737 | 91 | ) |
738 | 91 | .await |
739 | 91 | .err_tip(|| "In SimpleSchedulerStateManager::filter_operations"0 )?0 |
740 | 91 | .and_then(|awaited_action_subscriber| async move {42 |
741 | 42 | let awaited_action = awaited_action_subscriber |
742 | 42 | .borrow() |
743 | 42 | .await |
744 | 42 | .err_tip(|| "In SimpleSchedulerStateManager::filter_operations"0 )?0 ; |
745 | 42 | Ok((awaited_action_subscriber, awaited_action)) |
746 | 91 | }84 ) |
747 | 91 | .try_filter_map(move |(subscriber, awaited_action)| { |
748 | 42 | let filter = filter.clone(); |
749 | 42 | async move { |
750 | 42 | if self.apply_filter_predicate(&awaited_action, &filter).await { Branch (750:24): [True: 0, False: 0]
Branch (750:24): [True: 0, False: 0]
Branch (750:24): [True: 0, False: 0]
Branch (750:24): [True: 0, False: 0]
Branch (750:24): [Folded - Ignored]
Branch (750:24): [True: 0, False: 0]
Branch (750:24): [True: 40, False: 0]
Branch (750:24): [True: 0, False: 0]
Branch (750:24): [True: 2, False: 0]
|
751 | 42 | Ok(Some(subscriber)) |
752 | | } else { |
753 | 0 | Ok(None) |
754 | | } |
755 | 42 | } |
756 | 91 | }42 ) |
757 | 91 | .map(move |result| -> Box<dyn ActionStateResult> { |
758 | 42 | result.map_or_else( |
759 | 42 | |e| -> Box<dyn ActionStateResult> { Box::new(ErrorActionStateResult(e)) }0 , |
760 | 42 | |v| -> Box<dyn ActionStateResult> { to_action_state_result(v) }, |
761 | 42 | ) |
762 | 91 | }); |
763 | 91 | Ok(Box::pin(stream)) |
764 | 594 | } |
765 | | } |
766 | | |
767 | | #[async_trait] |
768 | | impl<T, I, NowFn> ClientStateManager for SimpleSchedulerStateManager<T, I, NowFn> |
769 | | where |
770 | | T: AwaitedActionDb, |
771 | | I: InstantWrapper, |
772 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
773 | | { |
774 | | async fn add_action( |
775 | | &self, |
776 | | client_operation_id: OperationId, |
777 | | action_info: Arc<ActionInfo>, |
778 | 25 | ) -> Result<Box<dyn ActionStateResult>, Error> { |
779 | 25 | let sub = self |
780 | 25 | .inner_add_operation(client_operation_id.clone(), action_info.clone()) |
781 | 25 | .await?0 ; |
782 | | |
783 | 25 | Ok(Box::new(ClientActionStateResult::new( |
784 | 25 | sub, |
785 | 25 | self.weak_self.clone(), |
786 | 25 | self.no_event_action_timeout, |
787 | 25 | self.now_fn.clone(), |
788 | 25 | ))) |
789 | 50 | } |
790 | | |
791 | | async fn filter_operations<'a>( |
792 | | &'a self, |
793 | | filter: OperationFilter, |
794 | 503 | ) -> Result<ActionStateResultStream<'a>, Error> { |
795 | 503 | self.inner_filter_operations(filter, move |rx| { |
796 | 3 | Box::new(ClientActionStateResult::new( |
797 | 3 | rx, |
798 | 3 | self.weak_self.clone(), |
799 | 3 | self.no_event_action_timeout, |
800 | 3 | self.now_fn.clone(), |
801 | 3 | )) |
802 | 503 | }) |
803 | 503 | .await |
804 | 1.00k | } |
805 | | |
806 | 0 | fn as_known_platform_property_provider(&self) -> Option<&dyn KnownPlatformPropertyProvider> { |
807 | 0 | None |
808 | 0 | } |
809 | | } |
810 | | |
811 | | #[async_trait] |
812 | | impl<T, I, NowFn> WorkerStateManager for SimpleSchedulerStateManager<T, I, NowFn> |
813 | | where |
814 | | T: AwaitedActionDb, |
815 | | I: InstantWrapper, |
816 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
817 | | { |
818 | | async fn update_operation( |
819 | | &self, |
820 | | operation_id: &OperationId, |
821 | | worker_id: &WorkerId, |
822 | | update: UpdateOperationType, |
823 | 12 | ) -> Result<(), Error> { |
824 | 12 | self.inner_update_operation(operation_id, Some(worker_id), update) |
825 | 12 | .await |
826 | 24 | } |
827 | | } |
828 | | |
829 | | #[async_trait] |
830 | | impl<T, I, NowFn> MatchingEngineStateManager for SimpleSchedulerStateManager<T, I, NowFn> |
831 | | where |
832 | | T: AwaitedActionDb, |
833 | | I: InstantWrapper, |
834 | | NowFn: Fn() -> I + Clone + Send + Unpin + Sync + 'static, |
835 | | { |
836 | | async fn filter_operations<'a>( |
837 | | &'a self, |
838 | | filter: OperationFilter, |
839 | 91 | ) -> Result<ActionStateResultStream<'a>, Error> { |
840 | 91 | self.inner_filter_operations(filter, |rx| { |
841 | 42 | Box::new(MatchingEngineActionStateResult::new( |
842 | 42 | rx, |
843 | 42 | self.weak_self.clone(), |
844 | 42 | self.no_event_action_timeout, |
845 | 42 | self.now_fn.clone(), |
846 | 42 | )) |
847 | 91 | }) |
848 | 91 | .await |
849 | 182 | } |
850 | | |
851 | | async fn assign_operation( |
852 | | &self, |
853 | | operation_id: &OperationId, |
854 | | worker_id_or_reason_for_unassign: Result<&WorkerId, Error>, |
855 | 29 | ) -> Result<(), Error> { |
856 | 29 | let (maybe_worker_id, update) = match worker_id_or_reason_for_unassign { |
857 | 28 | Ok(worker_id) => ( |
858 | 28 | Some(worker_id), |
859 | 28 | UpdateOperationType::UpdateWithActionStage(ActionStage::Executing), |
860 | 28 | ), |
861 | 1 | Err(err) => (None, UpdateOperationType::UpdateWithError(err)), |
862 | | }; |
863 | 29 | self.inner_update_operation(operation_id, maybe_worker_id, update) |
864 | 29 | .await |
865 | 58 | } |
866 | | } |