tango_bench/
worker.rs

1//! Child worker mode: reads JSON-RPC commands from stdin, writes samples to commpage.
2
3use crate::{
4    commpage::{Commpage, Role},
5    protocol::{self, *},
6    AuxMetricEntry, Benchmark,
7};
8use anyhow::{anyhow, Result};
9use jsonrpc_types::v2::*;
10use serde::{de::DeserializeOwned, Serialize};
11use std::{
12    io::{self, BufRead, BufReader, Write},
13    process::ExitCode,
14};
15
16/// Entry point for a child process in worker mode.
17///
18/// Opens the commpage from the given shmem name, then enters the JSON-RPC
19/// dispatch loop reading commands from stdin and writing responses to stdout.
20pub(crate) fn run_worker(
21    shmem_name: &str,
22    role: Role,
23    benchmarks: Vec<Benchmark>,
24) -> Result<ExitCode> {
25    let commpage = Commpage::open(shmem_name)?;
26
27    let aux_metrics = crate::available_aux_metrics();
28    let mut state = WorkerState {
29        benchmarks,
30        commpage,
31        role,
32        aux_metrics,
33    };
34
35    let stdin = BufReader::new(io::stdin());
36    let mut out = io::stdout().lock();
37
38    // increasing priority of a test thread, to minimize effect of CPU scheduler
39    #[cfg(target_os = "macos")]
40    unsafe {
41        use libc::qos_class_t::QOS_CLASS_USER_INTERACTIVE;
42        libc::pthread_set_qos_class_self_np(QOS_CLASS_USER_INTERACTIVE, 0);
43    }
44
45    let mut lines = stdin.lines();
46    while let Some(Ok(line)) = lines.next() {
47        let response = match serde_json::from_str::<MethodCall>(&line) {
48            Ok(r) => match r.method.as_str() {
49                protocol::METHOD_LIST_BENCHMARKS => rpc_handle(&r, |()| state.list_benchmarks()),
50                protocol::METHOD_ESTIMATE_ITERATIONS => {
51                    rpc_handle(&r, |p| state.estimate_iterations(p))
52                }
53                protocol::METHOD_RUN_BENCHMARK => rpc_handle(&r, |p| state.run_benchmark(p)),
54                protocol::METHOD_SHUTDOWN => break,
55                _ => jsonrpc_error(&r.id, Error::method_not_found()),
56            },
57            Err(e) => Output::<Value>::failure(
58                Error {
59                    code: ErrorCode::ParseError,
60                    message: format!("Parse error: {e}"),
61                    data: None,
62                },
63                None,
64            ),
65        };
66        let _ = writeln!(out, "{}", serde_json::to_string(&response).unwrap());
67        let _ = out.flush();
68    }
69    Ok(ExitCode::SUCCESS)
70}
71
72struct WorkerState {
73    commpage: Commpage,
74    role: Role,
75    benchmarks: Vec<Benchmark>,
76    aux_metrics: Vec<AuxMetricEntry>,
77}
78
79impl WorkerState {
80    fn list_benchmarks(&self) -> Result<ListBenchmarksResult> {
81        let benchmarks = self
82            .benchmarks
83            .iter()
84            .map(|b| b.name().to_string())
85            .collect();
86        let aux_metrics = self.aux_metrics.iter().map(|m| m.id.to_string()).collect();
87        Ok(ListBenchmarksResult {
88            benchmarks,
89            aux_metrics,
90        })
91    }
92
93    fn prepare_sampler(
94        &mut self,
95        index: usize,
96        seed: u64,
97    ) -> Result<Box<dyn crate::ErasedSampler>> {
98        let len = self.benchmarks.len();
99        let bench = self
100            .benchmarks
101            .get_mut(index)
102            .ok_or_else(|| anyhow!("Index {index} out of range (have {len} benchmarks)"))?;
103        Ok(bench.prepare_state(seed))
104    }
105
106    fn estimate_iterations(
107        &mut self,
108        params: EstimateIterationsParams,
109    ) -> Result<EstimateIterationsResult> {
110        let mut sampler = self.prepare_sampler(params.index, params.seed)?;
111        let iterations = sampler.estimate_iterations(params.time_ms);
112        Ok(EstimateIterationsResult { iterations })
113    }
114
115    fn run_benchmark(&mut self, params: RunBenchmarkParams) -> Result<RunBenchmarkResult> {
116        let mut sampler = self.prepare_sampler(params.index, params.seed)?;
117        let mut samples = preallocate_vec(params.num_samples);
118
119        // Resolve requested aux metrics to their function pointers
120        let active_aux: Vec<&AuxMetricEntry> = params
121            .aux_metrics
122            .iter()
123            .filter_map(|id| self.aux_metrics.iter().find(|m| m.id == id))
124            .collect();
125
126        let mut sample_no = 0u64;
127
128        #[cfg(feature = "stack-randomize")]
129        let mut stack_randomizer = stack_randomizer::StackRandomizer::new(params.seed);
130
131        if cfg!(not(feature = "no-sync")) {
132            // We use value 0 to self synchronize two processes at the start of benchmark execution.
133            // So cursor value is a sample index being collected right now. Hence the number
134            // of collected samples is cursor_value - 1;
135            assert_eq!(
136                self.commpage.load_cursor_value(self.role),
137                0,
138                "Illegal cursor sate (=0 expected)"
139            );
140            self.commpage.set_cursor_value(self.role, sample_no + 1);
141            self.commpage
142                .wait_for_cursor_value(self.role.peer(), sample_no + 1);
143        }
144
145        // Start aux metrics before the sample
146        let mut aux_metrics: Vec<u64> = active_aux.iter().map(|m| (m.start)()).collect();
147
148        // Terminate conditions differs if the explicit number of samples given.
149        // Yes – collect given amount
150        //  No – run until STOP bit is not set
151        while (params.num_samples > 0 && sample_no < params.num_samples as u64)
152            || (params.num_samples == 0 && !self.commpage.is_stopped())
153        {
154            #[cfg(feature = "stack-randomize")]
155            let elapsed_ns = stack_randomizer.measure(&mut *sampler, params.iterations);
156            #[cfg(not(feature = "stack-randomize"))]
157            let elapsed_ns = sampler.measure(params.iterations);
158
159            samples.push(elapsed_ns);
160
161            // Advance cursor and wait for peer
162            sample_no += 1;
163            if cfg!(not(feature = "no-sync")) {
164                self.commpage.set_cursor_value(self.role, sample_no + 1);
165                if !self
166                    .commpage
167                    .wait_for_cursor_value(self.role.peer(), sample_no + 1)
168                {
169                    // Peer exited early
170                    break;
171                }
172            }
173        }
174
175        // Finish aux metrics after the sample
176        for (aux, value) in active_aux.iter().zip(aux_metrics.iter_mut()) {
177            *value = (aux.finish)(*value);
178        }
179
180        self.commpage.mark_done(self.role);
181
182        Ok(RunBenchmarkResult {
183            samples,
184            aux_metrics,
185        })
186    }
187}
188
189fn preallocate_vec(samples: usize) -> Vec<u64> {
190    if samples > 0 {
191        Vec::with_capacity(samples)
192    } else {
193        Vec::new()
194    }
195}
196
197fn rpc_handle<P: DeserializeOwned, R: Serialize>(
198    req: &MethodCall,
199    f: impl FnOnce(P) -> Result<R>,
200) -> Output {
201    let params_value = params_to_value(&req.params);
202    match serde_json::from_value::<P>(params_value) {
203        Ok(p) => match f(p) {
204            Ok(r) => jsonrpc_success(&req.id, r),
205            Err(e) => server_error(&req.id, &e.to_string()),
206        },
207        Err(e) => jsonrpc_error(&req.id, Error::invalid_params(e)),
208    }
209}
210
211fn jsonrpc_success<T: Serialize>(id: &Id, result: T) -> Output {
212    let result = serde_json::to_value(result).unwrap();
213    Output::success(result, id.clone())
214}
215
216fn jsonrpc_error(id: &Id, error: Error) -> Output {
217    Output::failure(error, Some(id.clone()))
218}
219
220fn server_error(id: &Id, message: &str) -> Output {
221    jsonrpc_error(
222        id,
223        Error {
224            code: ErrorCode::ServerError(-32000),
225            message: message.to_string(),
226            data: None,
227        },
228    )
229}
230
231fn params_to_value(params: &Option<Params>) -> Value {
232    params
233        .as_ref()
234        .map(|p| serde_json::to_value(p).unwrap_or_default())
235        .unwrap_or_default()
236}
237
238#[cfg(feature = "stack-randomize")]
239mod stack_randomizer {
240    use crate::ErasedSampler;
241    use alloca::with_alloca;
242    use rand::{distributions, rngs::SmallRng, Rng, SeedableRng};
243
244    pub struct StackRandomizer {
245        rng: SmallRng,
246        distr: distributions::Uniform<usize>,
247    }
248
249    impl StackRandomizer {
250        pub fn new(seed: u64) -> Self {
251            Self {
252                rng: SmallRng::seed_from_u64(seed),
253                distr: rand::distributions::Uniform::new(0, 64),
254            }
255        }
256
257        pub fn measure(&mut self, sampler: &mut dyn ErasedSampler, iterations: usize) -> u64 {
258            with_alloca(self.rng.sample(self.distr), |_| sampler.measure(iterations))
259        }
260    }
261}
tango_bench/worker.rs

tango_bench/
worker.rs
