diff --git a/src/async_reader.rs b/src/async_reader.rs index f8201a5..c94ea0d 100644 --- a/src/async_reader.rs +++ b/src/async_reader.rs @@ -9,6 +9,9 @@ use reqwest::{Client, IntoUrl}; #[cfg(any(feature = "http-async", feature = "mmap-async-tokio"))] use tokio::io::AsyncReadExt; +use crate::cache::SearchResult; +#[cfg(any(feature = "http-async", feature = "mmap-async-tokio"))] +use crate::cache::{DirectoryCache, NoCache}; use crate::directory::{Directory, Entry}; use crate::error::Error; use crate::header::{HEADER_SIZE, MAX_INITIAL_BYTES}; @@ -19,17 +22,27 @@ use crate::mmap::MmapBackend; use crate::tile::tile_id; use crate::{Compression, Header}; -pub struct AsyncPmTilesReader { +pub struct AsyncPmTilesReader { backend: B, + cache: C, header: Header, root_directory: Directory, } -impl AsyncPmTilesReader { - /// Creates a new reader from a specified source and validates the provided PMTiles archive is valid. +impl AsyncPmTilesReader { + /// Creates a new cached reader from a specified source and validates the provided PMTiles archive is valid. /// /// Note: Prefer using new_with_* methods. pub async fn try_from_source(backend: B) -> Result { + Self::try_from_cached_source(backend, NoCache).await + } +} + +impl AsyncPmTilesReader { + /// Creates a new reader from a specified source and validates the provided PMTiles archive is valid. + /// + /// Note: Prefer using new_with_* methods. + pub async fn try_from_cached_source(backend: B, cache: C) -> Result { // Read the first 127 and up to 16,384 bytes to ensure we can initialize the header and root directory. let mut initial_bytes = backend.read(0, MAX_INITIAL_BYTES).await?; if initial_bytes.len() < HEADER_SIZE { @@ -47,6 +60,7 @@ impl AsyncPmTilesReader { Ok(Self { backend, + cache, header, root_directory, }) @@ -142,11 +156,21 @@ impl AsyncPmTilesReader { // the recursion is done as two functions because it is a bit cleaner, // and it allows directory to be cached later without cloning it first. let offset = (self.header.leaf_offset + entry.offset) as _; - let length = entry.length as _; - let dir = self.read_directory(offset, length).await.ok()?; - let entry = dir.find_tile_id(tile_id); - if let Some(entry) = entry { + let entry = match self.cache.get_dir_entry(offset, tile_id) { + SearchResult::NotCached => { + // Cache miss - read from backend + let length = entry.length as _; + let dir = self.read_directory(offset, length).await.ok()?; + let entry = dir.find_tile_id(tile_id).cloned(); + self.cache.insert_dir(offset, dir); + entry + } + SearchResult::NotFound => None, + SearchResult::Found(entry) => Some(entry), + }; + + if let Some(ref entry) = entry { if entry.is_leaf() { return if depth <= 4 { self.find_entry_rec(tile_id, entry, depth + 1).await @@ -156,7 +180,7 @@ impl AsyncPmTilesReader { } } - entry.cloned() + entry } async fn read_directory(&self, offset: usize, length: usize) -> Result { @@ -188,26 +212,50 @@ impl AsyncPmTilesReader { } #[cfg(feature = "http-async")] -impl AsyncPmTilesReader { +impl AsyncPmTilesReader { /// Creates a new PMTiles reader from a URL using the Reqwest backend. /// /// Fails if [url] does not exist or is an invalid archive. (Note: HTTP requests are made to validate it.) pub async fn new_with_url(client: Client, url: U) -> Result { + Self::new_with_cached_url(client, url, NoCache).await + } +} + +#[cfg(feature = "http-async")] +impl AsyncPmTilesReader { + /// Creates a new PMTiles reader with cache from a URL using the Reqwest backend. + /// + /// Fails if [url] does not exist or is an invalid archive. (Note: HTTP requests are made to validate it.) + pub async fn new_with_cached_url( + client: Client, + url: U, + cache: C, + ) -> Result { let backend = HttpBackend::try_from(client, url)?; - Self::try_from_source(backend).await + Self::try_from_cached_source(backend, cache).await } } #[cfg(feature = "mmap-async-tokio")] -impl AsyncPmTilesReader { +impl AsyncPmTilesReader { /// Creates a new PMTiles reader from a file path using the async mmap backend. /// /// Fails if [p] does not exist or is an invalid archive. pub async fn new_with_path>(path: P) -> Result { + Self::new_with_cached_path(path, NoCache).await + } +} + +#[cfg(feature = "mmap-async-tokio")] +impl AsyncPmTilesReader { + /// Creates a new cached PMTiles reader from a file path using the async mmap backend. + /// + /// Fails if [p] does not exist or is an invalid archive. + pub async fn new_with_cached_path>(path: P, cache: C) -> Result { let backend = MmapBackend::try_from(path).await?; - Self::try_from_source(backend).await + Self::try_from_cached_source(backend, cache).await } } diff --git a/src/cache.rs b/src/cache.rs new file mode 100644 index 0000000..0e3d05d --- /dev/null +++ b/src/cache.rs @@ -0,0 +1,60 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use crate::directory::{Directory, Entry}; + +pub enum SearchResult { + NotCached, + NotFound, + Found(Entry), +} + +impl From> for SearchResult { + fn from(entry: Option<&Entry>) -> Self { + match entry { + Some(entry) => SearchResult::Found(entry.clone()), + None => SearchResult::NotFound, + } + } +} + +/// A cache for PMTiles directories. +pub trait DirectoryCache { + /// Get a directory from the cache, using the offset as a key. + fn get_dir_entry(&self, offset: usize, tile_id: u64) -> SearchResult; + + /// Insert a directory into the cache, using the offset as a key. + /// Note that cache must be internally mutable. + fn insert_dir(&self, offset: usize, directory: Directory); +} + +pub struct NoCache; + +impl DirectoryCache for NoCache { + #[inline] + fn get_dir_entry(&self, _offset: usize, _tile_id: u64) -> SearchResult { + SearchResult::NotCached + } + + #[inline] + fn insert_dir(&self, _offset: usize, _directory: Directory) {} +} + +/// A simple HashMap-based implementation of a PMTiles directory cache. +#[derive(Default)] +pub struct HashMapCache { + pub cache: Arc>>, +} + +impl DirectoryCache for HashMapCache { + fn get_dir_entry(&self, offset: usize, tile_id: u64) -> SearchResult { + if let Some(dir) = self.cache.lock().unwrap().get(&offset) { + return dir.find_tile_id(tile_id).into(); + } + SearchResult::NotCached + } + + fn insert_dir(&self, offset: usize, directory: Directory) { + self.cache.lock().unwrap().insert(offset, directory); + } +} diff --git a/src/directory.rs b/src/directory.rs index ccceab0..a9e42c5 100644 --- a/src/directory.rs +++ b/src/directory.rs @@ -5,7 +5,7 @@ use varint_rs::VarintReader; use crate::error::Error; -pub(crate) struct Directory { +pub struct Directory { entries: Vec, } @@ -81,7 +81,7 @@ impl TryFrom for Directory { } #[derive(Clone, Default, Debug)] -pub(crate) struct Entry { +pub struct Entry { pub(crate) tile_id: u64, pub(crate) offset: u64, pub(crate) length: u32, @@ -90,7 +90,7 @@ pub(crate) struct Entry { #[cfg(any(feature = "http-async", feature = "mmap-async-tokio"))] impl Entry { - pub fn is_leaf(&self) -> bool { + pub(crate) fn is_leaf(&self) -> bool { self.run_length == 0 } } diff --git a/src/lib.rs b/src/lib.rs index b0de680..b7c9ecf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,10 +1,14 @@ #![forbid(unsafe_code)] +mod tile; + +mod header; pub use crate::header::{Compression, Header, TileType}; mod directory; +pub use directory::{Directory, Entry}; + pub mod error; -mod header; #[cfg(feature = "http-async")] pub mod http; @@ -14,7 +18,9 @@ pub mod mmap; #[cfg(any(feature = "http-async", feature = "mmap-async-tokio"))] pub mod async_reader; -pub mod tile; + +#[cfg(any(feature = "http-async", feature = "mmap-async-tokio"))] +pub mod cache; #[cfg(test)] mod tests {