Coverage Report

Created: 2026-07-21 15:28

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/build/source/nativelink-worker/src/worker_utils.rs
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::hash::BuildHasher;
16
use std::collections::HashMap;
17
use std::io::{BufRead, BufReader, Cursor};
18
use std::process::Stdio;
19
20
use futures::future::try_join_all;
21
use nativelink_config::cas_server::WorkerProperty;
22
use nativelink_error::{Error, ResultExt, make_err, make_input_err};
23
use nativelink_proto::build::bazel::remote::execution::v2::platform::Property;
24
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::ConnectWorkerRequest;
25
use tokio::process;
26
use tracing::{info, warn};
27
28
#[expect(clippy::future_not_send)] // TODO(jhpratt) remove this
29
19
pub async fn make_connect_worker_request<S: BuildHasher>(
30
19
    worker_id_prefix: String,
31
19
    worker_properties: &HashMap<String, WorkerProperty, S>,
32
19
    extra_envs: &HashMap<String, String, S>,
33
19
    max_inflight_tasks: u64,
34
19
) -> Result<ConnectWorkerRequest, Error> {
35
19
    let mut futures = vec![];
36
19
    for (
property_name7
,
worker_property7
) in worker_properties {
37
7
        futures.push(async move {
38
7
            match worker_property {
39
1
                WorkerProperty::Values(values) => {
40
1
                    let mut props = Vec::with_capacity(values.len());
41
1
                    for value in values {
42
1
                        props.push(Property {
43
1
                            name: property_name.clone(),
44
1
                            value: value.clone(),
45
1
                        });
46
1
                    }
47
1
                    Ok(props)
48
                }
49
6
                WorkerProperty::QueryCmd(cmd) => {
50
6
                    let maybe_split_cmd = shlex::split(cmd);
51
6
                    let (
command5
,
args5
) = match &maybe_split_cmd {
52
5
                        Some(split_cmd) => (&split_cmd[0], &split_cmd[1..]),
53
                        None => {
54
1
                            return Err(make_input_err!(
55
1
                                "Could not parse the value of worker property: {}: '{}'",
56
1
                                property_name,
57
1
                                cmd
58
1
                            ));
59
                        }
60
                    };
61
5
                    let mut process = process::Command::new(command);
62
5
                    process.env_clear();
63
5
                    process.envs(extra_envs);
64
5
                    process.args(args);
65
5
                    process.stdin(Stdio::null());
66
5
                    let err_fn = || 
{1
67
1
                        format!("Error executing property_name {property_name} command: '{cmd}'")
68
1
                    };
69
5
                    info!(cmd, property_name, "Spawning process",);
70
5
                    let process_output = process.output().await.err_tip(err_fn)
?0
;
71
5
                    if !process_output.stderr.is_empty() {
72
2
                        warn!(
73
2
                            stderr = ?String::from_utf8_lossy(&process_output.stderr),
74
                            cmd = cmd,
75
                            property_name = property_name,
76
                            "Got stderr when running query cmd"
77
                        );
78
3
                    }
79
5
                    if !process_output.status.success() {
80
1
                        return Err(make_err!(
81
1
                            process_output.status.code().unwrap().into(),
82
1
                            "{}",
83
1
                            err_fn()
84
1
                        ));
85
4
                    }
86
4
                    let reader = BufReader::new(Cursor::new(process_output.stdout));
87
88
4
                    let mut props = vec![];
89
4
                    for 
value3
in reader.lines() {
90
3
                        props.push(Property {
91
3
                            name: property_name.clone(),
92
3
                            value: value
93
3
                                .err_tip(|| "Could split input by lines")
?0
94
3
                                .trim()
95
3
                                .to_string(),
96
                        });
97
                    }
98
4
                    Ok(props)
99
                }
100
            }
101
7
        });
102
    }
103
104
    Ok(ConnectWorkerRequest {
105
19
        worker_id_prefix,
106
19
        properties: try_join_all(futures).await
?2
.
into_iter17
().
flatten17
().
collect17
(),
107
17
        max_inflight_tasks,
108
    })
109
19
}