memtrace_utils/
pipe_io.rs

1use serde::{Deserialize, Serialize};
2use std::fs::File;
3use std::io;
4use std::io::{BufReader, BufWriter, Read, Write};
5use std::num::ParseIntError;
6use thiserror::Error;
7
8pub struct PipeReader {
9    reader: BufReader<File>,
10    buf: [u8; 1024],
11}
12
13#[derive(Debug, Error)]
14pub enum Error {
15    #[error("invalid format")]
16    InvalidFormat,
17    #[error("io error")]
18    IOError(#[from] io::Error),
19}
20
21impl From<ParseIntError> for Error {
22    fn from(_: ParseIntError) -> Self {
23        Self::InvalidFormat
24    }
25}
26
27#[derive(Debug, Serialize, Deserialize)]
28pub enum Record {
29    Version(u16),
30    Exec(String),
31    Image {
32        name: String,
33        start_address: usize,
34        size: usize,
35    },
36    PageInfo {
37        size: usize,
38        pages: usize,
39    },
40    Trace {
41        ip: usize,
42        parent_idx: usize,
43    },
44    Alloc {
45        ptr: usize,
46        size: usize,
47        parent_idx: usize,
48    },
49    Free {
50        ptr: usize,
51    },
52    Duration(u128),
53    RSS(usize),
54}
55
56impl PipeReader {
57    pub fn new(file: File) -> Self {
58        Self {
59            reader: BufReader::with_capacity(4096, file),
60            buf: [0; 1024],
61        }
62    }
63
64    pub fn read_record(&mut self) -> Option<Result<Record, Error>> {
65        let mut length_buf = [0u8; 2];
66        if let Err(_) = self.reader.read_exact(&mut length_buf) {
67            return None;
68        }
69        let len = u16::from_le_bytes(length_buf) as usize;
70
71        let mut buf = &mut self.buf[..len];
72        if let Err(e) = self.reader.read_exact(&mut buf) {
73            return Some(Err(e.into()));
74        }
75
76        let record = bincode::deserialize(&buf).map_err(|_| Error::InvalidFormat);
77
78        Some(record)
79    }
80}
81
82pub struct PipeWriter {
83    writer: BufWriter<File>,
84}
85
86impl PipeWriter {
87    pub fn new(file: File) -> Self {
88        Self {
89            writer: BufWriter::with_capacity(4096, file),
90        }
91    }
92
93    pub fn write_version(&mut self, version: u16) {
94        let record = Record::Version(version);
95        self.write_record(record)
96    }
97
98    pub fn write_image(&mut self, name: String, start_address: usize, size: usize) {
99        let record = Record::Image {
100            name,
101            start_address,
102            size,
103        };
104        self.write_record(record)
105    }
106
107    pub fn write_exec(&mut self, ex: &str) {
108        let record = Record::Exec(ex.to_string());
109        self.write_record(record)
110    }
111
112    pub fn write_page_info(&mut self, page_size: usize, phys_pages: usize) {
113        let record = Record::PageInfo {
114            size: page_size,
115            pages: phys_pages,
116        };
117        self.write_record(record)
118    }
119
120    pub fn write_trace(&mut self, ip: usize, parent_idx: usize) {
121        let record = Record::Trace { ip, parent_idx };
122        self.write_record(record)
123    }
124
125    pub fn write_alloc(&mut self, size: usize, parent_idx: usize, ptr: usize) {
126        let record = Record::Alloc {
127            ptr,
128            size,
129            parent_idx,
130        };
131        self.write_record(record)
132    }
133
134    pub fn write_free(&mut self, ptr: usize) {
135        let record = Record::Free { ptr };
136        self.write_record(record)
137    }
138
139    pub fn write_duration(&mut self, duration: u128) {
140        let record = Record::Duration(duration);
141        self.write_record(record)
142    }
143
144    pub fn write_rss(&mut self, rss: usize) {
145        let record = Record::RSS(rss);
146        self.write_record(record)
147    }
148
149    fn write_record(&mut self, record: Record) {
150        let s = bincode::serialize(&record).unwrap();
151        _ = self.writer.write_all(&(s.len() as u16).to_le_bytes());
152        _ = self.writer.write_all(&s);
153    }
154
155    pub fn flush(&mut self) {
156        _ = self.writer.flush();
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use crate::pipe_io::PipeReader;
163    use std::fs::OpenOptions;
164
165    #[test]
166    fn test_read_record() {
167        let file = OpenOptions::new().read(true).open("/tmp/trace").unwrap();
168        let mut reader = PipeReader::new(file);
169
170        let record = reader.read_record().unwrap();
171        println!("{:?}", record);
172
173        let record = reader.read_record().unwrap();
174        println!("{:?}", record);
175
176        let record = reader.read_record().unwrap();
177        println!("{:?}", record);
178    }
179}
memtrace_utils/pipe_io.rs

memtrace_utils/
pipe_io.rs
