Skip to main content

clankers_data/
writer.rs

1use std::collections::{BTreeMap, HashMap};
2use std::fs::File;
3use std::io::BufWriter;
4use std::path::Path;
5
6use mcap::records::MessageHeader;
7use mcap::write::Writer;
8
9use clankers_core::{RobotError, RobotResult, Timestamp};
10
11/// MCAP file writer for recording node I/O.
12pub struct McapWriter {
13    writer: Writer<BufWriter<File>>,
14    channels: HashMap<String, u16>,
15}
16
17impl McapWriter {
18    pub fn create(path: impl AsRef<Path>) -> RobotResult<Self> {
19        let file = File::create(path.as_ref())
20            .map_err(|e| RobotError::Data(format!("create mcap: {e}")))?;
21        let writer =
22            Writer::new(BufWriter::new(file)).map_err(|e| RobotError::Data(e.to_string()))?;
23        Ok(Self {
24            writer,
25            channels: HashMap::new(),
26        })
27    }
28
29    pub fn write_message(
30        &mut self,
31        topic: &str,
32        schema_name: &str,
33        encoding: &str,
34        data: &[u8],
35        log_time: Timestamp,
36    ) -> RobotResult<()> {
37        let channel_id = self.get_or_create_channel(topic, schema_name, encoding, data)?;
38
39        self.writer
40            .write_to_known_channel(
41                &MessageHeader {
42                    channel_id,
43                    sequence: 0,
44                    log_time: log_time.as_nanos(),
45                    publish_time: log_time.as_nanos(),
46                },
47                data,
48            )
49            .map_err(|e| RobotError::Data(e.to_string()))?;
50        Ok(())
51    }
52
53    pub fn finish(mut self) -> RobotResult<()> {
54        self.writer
55            .finish()
56            .map_err(|e| RobotError::Data(e.to_string()))?;
57        Ok(())
58    }
59
60    fn get_or_create_channel(
61        &mut self,
62        topic: &str,
63        schema_name: &str,
64        encoding: &str,
65        schema_data: &[u8],
66    ) -> RobotResult<u16> {
67        if let Some(&id) = self.channels.get(topic) {
68            return Ok(id);
69        }
70        let schema_id = self
71            .writer
72            .add_schema(schema_name, encoding, schema_data)
73            .map_err(|e| RobotError::Data(e.to_string()))?;
74        let channel_id = self
75            .writer
76            .add_channel(schema_id, topic, encoding, &BTreeMap::new())
77            .map_err(|e| RobotError::Data(e.to_string()))?;
78        self.channels.insert(topic.to_string(), channel_id);
79        Ok(channel_id)
80    }
81}