Skip to content

Commit

Permalink
DataLoaders 7: support for custom DataLoaders (#4566)
Browse files Browse the repository at this point in the history
Add support for registering custom loaders, the same way we offer
support for registering custom space views and store subscribers.

Checks:
- [x] cargo r -p custom_data_loader -- Cargo.toml

---

Part of a series of PRs to make it possible to load _any_ file from the
local filesystem, by any means, on web and native:
- #4516
- #4517 
- #4518 
- #4519 
- #4520 
- #4521 
- #4565
- #4566
- #4567
  • Loading branch information
teh-cmc authored Dec 19, 2023
1 parent b7f7a55 commit 47bec77
Show file tree
Hide file tree
Showing 8 changed files with 157 additions and 19 deletions.
8 changes: 8 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

22 changes: 20 additions & 2 deletions crates/re_data_source/src/data_loader/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -229,8 +229,26 @@ static BUILTIN_LOADERS: Lazy<Vec<Arc<dyn DataLoader>>> = Lazy::new(|| {

/// Iterator over all registered [`DataLoader`]s.
#[inline]
pub fn iter_loaders() -> impl ExactSizeIterator<Item = Arc<dyn DataLoader>> {
BUILTIN_LOADERS.clone().into_iter()
pub fn iter_loaders() -> impl Iterator<Item = Arc<dyn DataLoader>> {
BUILTIN_LOADERS
.clone()
.into_iter()
.chain(CUSTOM_LOADERS.read().clone())
}

/// Keeps track of all custom [`DataLoader`]s.
///
/// Use [`register_custom_data_loader`] to add new loaders.
static CUSTOM_LOADERS: Lazy<parking_lot::RwLock<Vec<Arc<dyn DataLoader>>>> =
Lazy::new(parking_lot::RwLock::default);

/// Register a custom [`DataLoader`].
///
/// Any time the Rerun Viewer opens a file or directory, this custom loader will be notified.
/// Refer to [`DataLoader`]'s documentation for more information.
#[inline]
pub fn register_custom_data_loader(loader: impl DataLoader + 'static) {
CUSTOM_LOADERS.write().push(Arc::new(loader));
}

// ---
Expand Down
4 changes: 2 additions & 2 deletions crates/re_data_source/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ mod web_sockets;
mod load_stdin;

pub use self::data_loader::{
iter_loaders, ArchetypeLoader, DataLoader, DataLoaderError, DirectoryLoader, LoadedData,
RrdLoader,
iter_loaders, register_custom_data_loader, ArchetypeLoader, DataLoader, DataLoaderError,
DirectoryLoader, LoadedData, RrdLoader,
};
pub use self::data_source::DataSource;
pub use self::load_file::{extension, load_from_file_contents};
Expand Down
24 changes: 9 additions & 15 deletions crates/re_data_source/src/load_file.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,14 +34,16 @@ pub fn load_from_path(

re_log::info!("Loading {path:?}…");

let store_info = prepare_store_info(store_id, file_source, path, path.is_dir());
let data = load(store_id, path, None)?;

// If we reach this point, then at least one compatible `DataLoader` has been found.
let store_info = prepare_store_info(store_id, file_source, path);
if let Some(store_info) = store_info {
if tx.send(store_info).is_err() {
return Ok(()); // other end has hung up.
}
}

let data = load(store_id, path, None)?;
send(store_id, data, tx);

Ok(())
Expand All @@ -67,14 +69,16 @@ pub fn load_from_file_contents(

re_log::info!("Loading {filepath:?}…");

let store_info = prepare_store_info(store_id, file_source, filepath, false);
let data = load(store_id, filepath, Some(contents))?;

// If we reach this point, then at least one compatible `DataLoader` has been found.
let store_info = prepare_store_info(store_id, file_source, filepath);
if let Some(store_info) = store_info {
if tx.send(store_info).is_err() {
return Ok(()); // other end has hung up.
}
}

let data = load(store_id, filepath, Some(contents))?;
send(store_id, data, tx);

Ok(())
Expand All @@ -92,20 +96,11 @@ pub fn extension(path: &std::path::Path) -> String {
.to_string()
}

/// Returns whether the given path is supported by builtin [`crate::DataLoader`]s.
///
/// This does _not_ access the filesystem.
#[inline]
pub fn is_associated_with_builtin_loader(path: &std::path::Path, is_dir: bool) -> bool {
is_dir || crate::is_supported_file_extension(&extension(path))
}

/// Prepares an adequate [`re_log_types::StoreInfo`] [`LogMsg`] given the input.
pub(crate) fn prepare_store_info(
store_id: &re_log_types::StoreId,
file_source: FileSource,
path: &std::path::Path,
is_dir: bool,
) -> Option<LogMsg> {
re_tracing::profile_function!(path.display().to_string());

Expand All @@ -114,10 +109,9 @@ pub(crate) fn prepare_store_info(
let app_id = re_log_types::ApplicationId(path.display().to_string());
let store_source = re_log_types::StoreSource::File { file_source };

let is_builtin = is_associated_with_builtin_loader(path, is_dir);
let is_rrd = crate::SUPPORTED_RERUN_EXTENSIONS.contains(&extension(path).as_str());

(!is_rrd && is_builtin).then(|| {
(!is_rrd).then(|| {
LogMsg::SetStoreInfo(SetStoreInfo {
row_id: re_log_types::RowId::new(),
info: re_log_types::StoreInfo {
Expand Down
13 changes: 13 additions & 0 deletions examples/rust/custom_data_loader/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
[package]
name = "custom_data_loader"
version = "0.12.0-alpha.1+dev"
edition = "2021"
rust-version = "1.72"
license = "MIT OR Apache-2.0"
publish = false

[dependencies]
rerun = { path = "../../../crates/rerun", features = ["native_viewer"] }

[build-dependencies]
re_build_tools = { path = "../../../crates/re_build_tools" }
20 changes: 20 additions & 0 deletions examples/rust/custom_data_loader/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
---
title: Custom data-loader
tags: [data-loader, extension]
rust: https://github.com/rerun-io/rerun/tree/latest/examples/rust/custom_data_loader/src/main.rs?speculative-link
---

<picture>
<img src="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/full.png" alt="">
<source media="(max-width: 480px)" srcset="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/480w.png">
<source media="(max-width: 768px)" srcset="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/768w.png">
<source media="(max-width: 1024px)" srcset="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/1024w.png">
<source media="(max-width: 1200px)" srcset="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/1200w.png">
</picture>

This example demonstrates how to implement and register a `DataLoader` into the Rerun Viewer in order to add support for loading arbitrary files.

Usage:
```sh
$ cargo r -p custom_data_loader -- path/to/some/file
```
3 changes: 3 additions & 0 deletions examples/rust/custom_data_loader/build.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
fn main() {
re_build_tools::export_build_info_vars_for_crate("custom_data_loader");
}
82 changes: 82 additions & 0 deletions examples/rust/custom_data_loader/src/main.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
//! This example demonstrates how to implement and register a [`re_data_source::DataLoader`] into
//! the Rerun Viewer in order to add support for loading arbitrary files.
//!
//! Usage:
//! ```sh
//! $ cargo r -p custom_data_loader -- path/to/some/file
//! ```
use rerun::{
external::{anyhow, re_build_info, re_data_source, re_log, tokio},
log::{DataRow, RowId},
EntityPath, TimePoint,
};

#[tokio::main]
async fn main() -> anyhow::Result<std::process::ExitCode> {
re_log::setup_native_logging();

re_data_source::register_custom_data_loader(HashLoader);

let build_info = re_build_info::build_info!();
rerun::run(build_info, rerun::CallSource::Cli, std::env::args())
.await
.map(std::process::ExitCode::from)
}

// ---

/// A custom [`re_data_source::DataLoader`] that logs the hash of file as a [`rerun::TextDocument`].
struct HashLoader;

impl re_data_source::DataLoader for HashLoader {
fn name(&self) -> String {
"rerun.data_loaders.HashLoader".into()
}

fn load_from_path(
&self,
_store_id: rerun::external::re_log_types::StoreId,
path: std::path::PathBuf,
tx: std::sync::mpsc::Sender<re_data_source::LoadedData>,
) -> Result<(), re_data_source::DataLoaderError> {
let contents = std::fs::read(&path)?;
if path.is_dir() {
return Err(re_data_source::DataLoaderError::Incompatible(path)); // simply not interested
}
hash_and_log(&tx, &path, &contents)
}

fn load_from_file_contents(
&self,
_store_id: rerun::external::re_log_types::StoreId,
filepath: std::path::PathBuf,
contents: std::borrow::Cow<'_, [u8]>,
tx: std::sync::mpsc::Sender<re_data_source::LoadedData>,
) -> Result<(), re_data_source::DataLoaderError> {
hash_and_log(&tx, &filepath, &contents)
}
}

fn hash_and_log(
tx: &std::sync::mpsc::Sender<re_data_source::LoadedData>,
filepath: &std::path::Path,
contents: &[u8],
) -> Result<(), re_data_source::DataLoaderError> {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

let mut h = DefaultHasher::new();
contents.hash(&mut h);

let doc = rerun::TextDocument::new(format!("{:08X}", h.finish()))
.with_media_type(rerun::MediaType::TEXT);

let entity_path = EntityPath::from_file_path(filepath);
let entity_path = format!("{entity_path}/hashed").into();
let row = DataRow::from_archetype(RowId::new(), TimePoint::timeless(), entity_path, &doc)?;

tx.send(row.into()).ok();

Ok(())
}

0 comments on commit 47bec77

Please sign in to comment.