risingwave_object_store/benchmark/
trace.rs

1// Copyright 2026 RisingWave Labs
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
15use std::fs::OpenOptions;
16use std::io::{BufWriter, Write};
17use std::path::Path;
18use std::sync::Arc;
19use std::sync::atomic::{AtomicU64, Ordering};
20use std::sync::mpsc::{self, SyncSender, TrySendError};
21use std::thread::JoinHandle;
22use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
23
24use anyhow::{Context, Result};
25use serde::Serialize;
26
27use super::ReadMode;
28
29/// Nonblocking, bounded diagnostics. It never consumes or polls the response stream.
30#[derive(Clone, Debug)]
31pub struct ReadTrace {
32    sender: SyncSender<Message>,
33    ids: Arc<AtomicU64>,
34    dropped: Arc<AtomicU64>,
35}
36
37#[derive(Debug)]
38enum Message {
39    Record(Record),
40    Finish,
41}
42
43#[derive(Debug, Serialize)]
44pub(super) struct Record {
45    pub request_id: u64,
46    pub event: &'static str,
47    pub started_unix_us: u128,
48    pub elapsed_us: u128,
49    pub offset: usize,
50    pub bytes: usize,
51    pub status: &'static str,
52    #[serde(skip_serializing_if = "Option::is_none")]
53    pub request: Option<RequestInfo>,
54}
55
56#[derive(Debug, Serialize)]
57pub(super) struct RequestInfo {
58    pub phase: String,
59    pub key: String,
60    pub reader_id: usize,
61    pub mode: ReadMode,
62    pub range_start: usize,
63    pub range_length: usize,
64}
65
66#[derive(Debug, Serialize)]
67pub struct TraceSummary {
68    pub records_written: u64,
69    pub records_dropped: u64,
70}
71
72/// Keep this guard until all benchmark readers finish, then call `finish` to check
73/// write errors and dropped records. Drop also drains the writer on early errors.
74pub struct TraceWriter {
75    trace: ReadTrace,
76    worker: Option<JoinHandle<Result<u64>>>,
77}
78
79impl TraceWriter {
80    /// Refuse to overwrite existing files. Parent directories must already exist.
81    pub fn create(path: &Path) -> Result<Self> {
82        let file = OpenOptions::new()
83            .write(true)
84            .create_new(true)
85            .open(path)
86            .context("create benchmark read trace (destination must not exist)")?;
87        let (sender, receiver) = mpsc::sync_channel(8192);
88        let worker = std::thread::Builder::new()
89            .name("bench-read-trace".into())
90            .spawn(move || {
91                let mut file = BufWriter::new(file);
92                let mut count = 0;
93                let mut flushed = Instant::now();
94                loop {
95                    match receiver.recv_timeout(Duration::from_secs(1)) {
96                        Ok(Message::Record(record)) => {
97                            serde_json::to_writer(&mut file, &record)?;
98                            file.write_all(b"\n")?;
99                            count += 1;
100                        }
101                        Ok(Message::Finish) | Err(mpsc::RecvTimeoutError::Disconnected) => break,
102                        Err(mpsc::RecvTimeoutError::Timeout) => {}
103                    }
104                    if flushed.elapsed() >= Duration::from_secs(1) {
105                        file.flush()?;
106                        flushed = Instant::now();
107                    }
108                }
109                file.flush()?;
110                Ok(count)
111            })
112            .context("start benchmark trace writer")?;
113        Ok(Self {
114            trace: ReadTrace {
115                sender,
116                ids: Arc::new(AtomicU64::new(0)),
117                dropped: Arc::new(AtomicU64::new(0)),
118            },
119            worker: Some(worker),
120        })
121    }
122
123    pub fn trace(&self) -> ReadTrace {
124        self.trace.clone()
125    }
126
127    pub fn finish(mut self) -> Result<TraceSummary> {
128        self.join()
129    }
130
131    fn join(&mut self) -> Result<TraceSummary> {
132        // Readers have finished; blocking here cannot affect their timings.
133        let _ = self.trace.sender.send(Message::Finish);
134        let records_written = self
135            .worker
136            .take()
137            .expect("writer joined once")
138            .join()
139            .map_err(|_| anyhow::anyhow!("benchmark trace writer panicked"))?
140            .context("write benchmark read trace")?;
141        Ok(TraceSummary {
142            records_written,
143            records_dropped: self.trace.dropped.load(Ordering::Relaxed),
144        })
145    }
146}
147
148impl Drop for TraceWriter {
149    fn drop(&mut self) {
150        if self.worker.is_some() {
151            let _ = self.join();
152        }
153    }
154}
155
156impl ReadTrace {
157    pub(super) fn request_id(&self) -> u64 {
158        self.ids.fetch_add(1, Ordering::Relaxed)
159    }
160
161    pub(super) fn record(&self, record: Record) {
162        if let Err(TrySendError::Full(_) | TrySendError::Disconnected(_)) =
163            self.sender.try_send(Message::Record(record))
164        {
165            self.dropped.fetch_add(1, Ordering::Relaxed);
166        }
167    }
168}
169
170pub(super) fn unix_us() -> u128 {
171    SystemTime::now()
172        .duration_since(UNIX_EPOCH)
173        .unwrap_or_default()
174        .as_micros()
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    fn record() -> Record {
182        Record {
183            request_id: 0,
184            event: "next",
185            started_unix_us: 0,
186            elapsed_us: 1,
187            offset: 0,
188            bytes: 1,
189            status: "data",
190            request: None,
191        }
192    }
193
194    #[test]
195    fn bounded_writer_counts_loss_and_preserves_existing_files() {
196        let (sender, receiver) = mpsc::sync_channel(1);
197        let trace = ReadTrace {
198            sender,
199            ids: Arc::new(AtomicU64::new(0)),
200            dropped: Arc::new(AtomicU64::new(0)),
201        };
202        trace.record(record());
203        trace.record(record());
204        assert_eq!(trace.dropped.load(Ordering::Relaxed), 1);
205        drop(receiver);
206        trace.record(record());
207        assert_eq!(trace.dropped.load(Ordering::Relaxed), 2);
208
209        let directory = tempfile::tempdir().unwrap();
210        let path = directory.path().join("trace.jsonl");
211        {
212            let writer = TraceWriter::create(&path).unwrap();
213            writer.trace().record(record());
214            assert!(TraceWriter::create(&path).is_err());
215            // Early-return Drop still drains all queued records.
216        }
217        let rows = std::fs::read_to_string(&path).unwrap();
218        assert_eq!(rows.lines().count(), 1);
219        assert!(TraceWriter::create(&path).is_err());
220        assert_eq!(std::fs::read_to_string(path).unwrap(), rows);
221    }
222}
risingwave_object_store/benchmark/trace.rs

risingwave_object_store/benchmark/
trace.rs
