fluxbench_core/worker/
main_loop.rs

1//! Worker main loop.
2//!
3//! On Unix, the IPC fds come from `FLUX_IPC_FD`, and a SIGTERM handler requests a graceful stop.
4//! On non-Unix, IPC falls back to stdin/stdout and there is no signal handling.
5//!
6//! Samples stream to the supervisor during measurement.
7//! So a timeout or crash keeps the samples sent before it.
8
9use super::flusher::Flusher;
10use super::panic_report;
11use super::stream::{SharedStream, StreamSink};
12use super::transport::{self, IpcReader};
13use crate::measure::{cycles_overhead, pin_to_cpu, timer_overhead_ns};
14use crate::shutdown::{install_sigterm_handler, shutdown_requested};
15use crate::{Bencher, BenchmarkDef, MeasureConfig, run_benchmark_loop};
16use fluxbench_ipc::{
17    BenchmarkConfig, FailureKind, PIN_CPU_ENV, SupervisorCommand, WorkerCapabilities, WorkerMessage,
18};
19
20/// Pin to the CPU in `PIN_CPU_ENV`. An absent var skips pinning.
21///
22/// A bad value or a failed pin prints one stderr line. The benchmark still runs.
23fn pin_from_env() {
24    let Ok(value) = std::env::var(PIN_CPU_ENV) else {
25        return;
26    };
27    match value.parse::<usize>() {
28        Ok(cpu) => {
29            if let Err(e) = pin_to_cpu(cpu) {
30                eprintln!(
31                    "fluxbench: warning: pin to CPU {cpu} failed: {e}. Timings run unpinned."
32                );
33            }
34        }
35        Err(_) => eprintln!(
36            "fluxbench: warning: invalid {PIN_CPU_ENV}={value:?} (expected a CPU index). Timings run unpinned."
37        ),
38    }
39}
40
41/// Worker main loop
42pub struct WorkerMain {
43    reader: IpcReader,
44    stream: SharedStream,
45}
46
47impl WorkerMain {
48    /// Create a worker on the fds in `FLUX_IPC_FD`, else on stdin/stdout.
49    pub fn new() -> Self {
50        let (reader, writer) = transport::open();
51        Self {
52            reader,
53            stream: SharedStream::new(writer),
54        }
55    }
56
57    /// Run the worker main loop
58    pub fn run(&mut self) -> Result<(), Box<dyn std::error::Error>> {
59        install_sigterm_handler();
60        panic_report::install();
61
62        self.stream
63            .send(&WorkerMessage::Hello(WorkerCapabilities::default()))?;
64
65        // Start before pinning, so the flusher does not share the pinned core.
66        let _flusher = Flusher::start(self.stream.clone());
67
68        // Pin to the CPU the supervisor assigned, for stable TSC
69        pin_from_env();
70
71        // Calibrate timer overhead on the pinned core before any benchmark runs
72        let _ = timer_overhead_ns();
73        let _ = cycles_overhead();
74
75        loop {
76            if shutdown_requested() {
77                break;
78            }
79
80            let command: SupervisorCommand = self.reader.read()?;
81
82            match command {
83                SupervisorCommand::Run { bench_id, config } => {
84                    self.run_benchmark(&bench_id, &config)?;
85                    if shutdown_requested() {
86                        break;
87                    }
88                }
89                SupervisorCommand::Abort | SupervisorCommand::Shutdown => break,
90                SupervisorCommand::Ping => {}
91            }
92        }
93
94        Ok(())
95    }
96
97    /// Run one benchmark and send its samples and outcome.
98    fn run_benchmark(
99        &mut self,
100        bench_id: &str,
101        config: &BenchmarkConfig,
102    ) -> Result<(), Box<dyn std::error::Error>> {
103        let bench = inventory::iter::<BenchmarkDef>
104            .into_iter()
105            .find(|b| b.id == bench_id);
106
107        let Some(bench) = bench else {
108            self.stream.send(&WorkerMessage::Failure {
109                kind: FailureKind::Unknown,
110                message: format!(
111                    "benchmark {bench_id:?} not found in the worker binary. Rebuild the bench binary and run again."
112                ),
113                backtrace: None,
114            })?;
115            return Ok(());
116        };
117
118        self.stream.begin(bench_id);
119        panic_report::clear();
120
121        let sink = StreamSink(self.stream.clone());
122        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
123            let bencher = Bencher::with_config(MeasureConfig::from(config)).with_sample_sink(sink);
124            run_benchmark_loop(bencher, |b| (bench.runner_fn)(b))
125        }));
126
127        let last = match result {
128            Ok(Ok(bench_result)) => {
129                if bench_result.truncated {
130                    let reason = if shutdown_requested() {
131                        "stop requested"
132                    } else {
133                        "deadline reached"
134                    };
135                    eprintln!(
136                        "fluxbench: warning: {bench_id}: {reason}, {} of {} samples kept",
137                        bench_result.samples.len(),
138                        config.sample_count
139                    );
140                }
141                // Samples already went out through the sink.
142                WorkerMessage::Complete {
143                    total_iterations: bench_result.iterations,
144                    total_duration_nanos: bench_result.total_time_ns,
145                    throughput: bench_result.throughput,
146                }
147            }
148            Ok(Err(error)) => WorkerMessage::Failure {
149                kind: FailureKind::Unknown,
150                message: error.to_string(),
151                backtrace: None,
152            },
153            Err(payload) => match panic_report::take() {
154                Some(report) => WorkerMessage::Failure {
155                    kind: FailureKind::Panic,
156                    message: report.message_with_location(),
157                    backtrace: Some(report.backtrace),
158                },
159                None => WorkerMessage::Failure {
160                    kind: FailureKind::Panic,
161                    message: panic_report::payload_message(payload.as_ref()),
162                    backtrace: None,
163                },
164            },
165        };
166
167        self.stream.end(&last)?;
168        Ok(())
169    }
170}
171
172impl Default for WorkerMain {
173    fn default() -> Self {
174        Self::new()
175    }
176}
fluxbench_core/worker/main_loop.rs

fluxbench_core/worker/
main_loop.rs
