Coverage Report

Created: 2025-05-30 16:37

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/build/source/nativelink-worker/src/local_worker.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 core::pin::Pin;
16
use core::str;
17
use core::sync::atomic::{AtomicU64, Ordering};
18
use core::time::Duration;
19
use std::process::Stdio;
20
use std::sync::{Arc, Weak};
21
22
use futures::future::BoxFuture;
23
use futures::stream::FuturesUnordered;
24
use futures::{Future, FutureExt, StreamExt, TryFutureExt, select};
25
use nativelink_config::cas_server::LocalWorkerConfig;
26
use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err};
27
use nativelink_metric::{MetricsComponent, RootMetricsComponent};
28
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_worker::Update;
29
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::worker_api_client::WorkerApiClient;
30
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{
31
    ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForWorker, execute_result,
32
};
33
use nativelink_store::fast_slow_store::FastSlowStore;
34
use nativelink_util::action_messages::{ActionResult, ActionStage, OperationId};
35
use nativelink_util::common::fs;
36
use nativelink_util::digest_hasher::DigestHasherFunc;
37
use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime};
38
use nativelink_util::shutdown_guard::ShutdownGuard;
39
use nativelink_util::store_trait::Store;
40
use nativelink_util::{spawn, tls_utils};
41
use opentelemetry::context::Context;
42
use tokio::process;
43
use tokio::sync::{broadcast, mpsc};
44
use tokio::time::sleep;
45
use tokio_stream::wrappers::UnboundedReceiverStream;
46
use tonic::Streaming;
47
use tracing::{Level, debug, error, info, info_span, instrument, warn};
48
49
use crate::running_actions_manager::{
50
    ExecutionConfiguration, Metrics as RunningActionManagerMetrics, RunningAction,
51
    RunningActionsManager, RunningActionsManagerArgs, RunningActionsManagerImpl,
52
};
53
use crate::worker_api_client_wrapper::{WorkerApiClientTrait, WorkerApiClientWrapper};
54
use crate::worker_utils::make_connect_worker_request;
55
56
/// Amount of time to wait if we have actions in transit before we try to
57
/// consider an error to have occurred.
58
const ACTIONS_IN_TRANSIT_TIMEOUT_S: f32 = 10.;
59
60
/// If we lose connection to the worker api server we will wait this many seconds
61
/// before trying to connect.
62
const CONNECTION_RETRY_DELAY_S: f32 = 0.5;
63
64
/// Default endpoint timeout. If this value gets modified the documentation in
65
/// `cas_server.rs` must also be updated.
66
const DEFAULT_ENDPOINT_TIMEOUT_S: f32 = 5.;
67
68
/// Default maximum amount of time a task is allowed to run for.
69
/// If this value gets modified the documentation in `cas_server.rs` must also be updated.
70
const DEFAULT_MAX_ACTION_TIMEOUT: Duration = Duration::from_secs(1200); // 20 mins.
71
72
struct LocalWorkerImpl<'a, T: WorkerApiClientTrait, U: RunningActionsManager> {
73
    config: &'a LocalWorkerConfig,
74
    // According to the tonic documentation it is a cheap operation to clone this.
75
    grpc_client: T,
76
    worker_id: String,
77
    running_actions_manager: Arc<U>,
78
    // Number of actions that have been received in `Update::StartAction`, but
79
    // not yet processed by running_actions_manager's spawn. This number should
80
    // always be zero if there are no actions running and no actions being waited
81
    // on by the scheduler.
82
    actions_in_transit: Arc<AtomicU64>,
83
    metrics: Arc<Metrics>,
84
}
85
86
4
async fn preconditions_met(precondition_script: Option<String>) -> Result<(), Error> {
87
4
    let Some(
precondition_script1
) = &precondition_script else {
  Branch (87:9): [True: 1, False: 3]
  Branch (87:9): [Folded - Ignored]
  Branch (87:9): [Folded - Ignored]
88
        // No script means we are always ok to proceed.
89
3
        return Ok(());
90
    };
91
    // TODO: Might want to pass some information about the command to the
92
    //       script, but at this point it's not even been downloaded yet,
93
    //       so that's not currently possible.  Perhaps we'll move this in
94
    //       future to pass useful information through?  Or perhaps we'll
95
    //       have a pre-condition and a pre-execute script instead, although
96
    //       arguably entrypoint already gives us that.
97
1
    let precondition_process = process::Command::new(precondition_script)
98
1
        .kill_on_drop(true)
99
1
        .stdin(Stdio::null())
100
1
        .stdout(Stdio::piped())
101
1
        .stderr(Stdio::null())
102
1
        .env_clear()
103
1
        .spawn()
104
1
        .err_tip(|| format!(
"Could not execute precondition command {precondition_script:?}"0
))
?0
;
105
1
    let output = precondition_process.wait_with_output().await
?0
;
106
1
    if output.status.code() == Some(0) {
  Branch (106:8): [True: 0, False: 1]
  Branch (106:8): [Folded - Ignored]
  Branch (106:8): [Folded - Ignored]
107
0
        Ok(())
108
    } else {
109
1
        Err(make_err!(
110
1
            Code::ResourceExhausted,
111
1
            "Preconditions script returned status {} - {}",
112
1
            output.status,
113
1
            str::from_utf8(&output.stdout).unwrap_or("")
114
1
        ))
115
    }
116
4
}
117
118
impl<'a, T: WorkerApiClientTrait, U: RunningActionsManager> LocalWorkerImpl<'a, T, U> {
119
5
    fn new(
120
5
        config: &'a LocalWorkerConfig,
121
5
        grpc_client: T,
122
5
        worker_id: String,
123
5
        running_actions_manager: Arc<U>,
124
5
        metrics: Arc<Metrics>,
125
5
    ) -> Self {
126
5
        Self {
127
5
            config,
128
5
            grpc_client,
129
5
            worker_id,
130
5
            running_actions_manager,
131
5
            // Number of actions that have been received in `Update::StartAction`, but
132
5
            // not yet processed by running_actions_manager's spawn. This number should
133
5
            // always be zero if there are no actions running and no actions being waited
134
5
            // on by the scheduler.
135
5
            actions_in_transit: Arc::new(AtomicU64::new(0)),
136
5
            metrics,
137
5
        }
138
5
    }
139
140
    /// Starts a background spawn/thread that will send a message to the server every `timeout / 2`.
141
5
    async fn start_keep_alive(&self) -> Result<(), Error> 
{4
142
        // According to tonic's documentation this call should be cheap and is the same stream.
143
4
        let mut grpc_client = self.grpc_client.clone();
144
145
        loop {
146
4
            let timeout = self
147
4
                .config
148
4
                .worker_api_endpoint
149
4
                .timeout
150
4
                .unwrap_or(DEFAULT_ENDPOINT_TIMEOUT_S);
151
            // We always send 2 keep alive requests per timeout. Http2 should manage most of our
152
            // timeout issues, this is a secondary check to ensure we can still send data.
153
4
            sleep(Duration::from_secs_f32(timeout / 2.)).await;
154
0
            if let Err(e) = grpc_client
  Branch (154:20): [True: 0, False: 0]
  Branch (154:20): [Folded - Ignored]
  Branch (154:20): [Folded - Ignored]
155
0
                .keep_alive(KeepAliveRequest {
156
0
                    worker_id: self.worker_id.clone(),
157
0
                })
158
0
                .await
159
            {
160
0
                return Err(make_err!(
161
0
                    Code::Internal,
162
0
                    "Failed to send KeepAlive in LocalWorker : {:?}",
163
0
                    e
164
0
                ));
165
0
            }
166
        }
167
0
    }
168
169
5
    async fn run(
170
5
        &self,
171
5
        update_for_worker_stream: Streaming<UpdateForWorker>,
172
5
        shutdown_rx: &mut broadcast::Receiver<ShutdownGuard>,
173
5
    ) -> Result<(), Error> {
174
        // This big block of logic is designed to help simplify upstream components. Upstream
175
        // components can write standard futures that return a `Result<(), Error>` and this block
176
        // will forward the error up to the client and disconnect from the scheduler.
177
        // It is a common use case that an item sent through update_for_worker_stream will always
178
        // have a response but the response will be triggered through a callback to the scheduler.
179
        // This can be quite tricky to manage, so what we have done here is given access to a
180
        // `futures` variable which because this is in a single thread as well as a channel that you
181
        // send a future into that makes it into the `futures` variable.
182
        // This means that if you want to perform an action based on the result of the future
183
        // you use the `.map()` method and the new action will always come to live in this spawn,
184
        // giving mutable access to stuff in this struct.
185
        // NOTE: If you ever return from this function it will disconnect from the scheduler.
186
5
        let mut futures = FuturesUnordered::new();
187
5
        futures.push(self.start_keep_alive().boxed());
188
189
5
        let (add_future_channel, add_future_rx) = mpsc::unbounded_channel();
190
5
        let mut add_future_rx = UnboundedReceiverStream::new(add_future_rx).fuse();
191
192
5
        let mut update_for_worker_stream = update_for_worker_stream.fuse();
193
194
        loop {
195
16
            select! {
196
16
                
maybe_update6
= update_for_worker_stream.next() => {
197
6
                    match maybe_update
198
6
                        .err_tip(|| "UpdateForWorker stream closed early")
?1
199
5
                        .err_tip(|| "Got error in UpdateForWorker stream")
?0
200
                        .update
201
5
                        .err_tip(|| "Expected update to exist in UpdateForWorker")
?0
202
                    {
203
                        Update::ConnectionResult(_) => {
204
0
                            return Err(make_input_err!(
205
0
                                "Got ConnectionResult in LocalWorker::run which should never happen"
206
0
                            ));
207
                        }
208
                        // TODO(aaronmondal) We should possibly do something with this notification.
209
0
                        Update::Disconnect(()) => {
210
0
                            self.metrics.disconnects_received.inc();
211
0
                        }
212
0
                        Update::KeepAlive(()) => {
213
0
                            self.metrics.keep_alives_received.inc();
214
0
                        }
215
1
                        Update::KillOperationRequest(kill_operation_request) => {
216
1
                            let operation_id = OperationId::from(kill_operation_request.operation_id);
217
1
                            if let Err(
err0
) = self.running_actions_manager.kill_operation(&operation_id).await {
  Branch (217:36): [True: 0, False: 1]
  Branch (217:36): [Folded - Ignored]
  Branch (217:36): [Folded - Ignored]
218
0
                                error!(
219
                                    ?operation_id,
220
                                    ?err,
221
0
                                    "Failed to send kill request for operation"
222
                                );
223
1
                            }
224
                        }
225
4
                        Update::StartAction(start_execute) => {
226
4
                            self.metrics.start_actions_received.inc();
227
228
4
                            let execute_request = start_execute.execute_request.as_ref();
229
4
                            let operation_id = start_execute.operation_id.clone();
230
4
                            let maybe_instance_name = execute_request.map(|v| v.instance_name.clone());
231
4
                            let action_digest = execute_request.and_then(|v| v.action_digest.clone());
232
4
                            let digest_hasher = execute_request
233
4
                                .ok_or_else(|| make_input_err!("Expected execute_request to be set"))
234
4
                                .and_then(|v| DigestHasherFunc::try_from(v.digest_function))
235
4
                                .err_tip(|| "In LocalWorkerImpl::new()")
?0
;
236
237
4
                            let start_action_fut = {
238
4
                                let precondition_script_cfg = self.config.experimental_precondition_script.clone();
239
4
                                let actions_in_transit = self.actions_in_transit.clone();
240
4
                                let worker_id = self.worker_id.clone();
241
4
                                let running_actions_manager = self.running_actions_manager.clone();
242
4
                                self.metrics.clone().wrap(move |metrics| async move {
243
4
                                    metrics.preconditions.wrap(preconditions_met(precondition_script_cfg))
244
4
                                    .and_then(|()| 
running_actions_manager3
.
create_and_add_action3
(
worker_id3
,
start_execute3
))
245
4
                                    .map(move |r| {
246
                                        // Now that we either failed or registered our action, we can
247
                                        // consider the action to no longer be in transit.
248
4
                                        actions_in_transit.fetch_sub(1, Ordering::Release);
249
4
                                        r
250
4
                                    })
251
4
                                    .and_then(|action| 
{3
252
3
                                        debug!(
253
0
                                            operation_id = ?action.get_operation_id(),
254
0
                                            "Received request to run action"
255
                                        );
256
3
                                        action
257
3
                                            .clone()
258
3
                                            .prepare_action()
259
3
                                            .and_then(RunningAction::execute)
260
3
                                            .and_then(RunningAction::upload_results)
261
3
                                            .and_then(RunningAction::get_finished_result)
262
                                            // Note: We need ensure we run cleanup even if one of the other steps fail.
263
3
                                            .then(|result| async move 
{2
264
2
                                                if let Err(
e0
) = action.cleanup().await {
  Branch (264:56): [True: 0, False: 2]
  Branch (264:56): [Folded - Ignored]
  Branch (264:56): [Folded - Ignored]
265
0
                                                    return Result::<ActionResult, Error>::Err(e).merge(result);
266
2
                                                }
267
2
                                                result
268
4
                                            })
269
4
                                    
}3
).await
270
7
                                })
271
                            };
272
273
4
                            let make_publish_future = {
274
4
                                let mut grpc_client = self.grpc_client.clone();
275
276
4
                                let worker_id = self.worker_id.clone();
277
4
                                let running_actions_manager = self.running_actions_manager.clone();
278
3
                                move |res: Result<ActionResult, Error>| async move {
279
3
                                    let instance_name = maybe_instance_name
280
3
                                        .err_tip(|| "`instance_name` could not be resolved; this is likely an internal error in local_worker.")
?0
;
281
3
                                    match res {
282
2
                                        Ok(mut action_result) => {
283
                                            // Save in the action cache before notifying the scheduler that we've completed.
284
2
                                            if let Some(digest_info) = action_digest.clone().and_then(|action_digest| action_digest.try_into().ok()) {
  Branch (284:52): [True: 2, False: 0]
  Branch (284:52): [Folded - Ignored]
  Branch (284:52): [Folded - Ignored]
285
2
                                                if let Err(
err0
) = running_actions_manager.cache_action_result(digest_info, &mut action_result, digest_hasher).await {
  Branch (285:56): [True: 0, False: 2]
  Branch (285:56): [Folded - Ignored]
  Branch (285:56): [Folded - Ignored]
286
0
                                                    error!(
287
                                                        ?err,
288
                                                        ?action_digest,
289
0
                                                        "Error saving action in store",
290
                                                    );
291
2
                                                }
292
0
                                            }
293
2
                                            let action_stage = ActionStage::Completed(action_result);
294
2
                                            grpc_client.execution_response(
295
2
                                                ExecuteResult{
296
2
                                                    worker_id,
297
2
                                                    instance_name,
298
2
                                                    operation_id,
299
2
                                                    result: Some(execute_result::Result::ExecuteResponse(action_stage.into())),
300
2
                                                }
301
2
                                            )
302
2
                                            .await
303
0
                                            .err_tip(|| "Error while calling execution_response")?;
304
                                        },
305
1
                                        Err(e) => {
306
1
                                            grpc_client.execution_response(ExecuteResult{
307
1
                                                worker_id,
308
1
                                                instance_name,
309
1
                                                operation_id,
310
1
                                                result: Some(execute_result::Result::InternalError(e.into())),
311
1
                                            }).await.
err_tip0
(|| "Error calling execution_response with error")
?0
;
312
                                        },
313
                                    }
314
0
                                    Ok(())
315
3
                                }
316
                            };
317
318
4
                            self.actions_in_transit.fetch_add(1, Ordering::Release);
319
4
                            let futures_ref = &futures;
320
321
4
                            let add_future_channel = add_future_channel.clone();
322
323
4
                            info_span!(
324
                                "worker_start_action_ctx",
325
4
                                digest_function = %digest_hasher.to_string(),
326
4
                            ).in_scope(|| {
327
4
                                let _guard = Context::current_with_value(digest_hasher)
328
4
                                    .attach();
329
330
4
                                futures_ref.push(
331
4
                                    spawn!("worker_start_action", start_action_fut).map(move |res| 
{3
332
3
                                        let res = res.err_tip(|| "Failed to launch spawn")
?0
;
333
3
                                        if let Err(
err1
) = &res {
  Branch (333:48): [True: 1, False: 2]
  Branch (333:48): [Folded - Ignored]
  Branch (333:48): [Folded - Ignored]
334
1
                                            error!(?err, "Error executing action");
335
2
                                        }
336
3
                                        add_future_channel
337
3
                                            .send(make_publish_future(res).boxed())
338
3
                                            .map_err(|_| make_err!(
Code::Internal0
, "LocalWorker could not send future"))
?0
;
339
3
                                        Ok(())
340
3
                                    })
341
4
                                    .boxed()
342
                                );
343
4
                            });
344
                        }
345
                    }
346
                },
347
16
                
res3
= add_future_rx.next() => {
348
3
                    let fut = res.err_tip(|| "New future stream receives should never be closed")
?0
;
349
3
                    futures.push(fut);
350
                },
351
16
                
res3
= futures.next() =>
res3
.
err_tip3
(|| "Keep-alive should always pending. Likely unable to send data to scheduler")
?0
?0
,
352
16
                
complete_msg0
= shutdown_rx.recv().fuse() => {
353
0
                    warn!("Worker loop reveiced shutdown signal. Shutting down worker...",);
354
0
                    let mut grpc_client = self.grpc_client.clone();
355
0
                    let worker_id = self.worker_id.clone();
356
0
                    let running_actions_manager = self.running_actions_manager.clone();
357
0
                    let complete_msg_clone = complete_msg.map_err(|e| make_err!(Code::Internal, "Failed to receive shutdown message: {e:?}"))?.clone();
358
0
                    let shutdown_future = async move {
359
0
                        if let Err(e) = grpc_client.going_away(GoingAwayRequest { worker_id }).await {
  Branch (359:32): [True: 0, False: 0]
  Branch (359:32): [Folded - Ignored]
  Branch (359:32): [Folded - Ignored]
360
0
                            error!("Failed to send GoingAwayRequest: {e}",);
361
0
                            return Err(e.into());
362
0
                        }
363
0
                        running_actions_manager.complete_actions(complete_msg_clone).await;
364
0
                        Ok::<(), Error>(())
365
0
                    };
366
0
                    futures.push(shutdown_future.boxed());
367
                },
368
            };
369
        }
370
        // Unreachable.
371
1
    }
372
}
373
374
type ConnectionFactory<T> = Box<dyn Fn() -> BoxFuture<'static, Result<T, Error>> + Send + Sync>;
375
376
pub struct LocalWorker<T: WorkerApiClientTrait, U: RunningActionsManager> {
377
    config: Arc<LocalWorkerConfig>,
378
    running_actions_manager: Arc<U>,
379
    connection_factory: ConnectionFactory<T>,
380
    sleep_fn: Option<Box<dyn Fn(Duration) -> BoxFuture<'static, ()> + Send + Sync>>,
381
    metrics: Arc<Metrics>,
382
}
383
384
impl<T: WorkerApiClientTrait + core::fmt::Debug, U: RunningActionsManager + core::fmt::Debug>
385
    core::fmt::Debug for LocalWorker<T, U>
386
{
387
0
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
388
0
        f.debug_struct("LocalWorker")
389
0
            .field("config", &self.config)
390
0
            .field("running_actions_manager", &self.running_actions_manager)
391
0
            .field("metrics", &self.metrics)
392
0
            .finish_non_exhaustive()
393
0
    }
394
}
395
396
/// Creates a new `LocalWorker`. The `cas_store` must be an instance of
397
/// `FastSlowStore` and will be checked at runtime.
398
2
pub async fn new_local_worker(
399
2
    config: Arc<LocalWorkerConfig>,
400
2
    cas_store: Store,
401
2
    ac_store: Option<Store>,
402
2
    historical_store: Store,
403
2
) -> Result<LocalWorker<WorkerApiClientWrapper, RunningActionsManagerImpl>, Error> {
404
2
    let fast_slow_store = cas_store
405
2
        .downcast_ref::<FastSlowStore>(None)
406
2
        .err_tip(|| "Expected store for LocalWorker's store to be a FastSlowStore")
?0
407
2
        .get_arc()
408
2
        .err_tip(|| "FastSlowStore's Arc doesn't exist")
?0
;
409
410
2
    if let Ok(
path1
) = fs::canonicalize(&config.work_directory).await {
  Branch (410:12): [True: 1, False: 1]
  Branch (410:12): [Folded - Ignored]
  Branch (410:12): [Folded - Ignored]
411
1
        fs::remove_dir_all(path)
412
1
            .await
413
1
            .err_tip(|| "Could not remove work_directory in LocalWorker")
?0
;
414
1
    }
415
416
2
    fs::create_dir_all(&config.work_directory)
417
2
        .await
418
2
        .err_tip(|| format!(
"Could not make work_directory : {}"0
,
config.work_directory0
))
?0
;
419
2
    let entrypoint = if config.entrypoint.is_empty() {
  Branch (419:25): [True: 2, False: 0]
  Branch (419:25): [Folded - Ignored]
  Branch (419:25): [Folded - Ignored]
420
2
        None
421
    } else {
422
0
        Some(config.entrypoint.clone())
423
    };
424
2
    let max_action_timeout = if config.max_action_timeout == 0 {
  Branch (424:33): [True: 2, False: 0]
  Branch (424:33): [Folded - Ignored]
  Branch (424:33): [Folded - Ignored]
425
2
        DEFAULT_MAX_ACTION_TIMEOUT
426
    } else {
427
0
        Duration::from_secs(config.max_action_timeout as u64)
428
    };
429
2
    let running_actions_manager =
430
2
        Arc::new(RunningActionsManagerImpl::new(RunningActionsManagerArgs {
431
2
            root_action_directory: config.work_directory.clone(),
432
2
            execution_configuration: ExecutionConfiguration {
433
2
                entrypoint,
434
2
                additional_environment: config.additional_environment.clone(),
435
2
            },
436
2
            cas_store: fast_slow_store,
437
2
            ac_store,
438
2
            historical_store,
439
2
            upload_action_result_config: &config.upload_action_result,
440
2
            max_action_timeout,
441
2
            timeout_handled_externally: config.timeout_handled_externally,
442
2
        })
?0
);
443
2
    let local_worker = LocalWorker::new_with_connection_factory_and_actions_manager(
444
2
        config.clone(),
445
2
        running_actions_manager,
446
2
        Box::new(move || 
{0
447
0
            let config = config.clone();
448
0
            Box::pin(async move {
449
0
                let timeout = config
450
0
                    .worker_api_endpoint
451
0
                    .timeout
452
0
                    .unwrap_or(DEFAULT_ENDPOINT_TIMEOUT_S);
453
0
                let timeout_duration = Duration::from_secs_f32(timeout);
454
0
                let tls_config =
455
0
                    tls_utils::load_client_config(&config.worker_api_endpoint.tls_config)
456
0
                        .err_tip(|| "Parsing local worker TLS configuration")?;
457
0
                let endpoint =
458
0
                    tls_utils::endpoint_from(&config.worker_api_endpoint.uri, tls_config)
459
0
                        .map_err(|e| make_input_err!("Invalid URI for worker endpoint : {e:?}"))?
460
0
                        .connect_timeout(timeout_duration)
461
0
                        .timeout(timeout_duration);
462
463
0
                let transport = endpoint.connect().await.map_err(|e| {
464
0
                    make_err!(
465
0
                        Code::Internal,
466
                        "Could not connect to endpoint {}: {e:?}",
467
0
                        config.worker_api_endpoint.uri
468
                    )
469
0
                })?;
470
0
                Ok(WorkerApiClient::new(transport).into())
471
0
            })
472
0
        }),
473
2
        Box::new(move |d| 
Box::pin0
(
sleep0
(
d0
))),
474
    );
475
2
    Ok(local_worker)
476
2
}
477
478
impl<T: WorkerApiClientTrait, U: RunningActionsManager> LocalWorker<T, U> {
479
8
    pub fn new_with_connection_factory_and_actions_manager(
480
8
        config: Arc<LocalWorkerConfig>,
481
8
        running_actions_manager: Arc<U>,
482
8
        connection_factory: ConnectionFactory<T>,
483
8
        sleep_fn: Box<dyn Fn(Duration) -> BoxFuture<'static, ()> + Send + Sync>,
484
8
    ) -> Self {
485
8
        let metrics = Arc::new(Metrics::new(Arc::downgrade(
486
8
            running_actions_manager.metrics(),
487
        )));
488
8
        Self {
489
8
            config,
490
8
            running_actions_manager,
491
8
            connection_factory,
492
8
            sleep_fn: Some(sleep_fn),
493
8
            metrics,
494
8
        }
495
8
    }
496
497
    #[allow(
498
        clippy::missing_const_for_fn,
499
        reason = "False positive on stable, but not on nightly"
500
    )]
501
0
    pub fn name(&self) -> &String {
502
0
        &self.config.name
503
0
    }
504
505
8
    async fn register_worker(
506
8
        &self,
507
8
        client: &mut T,
508
8
    ) -> Result<(String, Streaming<UpdateForWorker>), Error> {
509
8
        let connect_worker_request =
510
8
            make_connect_worker_request(self.config.name.clone(), &self.config.platform_properties)
511
8
                .await
?0
;
512
8
        let 
mut update_for_worker_stream6
= client
513
8
            .connect_worker(connect_worker_request)
514
8
            .await
515
6
            .err_tip(|| "Could not call connect_worker() in worker")
?0
516
6
            .into_inner();
517
518
6
        let 
first_msg_update5
= update_for_worker_stream
519
6
            .next()
520
6
            .await
521
6
            .err_tip(|| "Got EOF expected UpdateForWorker")
?1
522
5
            .err_tip(|| "Got error when receiving UpdateForWorker")
?0
523
            .update;
524
525
5
        let worker_id = match first_msg_update {
526
5
            Some(Update::ConnectionResult(connection_result)) => connection_result.worker_id,
527
0
            other => {
528
0
                return Err(make_input_err!(
529
0
                    "Expected first response from scheduler to be a ConnectResult got : {:?}",
530
0
                    other
531
0
                ));
532
            }
533
        };
534
5
        Ok((worker_id, update_for_worker_stream))
535
6
    }
536
537
6
    #[instrument(skip(self), level = Level::INFO)]
538
    pub async fn run(
539
        mut self,
540
        mut shutdown_rx: broadcast::Receiver<ShutdownGuard>,
541
    ) -> Result<(), Error> {
542
        let sleep_fn = self
543
            .sleep_fn
544
            .take()
545
            .err_tip(|| "Could not unwrap sleep_fn in LocalWorker::run")?;
546
        let sleep_fn_pin = Pin::new(&sleep_fn);
547
2
        let error_handler = Box::pin(move |err| async move {
548
2
            error!(?err, "Error");
549
2
            (sleep_fn_pin)(Duration::from_secs_f32(CONNECTION_RETRY_DELAY_S)).await;
550
4
        });
551
552
        loop {
553
            // First connect to our endpoint.
554
            let mut client = match (self.connection_factory)().await {
555
                Ok(client) => client,
556
                Err(e) => {
557
                    (error_handler)(e).await;
558
                    continue; // Try to connect again.
559
                }
560
            };
561
562
            // Next register our worker with the scheduler.
563
            let (inner, update_for_worker_stream) = match self.register_worker(&mut client).await {
564
                Err(e) => {
565
                    (error_handler)(e).await;
566
                    continue; // Try to connect again.
567
                }
568
                Ok((worker_id, update_for_worker_stream)) => (
569
                    LocalWorkerImpl::new(
570
                        &self.config,
571
                        client,
572
                        worker_id,
573
                        self.running_actions_manager.clone(),
574
                        self.metrics.clone(),
575
                    ),
576
                    update_for_worker_stream,
577
                ),
578
            };
579
            info!(
580
                worker_id = %inner.worker_id,
581
                "Worker registered with scheduler"
582
            );
583
584
            // Now listen for connections and run all other services.
585
            if let Err(err) = inner.run(update_for_worker_stream, &mut shutdown_rx).await {
586
                'no_more_actions: {
587
                    // Ensure there are no actions in transit before we try to kill
588
                    // all our actions.
589
                    const ITERATIONS: usize = 1_000;
590
591
                    const ERROR_MSG: &str = "Actions in transit did not reach zero before we disconnected from the scheduler";
592
593
                    let sleep_duration = ACTIONS_IN_TRANSIT_TIMEOUT_S / ITERATIONS as f32;
594
                    for _ in 0..ITERATIONS {
595
                        if inner.actions_in_transit.load(Ordering::Acquire) == 0 {
596
                            break 'no_more_actions;
597
                        }
598
                        (sleep_fn_pin)(Duration::from_secs_f32(sleep_duration)).await;
599
                    }
600
                    error!(ERROR_MSG);
601
                    return Err(err.append(ERROR_MSG));
602
                }
603
                error!(?err, "Worker disconnected from scheduler");
604
                // Kill off any existing actions because if we re-connect, we'll
605
                // get some more and it might resource lock us.
606
                self.running_actions_manager.kill_all().await;
607
608
                (error_handler)(err).await; // Try to connect again.
609
            }
610
        }
611
        // Unreachable.
612
    }
613
}
614
615
#[derive(Debug, MetricsComponent)]
616
pub struct Metrics {
617
    #[metric(
618
        help = "Total number of actions sent to this worker to process. This does not mean it started them, it just means it received a request to execute it."
619
    )]
620
    start_actions_received: CounterWithTime,
621
    #[metric(help = "Total number of disconnects received from the scheduler.")]
622
    disconnects_received: CounterWithTime,
623
    #[metric(help = "Total number of keep-alives received from the scheduler.")]
624
    keep_alives_received: CounterWithTime,
625
    #[metric(
626
        help = "Stats about the calls to check if an action satisfies the config supplied script."
627
    )]
628
    preconditions: AsyncCounterWrapper,
629
    #[metric]
630
    #[allow(
631
        clippy::struct_field_names,
632
        reason = "TODO Fix this. Triggers on nightly"
633
    )]
634
    running_actions_manager_metrics: Weak<RunningActionManagerMetrics>,
635
}
636
637
impl RootMetricsComponent for Metrics {}
638
639
impl Metrics {
640
8
    fn new(running_actions_manager_metrics: Weak<RunningActionManagerMetrics>) -> Self {
641
8
        Self {
642
8
            start_actions_received: CounterWithTime::default(),
643
8
            disconnects_received: CounterWithTime::default(),
644
8
            keep_alives_received: CounterWithTime::default(),
645
8
            preconditions: AsyncCounterWrapper::default(),
646
8
            running_actions_manager_metrics,
647
8
        }
648
8
    }
649
}
650
651
impl Metrics {
652
4
    async fn wrap<U, T: Future<Output = U>, F: FnOnce(Arc<Self>) -> T>(
653
4
        self: Arc<Self>,
654
4
        fut: F,
655
4
    ) -> U {
656
4
        fut(self).await
657
3
    }
658
}