use futures::future::try_join_all;
use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
use std::path::Path;
use std::time::Duration;
use tokio::sync::RwLock;
use crate::Error;
use crate::MessageReader;
use crate::MessageWriter;
use crate::Metadata;
use crate::Result;
use crate::Topic;
use crate::message::MessageReadMut;
use crate::storage::{ConnectionType, Db3File, SplittedRosbag, find_db3_files};
use ConnectionType::{Atomic, Splitted};
type Topics = HashMap<String, Topic>;
pub struct Rosbag {
connections: ConnectionType,
metadata_path: std::path::PathBuf,
topics: RwLock<HashMap<String, Topic>>,
}
impl Rosbag {
pub async fn open<P: AsRef<std::path::Path>>(path: P) -> Result<Self> {
let path = path.as_ref();
if !path.is_dir() {
return Err(Error::BagDirNotFound(path.to_path_buf()));
}
let db3_path = find_db3_files(path)?
.into_iter()
.next()
.ok_or_else(|| Error::NoDb3File(path.to_path_buf()))?;
let db3 = Db3File::open(&db3_path).await?;
let topics = db3.get_topics().await?.into();
let conn = Atomic(db3);
let metadata_path = {
let mut path = path.to_path_buf();
path.push("metadata.yaml");
path
};
Ok(Self {
connections: conn,
metadata_path,
topics,
})
}
pub async fn open_splitted<P: AsRef<std::path::Path>, I: IntoIterator<Item = P>>(
paths: I,
metadata_path: P,
split_after: Duration,
) -> Result<Self> {
let connections = try_join_all(paths.into_iter().map(|p| Db3File::open(p))).await?;
let mut topics: Topics = HashMap::with_capacity(50);
for c in &connections {
c.get_topics()
.await?
.into_iter()
.for_each(|(topic_name, topic)| {
topics.insert(topic_name, topic);
});
}
let topics = RwLock::new(topics);
Ok(Self {
connections: Splitted(SplittedRosbag {
connections,
split_after,
}),
metadata_path: metadata_path.as_ref().to_path_buf(),
topics,
})
}
pub async fn build<P: AsRef<std::path::Path>>(path: P) -> BagBuilder<P> {
BagBuilder::new(path)
}
pub fn read(&self) -> MessageReader<'_> {
MessageReader::create(&self.connections)
}
pub fn write(&self) -> MessageWriter<'_> {
MessageWriter::create(&self.topics, &self.connections)
}
pub fn read_mut(&self) -> MessageReadMut {
todo!("Create a MessageReadMut reader for mutable access")
}
pub async fn topics(&self) -> tokio::sync::RwLockReadGuard<'_, Topics> {
self.topics.read().await
}
pub async fn get_metadata(&self) -> Result<Metadata> {
Metadata::create(
&&self.topics.read().await,
self.connections.iter(),
&self.metadata_path,
)
}
pub async fn get_index(&self) -> Result<Index> {
Ok(try_join_all(self.connections.iter().map(|c| c.get_index()))
.await?
.into_iter()
.flatten()
.collect())
}
}
#[derive(Serialize, Deserialize)]
pub struct IndexEntry {
pub topic_name: String,
pub timestamps: Vec<i64>,
}
pub type Index = Vec<IndexEntry>;
pub struct BagBuilder<P: AsRef<std::path::Path>> {
path: P,
split_after: Option<Duration>,
}
impl<P: AsRef<std::path::Path>> BagBuilder<P> {
pub fn new(path: P) -> Self {
Self {
path,
split_after: None,
}
}
pub async fn create(self) -> Result<Rosbag> {
let dir = self.path.as_ref();
let split_after = self.split_after;
let metadata_path = {
let mut path = dir.to_path_buf();
path.push("metadata.yaml");
path
};
create_dir_if_needed(dir)?;
ensure_empty(dir)?;
let db3_name = initial_db3_file_name(dir)?;
let db3_path = dir.join(&db3_name);
let connection = Db3File::init(db3_path).await?;
Ok(Rosbag {
connections: match split_after {
Some(split_after) => ConnectionType::Splitted(SplittedRosbag {
connections: vec![connection],
split_after,
}),
None => ConnectionType::Atomic(connection),
},
topics: std::collections::HashMap::new().into(),
metadata_path,
})
}
pub fn split_after(self, duration: Duration) -> Self {
Self {
path: self.path,
split_after: Some(duration),
}
}
}
fn create_dir_if_needed(dir: &Path) -> Result<()> {
std::fs::create_dir_all(dir)?;
Ok(())
}
fn ensure_empty(dir: &Path) -> Result<()> {
if dir.join("metadata.yaml").is_file() {
return Err(Error::BagHasMetadata(dir.to_path_buf()));
}
if !find_db3_files(dir)?.is_empty() {
return Err(Error::BagHasDb3Files(dir.to_path_buf()));
}
Ok(())
}
fn initial_db3_file_name(dir: &Path) -> Result<String> {
let dir_name = dir
.file_name()
.and_then(|n| n.to_str())
.ok_or_else(|| Error::NoDirectoryName(dir.to_path_buf()))?;
Ok(format!("{dir_name}_0.db3"))
}
#[cfg(test)]
mod test {
use std::env::current_dir;
use crate::Rosbag;
#[tokio::test]
async fn open_rosbag() {
let mut p = current_dir().expect("Could not get cwd");
p.push("test_bag");
Rosbag::open(p).await.expect("Could not open rosbag");
}
#[tokio::test]
async fn get_index() {
let mut p = current_dir().expect("Could not get cwd");
p.push("test_bag");
let bag = Rosbag::open(p).await.expect("Could not open rosbag");
let index = bag.get_index().await.expect("Could not get index");
assert_eq!(index.len(), 1, "Expected one index entry");
assert_eq!(index[0].topic_name, String::from("/image_raw"));
assert_eq!(index[0].timestamps.len(), 214, "Expected 214 timestamps");
assert!(
index[0].timestamps.windows(2).all(|w| w[0] <= w[1]),
"Timestamps are not sorted"
);
}
#[tokio::test]
async fn topics() {
let mut p = current_dir().expect("Could not get cwd");
p.push("test_bag");
let bag = Rosbag::open(p).await.expect("Could not open rosbag");
let topics = bag.topics().await;
assert_eq!(topics.len(), 1, "Expected one topic in rosbag");
assert_eq!(
*topics.keys().next().unwrap(),
String::from("/image_raw"),
"Wrong topic name"
);
assert_eq!(
topics.values().next().unwrap().r#type,
String::from("sensor_msgs/msg/Image"),
"Wrong topic name"
)
}
}