use crate::storage::sql;
use std::collections::HashMap;
use std::path::PathBuf;
use crate::message::filter::MessageFilter;
use crate::storage::{self, ConnectionType};
use crate::{Error, ReadIterator, Result, Topic};
use crate::{Message, storage::Db3File};
use cdr::CdrBe;
use futures::future::try_join_all;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
pub struct MessageWriter<'a> {
topics: &'a RwLock<HashMap<String, Topic>>,
connections: &'a ConnectionType,
}
impl<'a> MessageWriter<'a> {
pub fn create(
topics: &'a RwLock<HashMap<String, Topic>>,
connections: &'a crate::storage::ConnectionType,
) -> MessageWriter<'a> {
Self {
topics,
connections,
}
}
pub async fn new_topic(
&self,
name: &str,
r#type: &str,
type_description_hash: &str,
) -> Result<()> {
let mut topics_lock = self.topics.write().await;
if topics_lock.contains_key(name) {
return Err(Error::TopicAlreadyExists(name.to_string()));
}
let f = |db3: &'a Db3File| async move {
let mut rows = db3
.connection
.query(
sql::SQL_STATEMENT_INSERT_TOPIC,
(name, r#type, type_description_hash),
)
.await?;
let row = rows
.next()
.await?
.ok_or(turso::Error::QueryReturnedNoRows)?;
let id = row.get::<i64>(0)?;
db3.topic_ids.write().await.insert(name.to_string(), id);
Ok::<_, Error>(id)
};
let ids = try_join_all(self.connections.iter().map(f)).await?;
let id = *ids
.first()
.ok_or(Error::NoDb3File(std::path::PathBuf::new()))?;
if ids.iter().any(|other| *other != id) {
return Err(Error::TopicIdMismatch {
topic: name.to_string(),
ids,
});
}
topics_lock.insert(
name.to_string(),
Topic {
id,
name: name.to_string(),
r#type: r#type.to_string(),
serialization_format: "cdr".to_string(),
offered_qos_profiles: sql::OFFERED_QOS_PROFILES.to_string(),
},
);
Ok(())
}
pub async fn rm_topics<'b>(
&self,
topics: impl IntoIterator<Item = &'b str>,
) -> Result<HashMap<&'b str, Vec<&'a PathBuf>>> {
let mut topics_lock = self.topics.write().await;
let f = |(topic, db3): (&'b str, &'a Db3File)| async move {
let topic_id = db3
.topic_ids
.write()
.await
.remove(topic)
.ok_or(Error::TopicNotFound(topic.to_string()))?;
db3.connection
.query(storage::sql::SQL_STATEMENT_REMOVE_TOPIC, [topic])
.await
.map_err(|_| Error::TopicNotFound(topic.to_string()))?;
db3.connection
.query(
storage::sql::SQL_STATEMENT_REMOVE_TOPIC_MESSAGES,
[topic_id],
)
.await
.map_err(|_| Error::TopicIdNotFound(topic_id))?;
Ok::<_, Error>((topic, &db3.path))
};
let results = {
let mut map: HashMap<&'b str, Vec<&'a PathBuf>> = HashMap::new();
for (topic, db3) in try_join_all(topics.into_iter().flat_map(|topic| {
self.connections
.iter()
.map(move |db3| (topic, db3))
.map(f)
}))
.await?
{
map.entry(topic).or_default().push(db3);
}
map
};
for (topic, _) in &results {
topics_lock.remove(*topic);
}
Ok(results)
}
pub async fn insert<T: Serialize + for<'c> Deserialize<'c>>(
&self,
topic: &str,
timestamp: i64,
msg: T,
) -> Result<()> {
let data = cdr::serialize::<_, _, CdrBe>(&msg, cdr::Infinite)?;
let db3 = match &self.connections {
ConnectionType::Atomic(db3_file) => db3_file,
ConnectionType::Splitted(splitted_rosbag) => {
Self::connection_for_timestamp(&splitted_rosbag.connections, timestamp).await?
}
};
let topic_id = db3.topic_id(topic).await?;
db3.insert_message(topic_id, timestamp, data).await?;
Ok(())
}
async fn connection_for_timestamp(connections: &[Db3File], timestamp: i64) -> Result<&Db3File> {
let mut newest: Option<(i64, usize)> = None;
for (idx, db3_file) in connections.iter().enumerate() {
let (min, max) = db3_file.time_range().await?;
if min.is_some_and(|min| timestamp >= min) && max.is_some_and(|max| timestamp <= max) {
return Ok(db3_file);
}
if let Some(max) = max {
if newest.is_none_or(|(newest_max, _)| max > newest_max) {
newest = Some((max, idx));
}
}
}
let (_, idx) = newest.ok_or(Error::NoDb3File(std::path::PathBuf::new()))?;
Ok(&connections[idx])
}
pub async fn pop_by_filter<I: ReadIterator>(
&self,
message_filter: &MessageFilter<'_>,
) -> Result<I> {
let f = async |db3: &Db3File| db3.query_messages(&message_filter, true).await;
let rows = self.connections.run_on_connections(f).await?;
Ok(I::from_rows(rows))
}
pub async fn pop<M: Message>(&self) -> Result<Option<M>> {
for conn in self.connections.iter().rev() {
let mut ret = conn
.connection
.query(sql::SQL_STATEMENT_POP_MESSAGE, ())
.await?;
if let Some(row) = ret.next().await? {
return M::from_row(row).map(|m| Some(m));
}
}
Ok(None)
}
pub async fn pop_front<M: Message>(&self) -> Result<Option<M>> {
for conn in self.connections.iter() {
let mut ret = conn
.connection
.query(sql::SQL_STATEMENT_POP_FRONT_MESSAGE, ())
.await?;
if let Some(row) = ret.next().await? {
return M::from_row(row).map(|m| Some(m));
}
}
Ok(None)
}
}
#[cfg(test)]
mod test {
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use futures::StreamExt;
use super::*;
use crate::RawMessage;
use crate::RawReadIterator;
use crate::Rosbag;
fn unique_temp_dir() -> PathBuf {
let unique = format!(
"rs-bag-{}-{}",
std::process::id(),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("Time went backwards")
.as_nanos()
);
let dir = std::env::temp_dir().join(unique);
std::fs::create_dir_all(&dir).expect("Could not create temp dir");
dir
}
fn fixture_db3s() -> Vec<PathBuf> {
let mut src = std::env::current_dir().expect("Could not get cwd");
src.push("test_bag");
storage::find_db3_files(&src).expect("Could not find db3 files")
}
fn copy_fixture_db3s(dir: &PathBuf) {
for db3 in fixture_db3s() {
let dst = dir.join(db3.file_name().expect("Could not get file name"));
std::fs::copy(&db3, dst).expect("Could not copy db3 file");
}
}
async fn count_messages(mut it: RawReadIterator) -> usize {
let mut count = 0;
while let Some(result) = it.next().await {
result.expect("Reading next message failed");
count += 1;
}
count
}
async fn db3_message_count(path: &PathBuf) -> usize {
let db3 = Db3File::open(path).await.expect("Could not open db3 file");
db3.get_index()
.await
.expect("Could not get index")
.iter()
.map(|e| e.timestamps.len())
.sum()
}
#[tokio::test]
async fn test_pop_by_filter_drains_matching_messages() {
let bag_dir = unique_temp_dir();
copy_fixture_db3s(&bag_dir);
let bag = Rosbag::open(&bag_dir)
.await
.expect("Could not open rosbag copy");
let index = bag.get_index().await.expect("Could not get index");
let entry = index
.iter()
.find(|e| e.topic_name == "/image_raw")
.expect("No /image_raw topic in index");
let cutoff = entry.timestamps[entry.timestamps.len() / 2];
let expected_popped = entry.timestamps.iter().filter(|t| **t < cutoff).count();
let total = entry.timestamps.len();
let filter = MessageFilter::new()
.between(..cutoff)
.topics(["/image_raw"])
.expect("Could not construct MessageFilter");
let mut popped: RawReadIterator = bag
.write()
.pop_by_filter(&filter)
.await
.expect("pop_by_filter failed");
let mut popped_count = 0;
while let Some(result) = popped.next().await {
let msg = result.expect("Popping next message failed");
assert!(msg.timestamp < cutoff);
popped_count += 1;
}
assert_eq!(popped_count, expected_popped);
let reader = bag.read();
let drained: RawReadIterator = reader.read(filter).await.expect("Reading failed");
assert_eq!(count_messages(drained).await, 0);
let remaining: RawReadIterator = reader.read_all().await.expect("Reading failed");
assert_eq!(count_messages(remaining).await, total - expected_popped);
let _ = std::fs::remove_dir_all(&bag_dir);
}
#[tokio::test]
async fn test_pop_by_filter_none_drains_all() {
let bag_dir = unique_temp_dir();
copy_fixture_db3s(&bag_dir);
let bag = Rosbag::open(&bag_dir)
.await
.expect("Could not open rosbag copy");
let popped: RawReadIterator = bag
.write()
.pop_by_filter(&MessageFilter::NONE)
.await
.expect("pop_by_filter failed");
assert_eq!(count_messages(popped).await, 214);
let remaining: RawReadIterator = bag.read().read_all().await.expect("Reading failed");
assert_eq!(count_messages(remaining).await, 0);
let _ = std::fs::remove_dir_all(&bag_dir);
}
#[tokio::test]
async fn test_pop_removes_last_message_from_last_connection() {
let tmp = unique_temp_dir();
let path_a = tmp.join("bag_a_0.db3");
let path_b = tmp.join("bag_b_0.db3");
let fixture = fixture_db3s()[0].clone();
std::fs::copy(&fixture, &path_a).expect("Could not copy db3 file");
std::fs::copy(&fixture, &path_b).expect("Could not copy db3 file");
let bag = Rosbag::open_splitted(
[&path_a, &path_b],
&tmp.join("metadata.yaml"),
Duration::from_secs(100),
)
.await
.expect("Could not open splitted rosbag");
let popped: Option<RawMessage> = bag.write().pop().await.expect("pop failed");
let popped = popped.expect("Expected at least one message");
let max_timestamp = bag
.get_index()
.await
.expect("Could not get index")
.iter()
.find(|e| e.topic_name == "/image_raw")
.expect("No /image_raw topic in index")
.timestamps
.iter()
.copied()
.max()
.expect("Expected at least one timestamp");
assert_eq!(popped.timestamp, max_timestamp);
drop(bag);
let count_a = db3_message_count(&path_a).await;
let count_b = db3_message_count(&path_b).await;
assert_eq!(
count_a, 214,
"First connection should not be touched by pop"
);
assert_eq!(
count_b, 213,
"pop should remove the last message from the last connection"
);
let _ = std::fs::remove_dir_all(&tmp);
}
#[derive(Serialize, Deserialize)]
struct Probe {
value: u32,
}
async fn empty_db3_file(dir: &PathBuf) -> PathBuf {
let bag = Rosbag::build(dir)
.await
.create()
.await
.expect("Could not create empty bag");
drop(bag);
storage::find_db3_files(dir)
.expect("Could not find db3 files")
.first()
.expect("Expected one db3 file")
.clone()
}
async fn bag_db3_file_with_topic(dir: &PathBuf, topic: &str) -> PathBuf {
let bag = Rosbag::build(dir)
.await
.create()
.await
.expect("Could not create bag");
bag.write()
.new_topic(topic, "std_msgs/msg/UInt8", "hash")
.await
.expect("Could not create existing topic");
drop(bag);
storage::find_db3_files(dir)
.expect("Could not find db3 files")
.first()
.expect("Expected one db3 file")
.clone()
}
async fn topic_timestamps(bag: &Rosbag) -> Vec<i64> {
bag.get_index()
.await
.expect("Could not get index")
.iter()
.find(|e| e.topic_name == "/image_raw")
.expect("No /image_raw topic in index")
.timestamps
.clone()
}
#[tokio::test]
async fn test_pop_front_removes_first_message_in_read_order() {
let bag_dir = unique_temp_dir();
copy_fixture_db3s(&bag_dir);
let bag = Rosbag::open(&bag_dir)
.await
.expect("Could not open rosbag copy");
let writer = bag.write();
writer
.insert("/image_raw", 0, Probe { value: 1 })
.await
.expect("Could not insert probe");
writer
.insert("/image_raw", 0, Probe { value: 2 })
.await
.expect("Could not insert probe");
let mut probes: RawReadIterator = bag
.read()
.read(MessageFilter::new().between(0..1))
.await
.expect("Reading failed");
let first = probes
.next()
.await
.expect("Expected two probes")
.expect("Reading first probe failed");
let second = probes
.next()
.await
.expect("Expected two probes")
.expect("Reading second probe failed");
assert_eq!(first.timestamp, 0);
assert_eq!(second.timestamp, 0);
assert!(
first.id < second.id,
"Read order should yield the lower id first"
);
let popped: Option<RawMessage> = writer.pop_front().await.expect("pop_front failed");
let popped = popped.expect("Expected a message");
assert_eq!(
popped.id, first.id,
"pop_front must remove the message that read order yields first"
);
let mut remaining: RawReadIterator = bag
.read()
.read(MessageFilter::new().between(0..1))
.await
.expect("Reading failed");
let new_first = remaining
.next()
.await
.expect("Expected one remaining probe")
.expect("Reading remaining probe failed");
assert_eq!(new_first.id, second.id);
assert!(remaining.next().await.is_none());
let _ = std::fs::remove_dir_all(&bag_dir);
}
#[tokio::test]
async fn test_pop_front_on_empty_bag_returns_none() {
let bag_dir = unique_temp_dir();
let bag = Rosbag::build(&bag_dir)
.await
.create()
.await
.expect("Could not create empty bag");
let popped: Option<RawMessage> = bag.write().pop_front().await.expect("pop_front failed");
assert!(popped.is_none());
let popped: Option<RawMessage> = bag.write().pop().await.expect("pop failed");
assert!(popped.is_none());
let _ = std::fs::remove_dir_all(&bag_dir);
}
#[tokio::test]
async fn test_pop_front_skips_empty_leading_connection() {
let tmp = unique_temp_dir();
let empty_path = empty_db3_file(&tmp.join("empty")).await;
let fixture_path = tmp.join("bag_b_0.db3");
let fixture = fixture_db3s()[0].clone();
std::fs::copy(&fixture, &fixture_path).expect("Could not copy db3 file");
let bag = Rosbag::open_splitted(
[&empty_path, &fixture_path],
&tmp.join("metadata.yaml"),
Duration::from_secs(100),
)
.await
.expect("Could not open splitted rosbag");
let min_timestamp = topic_timestamps(&bag)
.await
.into_iter()
.min()
.expect("Expected at least one timestamp");
let popped: Option<RawMessage> = bag.write().pop_front().await.expect("pop_front failed");
let popped = popped.expect("Expected a message from the non-empty connection");
assert_eq!(popped.timestamp, min_timestamp);
drop(bag);
let count_empty = db3_message_count(&empty_path).await;
let count_fixture = db3_message_count(&fixture_path).await;
assert_eq!(count_empty, 0, "Empty connection should stay empty");
assert_eq!(
count_fixture, 213,
"pop_front should remove the first message from the first non-empty connection"
);
let _ = std::fs::remove_dir_all(&tmp);
}
#[tokio::test]
async fn test_pop_skips_empty_trailing_connection() {
let tmp = unique_temp_dir();
let fixture_path = tmp.join("bag_a_0.db3");
let empty_path = empty_db3_file(&tmp.join("empty")).await;
let fixture = fixture_db3s()[0].clone();
std::fs::copy(&fixture, &fixture_path).expect("Could not copy db3 file");
let bag = Rosbag::open_splitted(
[&fixture_path, &empty_path],
&tmp.join("metadata.yaml"),
Duration::from_secs(100),
)
.await
.expect("Could not open splitted rosbag");
let max_timestamp = topic_timestamps(&bag)
.await
.into_iter()
.max()
.expect("Expected at least one timestamp");
let popped: Option<RawMessage> = bag.write().pop().await.expect("pop failed");
let popped = popped.expect("Expected a message from the non-empty connection");
assert_eq!(popped.timestamp, max_timestamp);
drop(bag);
let count_fixture = db3_message_count(&fixture_path).await;
let count_empty = db3_message_count(&empty_path).await;
assert_eq!(
count_fixture, 213,
"pop should remove the last message from the last non-empty connection"
);
assert_eq!(count_empty, 0, "Empty connection should stay empty");
let _ = std::fs::remove_dir_all(&tmp);
}
#[tokio::test]
async fn test_new_topic_assigns_id_and_enables_insert() {
let bag_dir = unique_temp_dir();
let bag = Rosbag::build(&bag_dir)
.await
.create()
.await
.expect("Could not create empty bag");
let writer = bag.write();
writer
.new_topic("/probes", "std_msgs/msg/UInt8", "hash1")
.await
.expect("new_topic failed");
{
let topics = bag.topics().await;
let topic = topics.get("/probes").expect("Topic missing in memory");
assert_eq!(topic.id, 1, "First topic in a fresh bag must get id 1");
}
writer
.insert("/probes", 123, Probe { value: 7 })
.await
.expect("insert failed");
let mut messages: RawReadIterator = bag.read().read_all().await.expect("Reading failed");
let msg = messages
.next()
.await
.expect("Expected one message")
.expect("Reading message failed");
assert_eq!(msg.timestamp, 123);
assert!(messages.next().await.is_none());
writer
.new_topic("/probes2", "std_msgs/msg/UInt8", "hash2")
.await
.expect("new_topic failed");
{
let topics = bag.topics().await;
let topic = topics.get("/probes2").expect("Topic missing in memory");
assert_eq!(topic.id, 2, "Second topic must get id 2");
}
let err = writer
.new_topic("/probes", "std_msgs/msg/UInt8", "hash1")
.await
.expect_err("Expected duplicate topic to be rejected");
assert!(matches!(err, Error::TopicAlreadyExists(_)));
let _ = std::fs::remove_dir_all(&bag_dir);
}
#[tokio::test]
async fn test_new_topic_gets_next_id_after_existing_topic() {
let bag_dir = unique_temp_dir();
let bag = Rosbag::build(&bag_dir)
.await
.create()
.await
.expect("Could not create empty bag");
let writer = bag.write();
writer
.new_topic("/existing", "std_msgs/msg/UInt8", "hash0")
.await
.expect("Could not create existing topic");
let existing_id = {
let topics = bag.topics().await;
topics.get("/existing").expect("Topic missing").id
};
writer
.new_topic("/probes", "std_msgs/msg/UInt8", "hash")
.await
.expect("new_topic failed");
{
let topics = bag.topics().await;
let topic = topics.get("/probes").expect("Topic missing in memory");
assert_eq!(
topic.id,
existing_id + 1,
"New topic must get the next free id, not the existing maximum"
);
}
writer
.insert("/probes", 42, Probe { value: 1 })
.await
.expect("insert failed");
let probes: RawReadIterator = bag
.read()
.read(
MessageFilter::new()
.topics(["/probes"])
.expect("Could not construct MessageFilter"),
)
.await
.expect("Reading failed");
assert_eq!(count_messages(probes).await, 1);
let _ = std::fs::remove_dir_all(&bag_dir);
}
#[tokio::test]
async fn test_new_topic_on_splitted_bag_inserts_consistently() {
let tmp = unique_temp_dir();
let path_a = bag_db3_file_with_topic(&tmp.join("a"), "/existing").await;
let path_b = bag_db3_file_with_topic(&tmp.join("b"), "/existing").await;
let bag = Rosbag::open_splitted(
[&path_a, &path_b],
&tmp.join("metadata.yaml"),
Duration::from_secs(100),
)
.await
.expect("Could not open splitted rosbag");
bag.write()
.new_topic("/probes", "std_msgs/msg/UInt8", "hash")
.await
.expect("new_topic failed");
{
let topics = bag.topics().await;
let topic = topics.get("/probes").expect("Topic missing in memory");
assert_eq!(topic.id, 2);
}
drop(bag);
for path in [&path_a, &path_b] {
let db3 = Db3File::open(path).await.expect("Could not open db3 file");
let topics = db3.get_topics().await.expect("Could not get topics");
let topic = topics.get("/probes").expect("Topic missing in db3 file");
assert_eq!(topic.id, 2);
}
let _ = std::fs::remove_dir_all(&tmp);
}
#[tokio::test]
async fn test_new_topic_id_mismatch_errors() {
let tmp = unique_temp_dir();
let path_a = bag_db3_file_with_topic(&tmp.join("a"), "/existing").await;
let path_b = empty_db3_file(&tmp.join("b")).await;
let bag = Rosbag::open_splitted(
[&path_a, &path_b],
&tmp.join("metadata.yaml"),
Duration::from_secs(100),
)
.await
.expect("Could not open splitted rosbag");
let err = bag
.write()
.new_topic("/probes", "std_msgs/msg/UInt8", "hash")
.await
.expect_err("Expected diverging topic ids to be rejected");
assert!(matches!(err, Error::TopicIdMismatch { .. }));
let _ = std::fs::remove_dir_all(&tmp);
}
}