use crate::{Context, Mecha10Error, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::io::{Read, Write};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};
const MAGIC_BYTES: &[u8; 8] = b"METABAG1";
const VERSION: u32 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum Compression {
#[default]
None = 0,
Zstd = 1,
Gzip = 2,
}
impl From<u8> for Compression {
fn from(v: u8) -> Self {
match v {
1 => Compression::Zstd,
2 => Compression::Gzip,
_ => Compression::None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BagMetadata {
pub topics: Vec<String>,
pub message_counts: HashMap<String, u64>,
pub start_time_us: u64,
pub end_time_us: u64,
pub duration_sec: f64,
pub total_messages: u64,
pub compression: Compression,
#[serde(default)]
pub custom: HashMap<String, String>,
}
#[derive(Debug, Clone)]
struct MessageChunk {
timestamp_us: u64,
topic: String,
data: Vec<u8>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecorderConfig {
pub output_path: PathBuf,
#[serde(default)]
pub topics: Vec<String>,
#[serde(default)]
pub compression: Compression,
#[serde(default = "default_buffer_size")]
pub buffer_size: usize,
#[serde(default)]
pub max_file_size: u64,
#[serde(default)]
pub split_on_max_size: bool,
}
fn default_buffer_size() -> usize {
1000
}
impl Default for RecorderConfig {
fn default() -> Self {
Self {
output_path: PathBuf::from("recording.bag"),
topics: Vec::new(),
compression: Compression::None,
buffer_size: default_buffer_size(),
max_file_size: 0,
split_on_max_size: false,
}
}
}
pub struct BagRecorder {
config: RecorderConfig,
ctx: Arc<Context>,
buffer: Arc<RwLock<Vec<MessageChunk>>>,
metadata: Arc<RwLock<BagMetadata>>,
running: Arc<RwLock<bool>>,
write_task: Option<tokio::task::JoinHandle<()>>,
}
impl BagRecorder {
pub async fn new(ctx: Context, config: RecorderConfig) -> Result<Self> {
let metadata = BagMetadata {
topics: config.topics.clone(),
message_counts: HashMap::new(),
start_time_us: crate::prelude::now_micros(),
end_time_us: 0,
duration_sec: 0.0,
total_messages: 0,
compression: config.compression,
custom: HashMap::new(),
};
Ok(Self {
config,
ctx: Arc::new(ctx),
buffer: Arc::new(RwLock::new(Vec::new())),
metadata: Arc::new(RwLock::new(metadata)),
running: Arc::new(RwLock::new(false)),
write_task: None,
})
}
pub async fn start(&mut self) -> Result<()> {
*self.running.write().await = true;
info!("Starting bag recorder: {:?}", self.config.output_path);
self.write_header().await?;
let buffer = Arc::clone(&self.buffer);
let _metadata = Arc::clone(&self.metadata);
let running = Arc::clone(&self.running);
let output_path = self.config.output_path.clone();
let compression = self.config.compression;
let buffer_size = self.config.buffer_size;
let write_task = tokio::spawn(async move {
while *running.read().await {
tokio::time::sleep(Duration::from_millis(100)).await;
let mut buf = buffer.write().await;
if buf.len() >= buffer_size {
if let Err(e) = Self::write_chunks(&output_path, &buf, compression).await {
warn!("Failed to write chunks: {}", e);
}
buf.clear();
}
}
let buf = buffer.read().await;
if !buf.is_empty() {
if let Err(e) = Self::write_chunks(&output_path, &buf, compression).await {
warn!("Failed to write final chunks: {}", e);
}
}
});
self.write_task = Some(write_task);
for topic in &self.config.topics {
let topic_clone = topic.clone();
let buffer = Arc::clone(&self.buffer);
let metadata = Arc::clone(&self.metadata);
let running = Arc::clone(&self.running);
let ctx = Arc::clone(&self.ctx);
tokio::spawn(async move {
match ctx.subscribe_raw::<serde_json::Value>(&topic_clone).await {
Ok(mut rx) => {
while *running.read().await {
if let Some(msg) = rx.recv().await {
let mut data = Vec::new();
if ciborium::ser::into_writer(&msg, &mut data).is_ok() {
let chunk = MessageChunk {
timestamp_us: crate::prelude::now_micros(),
topic: topic_clone.clone(),
data,
};
buffer.write().await.push(chunk);
let mut meta = metadata.write().await;
*meta.message_counts.entry(topic_clone.clone()).or_insert(0) += 1;
meta.total_messages += 1;
}
}
}
}
Err(e) => {
warn!("Failed to subscribe to {}: {}", topic_clone, e);
}
}
});
}
Ok(())
}
pub async fn stop(&mut self) -> Result<()> {
*self.running.write().await = false;
info!("Stopping bag recorder");
if let Some(task) = self.write_task.take() {
task.await
.map_err(|e| Mecha10Error::Other(format!("Write task error: {}", e)))?;
}
let mut meta = self.metadata.write().await;
meta.end_time_us = crate::prelude::now_micros();
meta.duration_sec = (meta.end_time_us - meta.start_time_us) as f64 / 1_000_000.0;
self.write_header().await?;
info!(
"Recording complete: {} messages, {:.1}s",
meta.total_messages, meta.duration_sec
);
Ok(())
}
pub async fn metadata(&self) -> BagMetadata {
self.metadata.read().await.clone()
}
async fn write_header(&self) -> Result<()> {
let meta = self.metadata.read().await;
let mut file = std::fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(true)
.open(&self.config.output_path)
.map_err(Mecha10Error::Io)?;
file.write_all(MAGIC_BYTES).map_err(Mecha10Error::Io)?;
file.write_all(&VERSION.to_le_bytes()).map_err(Mecha10Error::Io)?;
file.write_all(&[self.config.compression as u8])
.map_err(Mecha10Error::Io)?;
let mut meta_bytes = Vec::new();
ciborium::ser::into_writer(&*meta, &mut meta_bytes)
.map_err(|e| Mecha10Error::Other(format!("Serialization error: {}", e)))?;
file.write_all(&(meta_bytes.len() as u64).to_le_bytes())
.map_err(Mecha10Error::Io)?;
file.write_all(&meta_bytes).map_err(Mecha10Error::Io)?;
Ok(())
}
async fn write_chunks(path: &PathBuf, chunks: &[MessageChunk], _compression: Compression) -> Result<()> {
let mut file = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(path)
.map_err(Mecha10Error::Io)?;
for chunk in chunks {
file.write_all(&chunk.timestamp_us.to_le_bytes())
.map_err(Mecha10Error::Io)?;
let topic_bytes = chunk.topic.as_bytes();
file.write_all(&(topic_bytes.len() as u32).to_le_bytes())
.map_err(Mecha10Error::Io)?;
file.write_all(topic_bytes).map_err(Mecha10Error::Io)?;
file.write_all(&(chunk.data.len() as u32).to_le_bytes())
.map_err(Mecha10Error::Io)?;
file.write_all(&chunk.data).map_err(Mecha10Error::Io)?;
}
Ok(())
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PlayerConfig {
pub bag_path: PathBuf,
#[serde(default)]
pub topics: Option<Vec<String>>,
#[serde(default = "default_rate")]
pub rate: f64,
#[serde(default)]
pub loop_playback: bool,
#[serde(default)]
pub start_paused: bool,
#[serde(default)]
pub skip_to_us: u64,
}
fn default_rate() -> f64 {
1.0
}
impl Default for PlayerConfig {
fn default() -> Self {
Self {
bag_path: PathBuf::from("recording.bag"),
topics: None,
rate: 1.0,
loop_playback: false,
start_paused: false,
skip_to_us: 0,
}
}
}
pub struct BagPlayer {
config: PlayerConfig,
ctx: Arc<Context>,
metadata: BagMetadata,
chunks: Vec<MessageChunk>,
paused: Arc<RwLock<bool>>,
playing: Arc<RwLock<bool>>,
}
impl BagPlayer {
pub async fn new(ctx: Context, config: PlayerConfig) -> Result<Self> {
info!("Loading bag file: {:?}", config.bag_path);
let (metadata, chunks) = Self::read_bag_file(&config.bag_path).await?;
info!(
"Loaded bag: {} messages, {:.1}s duration, {} topics",
metadata.total_messages,
metadata.duration_sec,
metadata.topics.len()
);
Ok(Self {
config,
ctx: Arc::new(ctx),
metadata,
chunks,
paused: Arc::new(RwLock::new(false)),
playing: Arc::new(RwLock::new(false)),
})
}
pub async fn play(&mut self) -> Result<()> {
if self.config.start_paused {
*self.paused.write().await = true;
}
*self.playing.write().await = true;
info!("Starting playback at {}x speed", self.config.rate);
loop {
let start_time = self.chunks.first().map(|c| c.timestamp_us).unwrap_or(0);
let playback_start = std::time::Instant::now();
for chunk in &self.chunks {
while *self.paused.read().await {
tokio::time::sleep(Duration::from_millis(100)).await;
}
if !*self.playing.read().await {
break;
}
if let Some(topics) = &self.config.topics {
if !topics.contains(&chunk.topic) {
continue;
}
}
let msg_time_offset = chunk.timestamp_us - start_time;
let real_time_offset = (msg_time_offset as f64 / self.config.rate) as u64;
let target_time = playback_start + Duration::from_micros(real_time_offset);
if let Some(sleep_duration) = target_time.checked_duration_since(std::time::Instant::now()) {
tokio::time::sleep(sleep_duration).await;
}
match ciborium::de::from_reader::<serde_json::Value, _>(chunk.data.as_slice()) {
Ok(msg) => {
if let Err(e) = self.ctx.publish_raw(&chunk.topic, &msg).await {
warn!("Failed to publish to {}: {}", chunk.topic, e);
} else {
debug!("Replayed message on {}", chunk.topic);
}
}
Err(e) => {
warn!("Failed to deserialize message: {}", e);
}
}
}
if !self.config.loop_playback {
break;
}
info!("Looping playback");
}
*self.playing.write().await = false;
info!("Playback complete");
Ok(())
}
async fn set_flag(flag: &RwLock<bool>, value: bool, message: &str) {
*flag.write().await = value;
info!("{}", message);
}
pub async fn pause(&self) {
Self::set_flag(&self.paused, true, "Playback paused").await;
}
pub async fn resume(&self) {
Self::set_flag(&self.paused, false, "Playback resumed").await;
}
pub async fn stop(&self) {
Self::set_flag(&self.playing, false, "Playback stopped").await;
}
pub fn metadata(&self) -> &BagMetadata {
&self.metadata
}
async fn read_bag_file(path: &PathBuf) -> Result<(BagMetadata, Vec<MessageChunk>)> {
let mut file = std::fs::File::open(path).map_err(Mecha10Error::Io)?;
let mut magic = [0u8; 8];
file.read_exact(&mut magic).map_err(Mecha10Error::Io)?;
if &magic != MAGIC_BYTES {
return Err(Mecha10Error::Other("Invalid bag file format".to_string()));
}
let mut version_bytes = [0u8; 4];
file.read_exact(&mut version_bytes).map_err(Mecha10Error::Io)?;
let version = u32::from_le_bytes(version_bytes);
if version != VERSION {
return Err(Mecha10Error::Other(format!("Unsupported bag version: {}", version)));
}
let mut compression_byte = [0u8; 1];
file.read_exact(&mut compression_byte).map_err(Mecha10Error::Io)?;
let _compression = Compression::from(compression_byte[0]);
let mut meta_len_bytes = [0u8; 8];
file.read_exact(&mut meta_len_bytes).map_err(Mecha10Error::Io)?;
let meta_len = u64::from_le_bytes(meta_len_bytes) as usize;
let mut meta_bytes = vec![0u8; meta_len];
file.read_exact(&mut meta_bytes).map_err(Mecha10Error::Io)?;
let metadata: BagMetadata = ciborium::de::from_reader(meta_bytes.as_slice())
.map_err(|e| Mecha10Error::Other(format!("Failed to parse metadata: {}", e)))?;
let mut chunks = Vec::new();
loop {
let mut timestamp_bytes = [0u8; 8];
if file.read_exact(&mut timestamp_bytes).is_err() {
break;
}
let timestamp_us = u64::from_le_bytes(timestamp_bytes);
let mut topic_len_bytes = [0u8; 4];
file.read_exact(&mut topic_len_bytes).map_err(Mecha10Error::Io)?;
let topic_len = u32::from_le_bytes(topic_len_bytes) as usize;
let mut topic_bytes = vec![0u8; topic_len];
file.read_exact(&mut topic_bytes).map_err(Mecha10Error::Io)?;
let topic = String::from_utf8(topic_bytes)
.map_err(|e| Mecha10Error::Other(format!("Invalid topic string: {}", e)))?;
let mut data_len_bytes = [0u8; 4];
file.read_exact(&mut data_len_bytes).map_err(Mecha10Error::Io)?;
let data_len = u32::from_le_bytes(data_len_bytes) as usize;
let mut data = vec![0u8; data_len];
file.read_exact(&mut data).map_err(Mecha10Error::Io)?;
chunks.push(MessageChunk {
timestamp_us,
topic,
data,
});
}
Ok((metadata, chunks))
}
}