Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Save recordings from web viewer #5488

Merged
merged 6 commits into from
Mar 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions crates/re_data_store/src/store_read.rs
Original file line number Diff line number Diff line change
Expand Up @@ -397,6 +397,7 @@ impl DataStore {

/// Sort all unsorted indices in the store.
pub fn sort_indices_if_needed(&self) {
re_tracing::profile_function!();
for index in self.tables.values() {
index.sort_indices_if_needed();
}
Expand Down
43 changes: 41 additions & 2 deletions crates/re_entity_db/src/entity_db.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,9 @@ use re_data_store::{
DataStore, DataStoreConfig, GarbageCollectionOptions, StoreEvent, StoreSubscriber,
};
use re_log_types::{
ApplicationId, DataCell, DataRow, DataTable, EntityPath, EntityPathHash, LogMsg, RowId,
SetStoreInfo, StoreId, StoreInfo, StoreKind, TimePoint, Timeline,
ApplicationId, DataCell, DataRow, DataTable, DataTableResult, EntityPath, EntityPathHash,
LogMsg, RowId, SetStoreInfo, StoreId, StoreInfo, StoreKind, TimePoint, TimeRange, TimeRangeF,
Timeline,
};
use re_types_core::{components::InstanceKey, Archetype, Loggable};

Expand Down Expand Up @@ -510,6 +511,44 @@ impl EntityDb {
self.store_info()
.map(|info| (info.application_id.0.as_str(), info.started))
}

/// Export the contents of the current database to a sequence of messages.
///
/// If `time_selection` is specified, then only data for that specific timeline over that
/// specific time range will be accounted for.
pub fn to_messages(
&self,
time_selection: Option<(Timeline, TimeRangeF)>,
) -> DataTableResult<Vec<LogMsg>> {
re_tracing::profile_function!();

self.store().sort_indices_if_needed();

let set_store_info_msg = self
.store_info_msg()
.map(|msg| LogMsg::SetStoreInfo(msg.clone()));

let time_filter = time_selection.map(|(timeline, range)| {
(
timeline,
TimeRange::new(range.min.floor(), range.max.ceil()),
)
});

let data_messages = self.store().to_data_tables(time_filter).map(|table| {
table
.to_arrow_msg()
.map(|msg| LogMsg::ArrowMsg(self.store_id().clone(), msg))
});

let messages: Result<Vec<_>, _> = set_store_info_msg
.map(re_log_types::DataTableResult::Ok)
.into_iter()
.chain(data_messages)
.collect();

messages
}
}

// ----------------------------------------------------------------------------
Expand Down
18 changes: 11 additions & 7 deletions crates/re_log_encoding/src/encoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,8 @@ impl<W: std::io::Write> Encoder<W> {
}

pub fn append(&mut self, message: &LogMsg) -> Result<(), EncodeError> {
re_tracing::profile_function!();

self.uncompressed.clear();
rmp_serde::encode::write_named(&mut self.uncompressed, message)?;

Expand Down Expand Up @@ -119,21 +121,23 @@ pub fn encode<'a>(
messages: impl Iterator<Item = &'a LogMsg>,
write: &mut impl std::io::Write,
) -> Result<(), EncodeError> {
re_tracing::profile_function!();
let mut encoder = Encoder::new(options, write)?;
for message in messages {
encoder.append(message)?;
}
Ok(())
}

pub fn encode_owned(
pub fn encode_as_bytes<'a>(
options: EncodingOptions,
messages: impl Iterator<Item = LogMsg>,
write: impl std::io::Write,
) -> Result<(), EncodeError> {
let mut encoder = Encoder::new(options, write)?;
messages: impl Iterator<Item = &'a LogMsg>,
) -> Result<Vec<u8>, EncodeError> {
re_tracing::profile_function!();
let mut bytes: Vec<u8> = vec![];
let mut encoder = Encoder::new(options, &mut bytes)?;
for message in messages {
encoder.append(&message)?;
encoder.append(message)?;
}
Ok(())
Ok(bytes)
}
3 changes: 0 additions & 3 deletions crates/re_log_encoding/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
#[cfg(feature = "decoder")]
pub mod decoder;
#[cfg(feature = "encoder")]
#[cfg(not(target_arch = "wasm32"))] // we do no yet support encoding LogMsgs in the browser
pub mod encoder;

#[cfg(feature = "encoder")]
Expand Down Expand Up @@ -121,7 +120,6 @@ impl FileHeader {
pub const SIZE: usize = 12;

#[cfg(feature = "encoder")]
#[cfg(not(target_arch = "wasm32"))] // we do no yet support encoding LogMsgs in the browser
pub fn encode(&self, write: &mut impl std::io::Write) -> Result<(), encoder::EncodeError> {
write
.write_all(&self.magic)
Expand Down Expand Up @@ -165,7 +163,6 @@ impl MessageHeader {
pub const SIZE: usize = 8;

#[cfg(feature = "encoder")]
#[cfg(not(target_arch = "wasm32"))] // we do no yet support encoding LogMsgs in the browser
pub fn encode(&self, write: &mut impl std::io::Write) -> Result<(), encoder::EncodeError> {
write
.write_all(&self.compressed_len.to_le_bytes())
Expand Down
4 changes: 2 additions & 2 deletions crates/re_ui/examples/re_ui_example.rs
Original file line number Diff line number Diff line change
Expand Up @@ -541,8 +541,8 @@ impl ExampleApp {
}

fn file_menu(ui: &mut egui::Ui, command_sender: &CommandSender) {
UICommand::Save.menu_button_ui(ui, command_sender);
UICommand::SaveSelection.menu_button_ui(ui, command_sender);
UICommand::SaveRecording.menu_button_ui(ui, command_sender);
UICommand::SaveRecordingSelection.menu_button_ui(ui, command_sender);
UICommand::Open.menu_button_ui(ui, command_sender);
UICommand::Quit.menu_button_ui(ui, command_sender);
}
Expand Down
21 changes: 7 additions & 14 deletions crates/re_ui/src/command.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,8 @@ pub trait UICommandSender {
pub enum UICommand {
// Listed in the order they show up in the command palette by default!
Open,
#[cfg(not(target_arch = "wasm32"))]
Save,
#[cfg(not(target_arch = "wasm32"))]
SaveSelection,
SaveRecording,
SaveRecordingSelection,
CloseCurrentRecording,
#[cfg(not(target_arch = "wasm32"))]
Quit,
Expand Down Expand Up @@ -95,12 +93,10 @@ impl UICommand {

pub fn text_and_tooltip(self) -> (&'static str, &'static str) {
match self {
#[cfg(not(target_arch = "wasm32"))]
Self::Save => ("Save…", "Save all data to a Rerun data file (.rrd)"),
Self::SaveRecording => ("Save recording…", "Save all data to a Rerun data file (.rrd)"),

#[cfg(not(target_arch = "wasm32"))]
Self::SaveSelection => (
"Save loop selection…",
Self::SaveRecordingSelection => (
"Save recording (current time selection only)…",
"Save data for the current loop selection to a Rerun data file (.rrd)",
),

Expand Down Expand Up @@ -238,7 +234,6 @@ impl UICommand {
KeyboardShortcut::new(Modifiers::COMMAND, key)
}

#[cfg(not(target_arch = "wasm32"))]
fn cmd_alt(key: Key) -> KeyboardShortcut {
KeyboardShortcut::new(Modifiers::COMMAND.plus(Modifiers::ALT), key)
}
Expand All @@ -248,10 +243,8 @@ impl UICommand {
}

match self {
#[cfg(not(target_arch = "wasm32"))]
Self::Save => Some(cmd(Key::S)),
#[cfg(not(target_arch = "wasm32"))]
Self::SaveSelection => Some(cmd_alt(Key::S)),
Self::SaveRecording => Some(cmd(Key::S)),
Self::SaveRecordingSelection => Some(cmd_alt(Key::S)),
Self::Open => Some(cmd(Key::O)),
Self::CloseCurrentRecording => None,

Expand Down
86 changes: 70 additions & 16 deletions crates/re_viewer/src/app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -442,18 +442,17 @@ impl App {
cmd: UICommand,
) {
match cmd {
#[cfg(not(target_arch = "wasm32"))]
UICommand::Save => {
UICommand::SaveRecording => {
save(self, store_context, None);
}
#[cfg(not(target_arch = "wasm32"))]
UICommand::SaveSelection => {
UICommand::SaveRecordingSelection => {
save(
self,
store_context,
self.state.loop_selection(store_context),
);
}

#[cfg(not(target_arch = "wasm32"))]
UICommand::Open => {
for file_path in open_file_dialog_native() {
Expand Down Expand Up @@ -1501,41 +1500,96 @@ async fn async_open_rrd_dialog() -> Vec<re_data_source::FileContents> {
file_contents
}

#[cfg(not(target_arch = "wasm32"))]
#[allow(clippy::needless_pass_by_ref_mut)]
fn save(
app: &mut App,
#[allow(unused_variables)] app: &mut App, // only used on native
store_context: Option<&StoreContext<'_>>,
loop_selection: Option<(re_entity_db::Timeline, re_log_types::TimeRangeF)>,
) {
use crate::saving::save_database_to_file;
re_tracing::profile_function!();

let Some(entity_db) = store_context.as_ref().and_then(|view| view.recording) else {
// NOTE: Can only happen if saving through the command palette.
re_log::error!("No data to save!");
return;
};

let file_name = "data.rrd";

let title = if loop_selection.is_some() {
"Save loop selection"
} else {
"Save"
};

if let Some(path) = rfd::FileDialog::new()
.set_file_name("data.rrd")
.set_title(title)
.save_file()
// Web
#[cfg(target_arch = "wasm32")]
{
let f = match save_database_to_file(entity_db, path, loop_selection) {
Ok(f) => f,
let messages = match entity_db.to_messages(loop_selection) {
Ok(messages) => messages,
Err(err) => {
re_log::error!("File saving failed: {err}");
return;
}
};
if let Err(err) = app.background_tasks.spawn_file_saver(f) {
// NOTE: Can only happen if saving through the command palette.
re_log::error!("File saving failed: {err}");

wasm_bindgen_futures::spawn_local(async move {
if let Err(err) = async_save_dialog(file_name, title, &messages).await {
re_log::error!("File saving failed: {err}");
}
});
}

// Native
#[cfg(not(target_arch = "wasm32"))]
{
let path = {
re_tracing::profile_scope!("file_dialog");
rfd::FileDialog::new()
.set_file_name(file_name)
.set_title(title)
.save_file()
};
if let Some(path) = path {
let messages = match entity_db.to_messages(loop_selection) {
Ok(messages) => messages,
Err(err) => {
re_log::error!("File saving failed: {err}");
return;
}
};
if let Err(err) = app
.background_tasks
.spawn_file_saver(move || crate::saving::encode_to_file(&path, messages.iter()))
{
// NOTE: Can only happen if saving through the command palette.
re_log::error!("File saving failed: {err}");
}
}
}
}

#[cfg(target_arch = "wasm32")]
async fn async_save_dialog(
file_name: &str,
title: &str,
messages: &[LogMsg],
) -> anyhow::Result<()> {
use anyhow::Context as _;

let file_handle = rfd::AsyncFileDialog::new()
.set_file_name(file_name)
.set_title(title)
.save_file()
.await;

let Some(file_handle) = file_handle else {
return Ok(()); // aborted
};

let bytes = re_log_encoding::encoder::encode_as_bytes(
re_log_encoding::EncodingOptions::COMPRESSED,
messages.iter(),
)?;
file_handle.write(&bytes).await.context("Failed to save")
}
64 changes: 9 additions & 55 deletions crates/re_viewer/src/saving.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,3 @@
#[cfg(not(target_arch = "wasm32"))]
use re_entity_db::EntityDb;

#[cfg(not(target_arch = "wasm32"))]
use re_log_types::ApplicationId;

Expand Down Expand Up @@ -65,60 +62,17 @@ pub fn default_blueprint_path(app_id: &ApplicationId) -> anyhow::Result<std::pat
}

#[cfg(not(target_arch = "wasm32"))]
/// Returns a closure that, when run, will save the contents of the current database
/// to disk, at the specified `path`.
///
/// If `time_selection` is specified, then only data for that specific timeline over that
/// specific time range will be accounted for.
pub fn save_database_to_file(
entity_db: &EntityDb,
path: std::path::PathBuf,
time_selection: Option<(re_entity_db::Timeline, re_log_types::TimeRangeF)>,
) -> anyhow::Result<impl FnOnce() -> anyhow::Result<std::path::PathBuf>> {
use re_data_store::TimeRange;

pub fn encode_to_file<'a>(
path: &std::path::Path,
messages: impl Iterator<Item = &'a re_log_types::LogMsg>,
) -> anyhow::Result<()> {
re_tracing::profile_function!();

entity_db.store().sort_indices_if_needed();

let set_store_info_msg = entity_db
.store_info_msg()
.map(|msg| LogMsg::SetStoreInfo(msg.clone()));

let time_filter = time_selection.map(|(timeline, range)| {
(
timeline,
TimeRange::new(range.min.floor(), range.max.ceil()),
)
});
let data_msgs: Result<Vec<_>, _> = entity_db
.store()
.to_data_tables(time_filter)
.map(|table| {
table
.to_arrow_msg()
.map(|msg| LogMsg::ArrowMsg(entity_db.store_id().clone(), msg))
})
.collect();

use anyhow::Context as _;
use re_log_types::LogMsg;
let data_msgs = data_msgs.with_context(|| "Failed to export to data tables")?;

let msgs = std::iter::once(set_store_info_msg)
.flatten() // option
.chain(data_msgs);

Ok(move || {
re_tracing::profile_scope!("save_to_file");

use anyhow::Context as _;
let file = std::fs::File::create(path.as_path())
.with_context(|| format!("Failed to create file at {path:?}"))?;
let mut file = std::fs::File::create(path)
.with_context(|| format!("Failed to create file at {path:?}"))?;

let encoding_options = re_log_encoding::EncodingOptions::COMPRESSED;
re_log_encoding::encoder::encode_owned(encoding_options, msgs, file)
.map(|_| path)
.context("Message encode")
})
let encoding_options = re_log_encoding::EncodingOptions::COMPRESSED;
re_log_encoding::encoder::encode(encoding_options, messages, &mut file)
.context("Message encode")
}
Loading
Loading