From 8bd98fac02fe141d196ca2c88e7240a792c71a00 Mon Sep 17 00:00:00 2001 From: Adrien Guillo Date: Fri, 21 Aug 2026 08:20:12 -0400 Subject: [PATCH] POC: Remove deeply nested stack of boxed `Storage::get_slice` futures --- quickwit/Cargo.lock | 8 +- quickwit/Cargo.toml | 2 +- quickwit/quickwit-cli/src/lib.rs | 2 +- .../quickwit-config/src/node_config/mod.rs | 2 +- .../quickwit-datafusion/src/storage_bridge.rs | 8 +- .../src/storage_directory.rs | 38 +-- .../quickwit-index-management/src/index.rs | 7 +- .../src/source/doc_file_reader.rs | 2 +- quickwit/quickwit-indexing/src/source/mod.rs | 2 +- quickwit/quickwit-search/src/fetch_docs.rs | 20 +- quickwit/quickwit-search/src/leaf.rs | 190 ++++++++------ quickwit/quickwit-search/src/lib.rs | 1 + .../quickwit-search/src/list_fields/leaf.rs | 18 +- quickwit/quickwit-search/src/list_terms.rs | 14 +- .../quickwit-storage/src/bundle_storage.rs | 58 +++-- quickwit/quickwit-storage/src/cache/mod.rs | 12 +- .../src/cache/storage_with_cache.rs | 45 ++-- .../quickwit-storage/src/counting_storage.rs | 76 +++--- quickwit/quickwit-storage/src/debouncer.rs | 202 ++++++++------- quickwit/quickwit-storage/src/lib.rs | 8 +- .../src/local_file_storage.rs | 79 +++--- .../src/object_storage/azure_blob_storage.rs | 61 +++-- .../object_storage/s3_compatible_storage.rs | 45 ++-- .../s3_compatible_storage_resolver.rs | 23 +- .../src/opendal_storage/base.rs | 46 ++-- .../opendal_storage/google_cloud_storage.rs | 21 +- .../src/opendal_storage/mod.rs | 2 +- .../quickwit-storage/src/prefix_storage.rs | 37 ++- quickwit/quickwit-storage/src/ram_storage.rs | 53 ++-- .../src/split_cache/download_task.rs | 2 +- .../quickwit-storage/src/split_cache/mod.rs | 7 +- quickwit/quickwit-storage/src/storage.rs | 116 ++++++++- .../quickwit-storage/src/storage_factory.rs | 231 ++++++++++++++++-- .../quickwit-storage/src/storage_resolver.rs | 70 ++---- .../src/timeout_and_retry_storage.rs | 162 ++++++------ 35 files changed, 1069 insertions(+), 601 deletions(-) diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index 042011984fc..546ed4677c6 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -6262,9 +6262,9 @@ dependencies = [ [[package]] name = "mockall" -version = "0.14.0" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f58d964098a5f9c6b63d0798e5372fd04708193510a7af313c22e9f29b7b620b" +checksum = "1a6ceddfe3ce334925e96bf420fdb2dcee5bed6c632a168ece622676dadeaf8a" dependencies = [ "cfg-if", "downcast", @@ -6276,9 +6276,9 @@ dependencies = [ [[package]] name = "mockall_derive" -version = "0.14.0" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca41ce716dda6a9be188b385aa78ee5260fc25cd3802cb2a8afdc6afbe6b6dbf" +checksum = "9cfe16fbe8a314aeec0b861ac24e60b1e123e97634bab045475b9d6a18416fd8" dependencies = [ "cfg-if", "proc-macro2", diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml index 8904ed1dd3c..edcdbd81306 100644 --- a/quickwit/Cargo.toml +++ b/quickwit/Cargo.toml @@ -180,7 +180,7 @@ metrics-exporter-prometheus = { version = "0.18", default-features = false } metrics-util = "0.20" mime_guess = "2.0" mini-moka = "0.10" -mockall = "0.14" +mockall = "0.15" mrecordlog = { git = "https://github.com/quickwit-oss/mrecordlog", rev = "3b3562ef" } new_string_template = "1.5" nom = "8.0" diff --git a/quickwit/quickwit-cli/src/lib.rs b/quickwit/quickwit-cli/src/lib.rs index c3dfd952cc3..cd9e68c1b0e 100644 --- a/quickwit/quickwit-cli/src/lib.rs +++ b/quickwit/quickwit-cli/src/lib.rs @@ -34,7 +34,7 @@ use quickwit_metastore::{IndexMetadataResponseExt, MetastoreResolver}; use quickwit_proto::metastore::{IndexMetadataRequest, MetastoreService, MetastoreServiceClient}; use quickwit_rest_client::models::Timeout; use quickwit_rest_client::rest_client::{DEFAULT_BASE_URL, QuickwitClient, QuickwitClientBuilder}; -use quickwit_storage::{StorageResolver, load_file}; +use quickwit_storage::{Storage, StorageResolver, load_file}; use reqwest::Url; use tabled::settings::object::Rows; use tabled::settings::panel::Header; diff --git a/quickwit/quickwit-config/src/node_config/mod.rs b/quickwit/quickwit-config/src/node_config/mod.rs index 5e02be1af39..5ab7b3c1388 100644 --- a/quickwit/quickwit-config/src/node_config/mod.rs +++ b/quickwit/quickwit-config/src/node_config/mod.rs @@ -633,7 +633,7 @@ impl std::fmt::Display for CachePolicy { /// This policy is inspired by this guidance. It does not track instanteneous throughput, but /// computes an overall timeout using the following formula: /// `timeout_offset + num_bytes_get_request / min_throughtput` -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] pub struct StorageTimeoutPolicy { pub min_throughtput_bytes_per_secs: u64, pub timeout_millis: u64, diff --git a/quickwit/quickwit-datafusion/src/storage_bridge.rs b/quickwit/quickwit-datafusion/src/storage_bridge.rs index d1c9b6a096e..9c6f43d0035 100644 --- a/quickwit/quickwit-datafusion/src/storage_bridge.rs +++ b/quickwit/quickwit-datafusion/src/storage_bridge.rs @@ -43,18 +43,18 @@ use object_store::{ Result as ObjectStoreResult, }; use quickwit_common::uri::Uri; -use quickwit_storage::{Storage, StorageResolver}; +use quickwit_storage::{ResolvedStorage, Storage, StorageResolver}; use tokio::sync::OnceCell; /// Adapts Quickwit's `Storage` trait to DataFusion's `ObjectStore` interface. /// /// Construction is sync and cheap: a `Uri` plus a `StorageResolver` handle -/// (resolver is `Clone`). The underlying `Arc` is materialised +/// (resolver is `Clone`). The underlying storage handle is materialised /// on the first async method call and cached for the wrapper's lifetime. pub struct QuickwitObjectStore { index_uri: Uri, storage_resolver: StorageResolver, - storage: OnceCell>, + storage: OnceCell>, } impl QuickwitObjectStore { @@ -68,7 +68,7 @@ impl QuickwitObjectStore { /// Returns the handle to the underlying `Storage`, resolving it via the /// `StorageResolver` if this is the first call. - async fn storage(&self) -> ObjectStoreResult<&Arc> { + async fn storage(&self) -> ObjectStoreResult<&Arc> { self.storage .get_or_try_init(|| async { self.storage_resolver diff --git a/quickwit/quickwit-directories/src/storage_directory.rs b/quickwit/quickwit-directories/src/storage_directory.rs index 31ee6e3f736..aa8251cac0f 100644 --- a/quickwit/quickwit-directories/src/storage_directory.rs +++ b/quickwit/quickwit-directories/src/storage_directory.rs @@ -20,24 +20,24 @@ use std::{fmt, io}; use async_trait::async_trait; use quickwit_common::uri::Uri; -use quickwit_storage::{OwnedBytes, Storage}; +use quickwit_storage::{OwnedBytes, Storage, StorageGetSlice}; use tantivy::directory::FileHandle; use tantivy::directory::error::OpenReadError; use tantivy::{Directory, HasLen}; use tracing::{error, instrument}; -struct StorageDirectoryFileHandle { - storage_directory: StorageDirectory, +struct StorageDirectoryFileHandle { + storage_directory: StorageDirectory, path: PathBuf, } -impl HasLen for StorageDirectoryFileHandle { +impl HasLen for StorageDirectoryFileHandle { fn len(&self) -> usize { unimplemented!() } } -impl fmt::Debug for StorageDirectoryFileHandle { +impl fmt::Debug for StorageDirectoryFileHandle { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!( f, @@ -48,7 +48,7 @@ impl fmt::Debug for StorageDirectoryFileHandle { } #[async_trait] -impl FileHandle for StorageDirectoryFileHandle { +impl FileHandle for StorageDirectoryFileHandle { fn read_bytes(&self, _byte_range: Range) -> io::Result { Err(unsupported_operation(&self.path)) } @@ -76,25 +76,29 @@ impl FileHandle for StorageDirectoryFileHandle { /// This directory is fetch slices of data to a possibly distant storage /// everytime `read_bytes` is called. #[derive(Clone)] -pub struct StorageDirectory { - storage: Arc, +pub struct StorageDirectory { + storage: T, } -impl Debug for StorageDirectory { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "StorageDirectory({:?})", self.uri()) +impl Debug for StorageDirectory { + fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter + .debug_tuple("StorageDirectory") + .field(&self.uri()) + .finish() } } -impl StorageDirectory { +impl StorageDirectory { /// Creates a new StorageDirectory, backed by the given `storage`. - pub fn new(storage: Arc) -> StorageDirectory { - StorageDirectory { storage } + pub fn new(storage: T) -> Self { + Self { storage } } /// Fetches a slice of byte from a file asynchronously. - pub async fn get_slice(&self, path: &Path, range: Range) -> io::Result { - let payload: OwnedBytes = self.storage.get_slice(path, range).await?; + pub async fn get_slice(&self, path: &Path, range: Range) -> io::Result + where T: StorageGetSlice { + let payload: OwnedBytes = self.storage.get_slice_unboxed(path, range).await?; Ok(payload) } @@ -116,7 +120,7 @@ fn unsupported_operation(path: &Path) -> io::Error { io::Error::other(format!("{error}: {}", path.display())) } -impl Directory for StorageDirectory { +impl Directory for StorageDirectory { fn get_file_handle(&self, path: &Path) -> Result, OpenReadError> { Ok(Arc::new(StorageDirectoryFileHandle { storage_directory: self.clone(), diff --git a/quickwit/quickwit-index-management/src/index.rs b/quickwit/quickwit-index-management/src/index.rs index af2d3bc34d5..d4e6bf1d68f 100644 --- a/quickwit/quickwit-index-management/src/index.rs +++ b/quickwit/quickwit-index-management/src/index.rs @@ -14,6 +14,7 @@ use std::collections::{HashMap, HashSet}; use std::path::Path; +use std::sync::Arc; use std::time::Duration; use futures_util::StreamExt; @@ -37,7 +38,7 @@ use quickwit_proto::metastore::{ }; use quickwit_proto::types::{IndexUid, SplitId}; use quickwit_proto::{ServiceError, ServiceErrorCode}; -use quickwit_storage::{StorageResolver, StorageResolverError}; +use quickwit_storage::{Storage, StorageResolver, StorageResolverError}; use thiserror::Error; use tracing::{error, info}; @@ -385,7 +386,7 @@ impl IndexService { .deserialize_index_metadata()?; let index_uid = index_metadata.index_uid.clone(); let index_config = index_metadata.into_index_config(); - let storage = self + let storage: Arc = self .storage_resolver .resolve(&index_config.index_uri) .await?; @@ -426,7 +427,7 @@ impl IndexService { .deserialize_index_metadata()?; let index_uid = index_metadata.index_uid.clone(); let index_config = index_metadata.into_index_config(); - let storage = self + let storage: Arc = self .storage_resolver .resolve(&index_config.index_uri) .await?; diff --git a/quickwit/quickwit-indexing/src/source/doc_file_reader.rs b/quickwit/quickwit-indexing/src/source/doc_file_reader.rs index 90bc99c01ad..a56d9bb4d07 100644 --- a/quickwit/quickwit-indexing/src/source/doc_file_reader.rs +++ b/quickwit/quickwit-indexing/src/source/doc_file_reader.rs @@ -23,7 +23,7 @@ use quickwit_common::uri::Uri; use quickwit_metastore::checkpoint::PartitionId; use quickwit_proto::metastore::SourceType; use quickwit_proto::types::Position; -use quickwit_storage::StorageResolver; +use quickwit_storage::{Storage, StorageResolver}; use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader}; use super::{BATCH_NUM_BYTES_LIMIT, BatchBuilder}; diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs index d6b76d32d48..b84b6f66f65 100644 --- a/quickwit/quickwit-indexing/src/source/mod.rs +++ b/quickwit/quickwit-indexing/src/source/mod.rs @@ -112,7 +112,7 @@ use quickwit_proto::metastore::{ MetastoreServiceClient, SourceType, }; use quickwit_proto::types::{IndexUid, NodeIdRef, PipelineUid, ShardId}; -use quickwit_storage::StorageResolver; +use quickwit_storage::{Storage, StorageResolver}; use serde_json::Value as JsonValue; pub use source_factory::{SourceFactory, SourceLoader, TypedSourceFactory}; pub use source_sink::SourceSink; diff --git a/quickwit/quickwit-search/src/fetch_docs.rs b/quickwit/quickwit-search/src/fetch_docs.rs index 29be8eb2be2..c1bf21cd610 100644 --- a/quickwit/quickwit-search/src/fetch_docs.rs +++ b/quickwit/quickwit-search/src/fetch_docs.rs @@ -22,7 +22,7 @@ use quickwit_doc_mapper::DocMapper; use quickwit_proto::search::{ FetchDocsResponse, PartialHit, SnippetRequest, SplitIdAndFooterOffsets, }; -use quickwit_storage::Storage; +use quickwit_storage::StorageGetSlice; use tantivy::query::Query; use tantivy::schema::document::CompactDocValue; use tantivy::schema::{Document as DocumentTrait, Field, TantivyDocument, Value}; @@ -38,10 +38,10 @@ const SNIPPET_MAX_NUM_CHARS: usize = 150; /// Given a list of global doc address, fetches all the documents and /// returns them as a hashmap. -async fn fetch_docs_to_map( +async fn fetch_docs_to_map( searcher_context: Arc, mut global_doc_addrs: Vec, - index_storage: Arc, + storage: T, splits: &[SplitIdAndFooterOffsets], doc_mapper: Arc, snippet_request_opt: Option<&SnippetRequest>, @@ -70,7 +70,7 @@ async fn fetch_docs_to_map( fetch_docs_in_split( searcher_context.clone(), global_doc_addrs, - index_storage.clone(), + storage.clone(), split_and_offset, doc_mapper.clone(), snippet_request_opt, @@ -100,10 +100,10 @@ async fn fetch_docs_to_map( /// This function takes a list of partial hits (possibly from different splits) /// and the storage associated to an index, fetches the document from /// the split document stores, and returns the full hits. -pub async fn fetch_docs( +pub async fn fetch_docs( searcher_context: Arc, partial_hits: Vec, - index_storage: Arc, + storage: T, splits: &[SplitIdAndFooterOffsets], doc_mapper: Arc, snippet_request_opt: Option<&SnippetRequest>, @@ -116,7 +116,7 @@ pub async fn fetch_docs( let mut global_doc_addr_to_doc_json = fetch_docs_to_map( searcher_context, global_doc_addrs, - index_storage, + storage, splits, doc_mapper, snippet_request_opt, @@ -154,10 +154,10 @@ struct Document { /// Fetching docs from a specific split. #[instrument(skip_all, fields(split_id = split.split_id, num_docs = global_doc_addrs.len()))] -async fn fetch_docs_in_split( +async fn fetch_docs_in_split( searcher_context: Arc, mut global_doc_addrs: Vec, - index_storage: Arc, + storage: T, split: &SplitIdAndFooterOffsets, doc_mapper: Arc, snippet_request_opt: Option<&SnippetRequest>, @@ -167,7 +167,7 @@ async fn fetch_docs_in_split( // when fetching docs as we will fetch them only once. let (mut index, _) = open_index_with_caches( &searcher_context, - index_storage, + storage, split, Some(doc_mapper.tokenizer_manager()), None, diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index 043da090a27..5e62b1acec1 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -17,7 +17,7 @@ use std::collections::binary_heap::PeekMut; use std::collections::{BinaryHeap, HashMap, HashSet}; use std::num::NonZeroUsize; use std::ops::Bound; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, RwLock}; @@ -45,7 +45,7 @@ use quickwit_query::query_ast::{ use quickwit_query::tokenizers::TokenizerManager; use quickwit_storage::{ BundleStorage, ByteRangeCache, CountingStorage, MemorySizedCache, OwnedBytes, SearchSplitCache, - Storage, StorageResolver, TimeoutAndRetryStorage, wrap_storage_with_cache, + Storage, StorageGetSlice, StorageResolver, TimeoutAndRetryStorage, wrap_storage_with_cache, }; use tantivy::aggregation::AggContextParams; use tantivy::aggregation::agg_req::{AggregationVariants, Aggregations}; @@ -124,8 +124,8 @@ fn greedy_batch_split( batches } -async fn get_split_footer_from_cache_or_fetch( - index_storage: Arc, +async fn get_split_footer_from_cache_or_fetch( + storage: &T, split_and_footer_offsets: &SplitIdAndFooterOffsets, footer_cache: &MemorySizedCache, ) -> anyhow::Result { @@ -136,8 +136,8 @@ async fn get_split_footer_from_cache_or_fetch( } } let split_file = PathBuf::from(format!("{}.split", split_and_footer_offsets.split_id)); - let footer_data_opt = index_storage - .get_slice( + let footer_data_opt = storage + .get_slice_unboxed( &split_file, split_and_footer_offsets.split_footer_start as usize ..split_and_footer_offsets.split_footer_end as usize, @@ -146,7 +146,7 @@ async fn get_split_footer_from_cache_or_fetch( .with_context(|| { format!( "failed to fetch hotcache and footer from {} for split `{}`", - index_storage.uri(), + storage.uri(), split_and_footer_offsets.split_id ) })?; @@ -159,55 +159,50 @@ async fn get_split_footer_from_cache_or_fetch( Ok(footer_data_opt) } -/// Returns hotcache_bytes and the split directory (`BundleStorage`) with cache layer: +/// A split bundle with an optional split-cache layer, kept as a concrete enum so the cache branch +/// does not force an inner `Arc`. +pub(crate) enum OpenedSplitBundle { + WithoutCache(BundleStorage), + WithCache(BundleStorage>), +} + +impl OpenedSplitBundle { + pub(crate) async fn get_all(&self, path: &Path) -> quickwit_storage::StorageResult { + match self { + Self::WithoutCache(storage) => storage.get_all(path).await, + Self::WithCache(storage) => storage.get_all(path).await, + } + } +} + +/// Returns hotcache bytes and the split bundle with its optional cache layer: /// - A split footer cache given by `SearcherContext.split_footer_cache`. -pub(crate) async fn open_split_bundle( +pub(crate) async fn open_split_bundle( searcher_context: &SearcherContext, - index_storage: Arc, + storage: T, split_and_footer_offsets: &SplitIdAndFooterOffsets, -) -> anyhow::Result<(OwnedBytes, BundleStorage)> { +) -> anyhow::Result<(OwnedBytes, OpenedSplitBundle)> { let split_file = PathBuf::from(format!("{}.split", split_and_footer_offsets.split_id)); let footer_data = get_split_footer_from_cache_or_fetch( - index_storage.clone(), + &storage, split_and_footer_offsets, &searcher_context.split_footer_cache, ) .await?; - // We wrap the top-level storage with the split cache. - // This is before the bundle storage: at this point, this storage is reading `.split` files. - let index_storage_with_split_cache = - if let Some(split_cache) = searcher_context.split_cache_opt.as_ref() { - SearchSplitCache::wrap_storage(split_cache.clone(), index_storage.clone()) - } else { - index_storage.clone() - }; - - let (bundle_storage, hotcache_bytes) = BundleStorage::open_from_split_bytes( - index_storage_with_split_cache, - split_file, - footer_data, - )?; - - Ok((hotcache_bytes, bundle_storage)) -} - -/// Add a storage proxy to retry `get_slice` requests if they are taking too long, -/// if configured in the searcher config. -/// -/// The goal here is too ensure a low latency. -fn configure_storage_retries( - searcher_context: &SearcherContext, - index_storage: Arc, -) -> Arc { - if let Some(storage_timeout_policy) = &searcher_context.searcher_config.storage_timeout_policy { - Arc::new(TimeoutAndRetryStorage::new( - index_storage, - storage_timeout_policy.clone(), - )) - } else { - index_storage - } + // The split cache sits before the bundle storage because it reads complete `.split` files. + if let Some(split_cache) = searcher_context.split_cache_opt.clone() { + let storage = SearchSplitCache::wrap_storage(split_cache, storage); + let (bundle_storage, hotcache_bytes) = + BundleStorage::open_from_split_bytes(storage, split_file, footer_data)?; + return Ok((hotcache_bytes, OpenedSplitBundle::WithCache(bundle_storage))); + } + let (bundle_storage, hotcache_bytes) = + BundleStorage::open_from_split_bytes(storage, split_file, footer_data)?; + Ok(( + hotcache_bytes, + OpenedSplitBundle::WithoutCache(bundle_storage), + )) } /// Opens a `tantivy::Index` for the given split with several cache layers: @@ -215,28 +210,85 @@ fn configure_storage_retries( /// - A fast fields cache given by `SearcherContext.storage_long_term_cache`. /// - An ephemeral unbounded cache directory (whose lifetime is tied to the returned `Index` if no /// `ByteRangeCache` is provided). -pub(crate) async fn open_index_with_caches( +pub(crate) async fn open_index_with_caches( searcher_context: &SearcherContext, - index_storage: Arc, + storage: T, split_and_footer_offsets: &SplitIdAndFooterOffsets, tokenizer_manager: Option<&TokenizerManager>, ephemeral_unbounded_cache: Option, ) -> anyhow::Result<(Index, HotDirectory)> { - let index_storage_with_retry_on_timeout = - configure_storage_retries(searcher_context, index_storage); - - let (hotcache_bytes, bundle_storage) = open_split_bundle( + if let Some(storage_timeout_policy) = searcher_context.searcher_config.storage_timeout_policy { + let storage = TimeoutAndRetryStorage::new(storage, storage_timeout_policy); + return open_index_with_split_cache( + searcher_context, + storage, + split_and_footer_offsets, + tokenizer_manager, + ephemeral_unbounded_cache, + ) + .await; + } + open_index_with_split_cache( searcher_context, - index_storage_with_retry_on_timeout, + storage, + split_and_footer_offsets, + tokenizer_manager, + ephemeral_unbounded_cache, + ) + .await +} + +async fn open_index_with_split_cache( + searcher_context: &SearcherContext, + storage: T, + split_and_footer_offsets: &SplitIdAndFooterOffsets, + tokenizer_manager: Option<&TokenizerManager>, + ephemeral_unbounded_cache: Option, +) -> anyhow::Result<(Index, HotDirectory)> { + let split_file = PathBuf::from(format!("{}.split", split_and_footer_offsets.split_id)); + let footer_data = get_split_footer_from_cache_or_fetch( + &storage, split_and_footer_offsets, + &searcher_context.split_footer_cache, ) .await?; - let bundle_storage_with_cache = wrap_storage_with_cache( - searcher_context.fast_fields_cache.clone(), - Arc::new(bundle_storage), - ); + if let Some(split_cache) = searcher_context.split_cache_opt.clone() { + let storage = SearchSplitCache::wrap_storage(split_cache, storage); + return open_index_from_split_storage( + searcher_context, + storage, + split_file, + footer_data, + tokenizer_manager, + ephemeral_unbounded_cache, + ); + } + open_index_from_split_storage( + searcher_context, + storage, + split_file, + footer_data, + tokenizer_manager, + ephemeral_unbounded_cache, + ) +} +fn open_index_from_split_storage( + searcher_context: &SearcherContext, + storage: T, + split_file: PathBuf, + footer_data: OwnedBytes, + tokenizer_manager: Option<&TokenizerManager>, + ephemeral_unbounded_cache: Option, +) -> anyhow::Result<(Index, HotDirectory)> { + let (bundle_storage, hotcache_bytes) = + BundleStorage::open_from_split_bytes(storage, split_file, footer_data)?; + let bundle_storage_with_cache = + wrap_storage_with_cache(searcher_context.fast_fields_cache.clone(), bundle_storage); + + // Keep the complete storage stack concrete. Tantivy erases the `FileHandle` once when the + // directory is opened; `get_slice` remains statically dispatched below that boundary. let directory = StorageDirectory::new(bundle_storage_with_cache); let hot_directory = if let Some(cache) = ephemeral_unbounded_cache { @@ -652,10 +704,10 @@ pub(crate) fn term_absence_cache_key(term: &Term) -> String { } /// Apply a leaf search on a single split. -async fn leaf_search_single_split( +async fn leaf_search_single_split( search_request: SearchRequest, ctx: Arc, - storage: Arc, + storage: T, split: SplitIdAndFooterOffsets, search_permit: &mut SearchPermit, ) -> crate::Result> { @@ -1945,10 +1997,10 @@ async fn schedule_search_tasks( /// [PartialHit] candidates. The root will be in /// charge to consolidate, identify the actual final top hits to display, and /// fetch the actual documents to convert the partial hits into actual Hits. -pub async fn single_doc_mapping_leaf_search( +pub async fn single_doc_mapping_leaf_search( searcher_context: Arc, request: Arc, - index_storage: Arc, + storage: T, splits: Vec, doc_mapper: Arc, ) -> Result { @@ -2000,7 +2052,7 @@ pub async fn single_doc_mapping_leaf_search( &searcher_context, &request, &doc_mapper, - index_storage.uri().clone(), + storage.uri().clone(), offloaded_search_tasks, &incremental_merge_collector_arc, ); @@ -2015,7 +2067,7 @@ pub async fn single_doc_mapping_leaf_search( }); let run_local_search_tasks_fut = run_local_search_tasks( local_search_tasks, - index_storage, + storage, split_filter_arc, leaf_search_context, ); @@ -2067,9 +2119,9 @@ pub async fn single_doc_mapping_leaf_search( Ok(leaf_search_response) } -async fn run_local_search_tasks( +async fn run_local_search_tasks( local_search_tasks: Vec, - index_storage: Arc, + storage: T, split_filter_arc: Arc>, leaf_search_context: Arc, ) { @@ -2108,7 +2160,7 @@ async fn run_local_search_tasks( leaf_search_single_split_wrapper( simplified_search_request, leaf_search_context.clone(), - index_storage.clone(), + storage.clone(), split.clone(), leaf_split_search_permit, ) @@ -2263,10 +2315,10 @@ struct LeafSearchContext { split_filter: Arc>, } -async fn leaf_search_single_split_wrapper( +async fn leaf_search_single_split_wrapper( request: SearchRequest, ctx: Arc, - index_storage: Arc, + storage: T, split: SplitIdAndFooterOffsets, mut search_permit: SearchPermit, ) { @@ -2275,7 +2327,7 @@ async fn leaf_search_single_split_wrapper( leaf_search_single_split( request, ctx.clone(), - index_storage, + storage, split.clone(), &mut search_permit, ) diff --git a/quickwit/quickwit-search/src/lib.rs b/quickwit/quickwit-search/src/lib.rs index 2d891dbfa65..f386cb3a0b7 100644 --- a/quickwit/quickwit-search/src/lib.rs +++ b/quickwit/quickwit-search/src/lib.rs @@ -13,6 +13,7 @@ // limitations under the License. //! This projects implements quickwit's search API. +#![recursion_limit = "256"] #![warn(missing_docs)] #![allow(clippy::bool_assert_comparison)] #![deny(clippy::disallowed_methods)] diff --git a/quickwit/quickwit-search/src/list_fields/leaf.rs b/quickwit/quickwit-search/src/list_fields/leaf.rs index cd416037b14..8e8058720f2 100644 --- a/quickwit/quickwit-search/src/list_fields/leaf.rs +++ b/quickwit/quickwit-search/src/list_fields/leaf.rs @@ -22,7 +22,7 @@ use quickwit_proto::search::{ ListFieldsEntry, ListFieldsMetadata, ListFieldsResponse, SplitIdAndFooterOffsets, }; use quickwit_proto::types::IndexId; -use quickwit_storage::{OwnedBytes, Storage}; +use quickwit_storage::{OwnedBytes, StorageGetSlice}; use tracing::{Span, instrument}; use crate::leaf::open_split_bundle; @@ -41,12 +41,12 @@ const MAX_CONCURRENT_DOWNLOADS: usize = 500; /// /// Returns field metadata from the assigned splits. #[instrument(skip_all, fields(index_id = %index_id, num_splits = split_footers.len()))] -pub async fn leaf_list_fields( +pub async fn leaf_list_fields( index_id: IndexId, field_patterns_strs: &[String], split_footers: Vec, searcher_ctx: Arc, - storage: Arc, + storage: T, ) -> crate::Result { if split_footers.is_empty() { return Ok(ListFieldsResponse::default()); @@ -73,12 +73,12 @@ pub async fn leaf_list_fields( /// Two-stage pipeline: download blobs with bounded I/O concurrency, then hand each one off to the /// CPU pool as soon as it's ready. The CPU stage is left unbounded here because `run_cpu_intensive` /// already serializes work through the search thread pool. -async fn get_and_process_fields_metadata( +async fn get_and_process_fields_metadata( index_id: &str, field_patterns: &FieldPatterns, split_footers: Vec, searcher_ctx: Arc, - storage: Arc, + storage: T, ) -> anyhow::Result>> { stream::iter(split_footers) .map(|split_footer| { @@ -112,10 +112,10 @@ async fn get_and_process_fields_metadata( /// on miss, from storage (populating the cache as a side effect). The cache stores the bytes /// exactly as fetched (proto + zstd). #[instrument(skip_all, fields(split_id = %split_footer.split_id))] -async fn get_fields_metadata( +async fn get_fields_metadata( split_footer: &SplitIdAndFooterOffsets, searcher_ctx: &SearcherContext, - storage: Arc, + storage: T, ) -> anyhow::Result { if let Some(serialized_entries) = searcher_ctx.list_fields_cache.get(&split_footer.split_id) { return Ok(serialized_entries); @@ -131,9 +131,9 @@ async fn get_fields_metadata( /// Downloads the serialized fields metadata blob for one split. #[instrument(skip_all, fields(split_id = %split_footer.split_id))] -async fn fetch_fields_metadata( +async fn fetch_fields_metadata( searcher_ctx: &SearcherContext, - storage: Arc, + storage: T, split_footer: &SplitIdAndFooterOffsets, ) -> anyhow::Result { let (_, split_bundle) = open_split_bundle(searcher_ctx, storage, split_footer).await?; diff --git a/quickwit/quickwit-search/src/list_terms.rs b/quickwit/quickwit-search/src/list_terms.rs index 9c6e025ac92..46e6a7dc8a6 100644 --- a/quickwit/quickwit-search/src/list_terms.rs +++ b/quickwit/quickwit-search/src/list_terms.rs @@ -29,7 +29,7 @@ use quickwit_proto::search::{ SplitIdAndFooterOffsets, SplitSearchError, }; use quickwit_proto::types::IndexUid; -use quickwit_storage::{ByteRangeCache, Storage}; +use quickwit_storage::{ByteRangeCache, StorageGetSlice}; use tantivy::schema::{Field, FieldType}; use tantivy::{ReloadPolicy, Term}; use tracing::{debug, error, info, instrument}; @@ -208,10 +208,10 @@ pub fn jobs_to_leaf_requests( /// Apply a leaf list terms on a single split. #[instrument(skip_all, fields(split_id = split.split_id))] #[allow(deprecated)] -async fn leaf_list_terms_single_split( +async fn leaf_list_terms_single_split( searcher_context: &SearcherContext, search_request: &ListTermsRequest, - storage: Arc, + storage: T, split: SplitIdAndFooterOffsets, ) -> crate::Result { let cache = ByteRangeCache::with_infinite_capacity(); @@ -319,10 +319,10 @@ fn term_to_data(field: Field, field_type: &FieldType, field_value: &[u8]) -> Vec /// `leaf` step of list terms. #[instrument(skip_all)] -pub async fn leaf_list_terms( +pub async fn leaf_list_terms( searcher_context: Arc, request: &ListTermsRequest, - index_storage: Arc, + storage: T, splits: &[SplitIdAndFooterOffsets], ) -> Result { info!(split_offsets = ?PrettySample::new(splits, 5)); @@ -353,7 +353,7 @@ pub async fn leaf_list_terms( .iter() .zip(permits.into_iter()) .map(|(split, search_permit_recv)| { - let index_storage_clone = index_storage.clone(); + let storage_clone = storage.clone(); let searcher_context_clone = searcher_context.clone(); async move { let leaf_split_search_permit = search_permit_recv.await; @@ -363,7 +363,7 @@ pub async fn leaf_list_terms( let leaf_search_single_split_res = leaf_list_terms_single_split( &searcher_context_clone, request, - index_storage_clone, + storage_clone, split.clone(), ) .await; diff --git a/quickwit/quickwit-storage/src/bundle_storage.rs b/quickwit/quickwit-storage/src/bundle_storage.rs index 99a17d47a3f..691e1b842bb 100644 --- a/quickwit/quickwit-storage/src/bundle_storage.rs +++ b/quickwit/quickwit-storage/src/bundle_storage.rs @@ -33,40 +33,44 @@ use tracing::error; use crate::storage::SendableAsync; use crate::{ - BulkDeleteError, OwnedBytes, Storage, StorageError, StorageResult, VersionedComponent, + BulkDeleteError, OwnedBytes, Storage, StorageError, StorageGetSlice, StorageResult, + VersionedComponent, }; /// BundleStorage bundles together multiple files into a single file. -pub struct BundleStorage { - storage: Arc, +#[derive(Clone)] +pub struct BundleStorage { + storage: T, /// The file path of the bundle in the storage. bundle_filepath: PathBuf, file_ranges: BundleFileRanges, } -impl BundleStorage { +impl BundleStorage { /// Opens a split bundle by locating its footer directly from object storage. /// /// New splits expose the footer start in a fixed-size trailer. Legacy splits are supported by /// walking backward through the hotcache and bundle-metadata length fields. pub async fn open_from_storage( - storage: Arc, + storage: T, bundle_filepath: PathBuf, ) -> anyhow::Result<(Self, OwnedBytes, Range)> { let split_len = storage.file_num_bytes(&bundle_filepath).await?; let split_footer_range = - locate_split_footer_range(storage.as_ref(), &bundle_filepath, split_len).await?; + locate_split_footer_range(&storage, &bundle_filepath, split_len).await?; let split_footer_start = usize::try_from(split_footer_range.start)?; let split_footer_end = usize::try_from(split_footer_range.end)?; let split_footer_range_usize = split_footer_start..split_footer_end; let split_footer_bytes = storage - .get_slice(&bundle_filepath, split_footer_range_usize) + .get_slice_unboxed(&bundle_filepath, split_footer_range_usize) .await?; let (bundle_storage, hotcache) = Self::open_from_split_bytes(storage, bundle_filepath, split_footer_bytes)?; Ok((bundle_storage, hotcache, split_footer_range)) } +} +impl BundleStorage { /// Opens a BundleStorage. /// /// The provided bytes must include the footer bytes at the end of the slice, but they can have @@ -74,7 +78,7 @@ impl BundleStorage { /// /// Returns (Self, Hotcache) pub fn open_from_split_bytes( - storage: Arc, + storage: T, bundle_filepath: PathBuf, split_bytes: OwnedBytes, ) -> anyhow::Result<(Self, OwnedBytes)> { @@ -136,8 +140,8 @@ fn deserialize_split_footer_trailer(trailer: &[u8]) -> anyhow::Result( + storage: &S, split_path: &Path, split_len: u64, ) -> anyhow::Result> { @@ -288,7 +292,7 @@ impl BundleFileRanges { } #[async_trait] -impl Storage for BundleStorage { +impl Storage for BundleStorage { async fn check_connectivity(&self) -> anyhow::Result<()> { if !self .storage @@ -329,15 +333,7 @@ impl Storage for BundleStorage { path: &Path, range: Range, ) -> crate::StorageResult { - let file_range = self.file_ranges.get(path).ok_or_else(|| { - crate::StorageErrorKind::NotFound - .with_error(anyhow::anyhow!("missing file `{}`", path.display())) - })?; - let new_range = - file_range.start as usize + range.start..file_range.start as usize + range.end; - self.storage - .get_slice(&self.bundle_filepath, new_range) - .await + self.get_slice_unboxed(path, range).await } async fn get_slice_stream( @@ -390,13 +386,31 @@ impl Storage for BundleStorage { } } -impl HasLen for BundleStorage { +impl StorageGetSlice for BundleStorage { + async fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> StorageResult { + let file_range = self.file_ranges.get(path).ok_or_else(|| { + crate::StorageErrorKind::NotFound + .with_error(anyhow::anyhow!("missing file `{}`", path.display())) + })?; + let translated_range = + file_range.start as usize + range.start..file_range.start as usize + range.end; + self.storage + .get_slice_unboxed(&self.bundle_filepath, translated_range) + .await + } +} + +impl HasLen for BundleStorage { fn len(&self) -> usize { unimplemented!() } } -impl fmt::Debug for BundleStorage { +impl fmt::Debug for BundleStorage { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!( f, diff --git a/quickwit/quickwit-storage/src/cache/mod.rs b/quickwit/quickwit-storage/src/cache/mod.rs index 6a2bf38abf8..6453e97aeaa 100644 --- a/quickwit/quickwit-storage/src/cache/mod.rs +++ b/quickwit/quickwit-storage/src/cache/mod.rs @@ -30,7 +30,7 @@ pub use storage_with_cache::StorageWithCache; pub use self::byte_range_cache::{ByteRangeCache, FileByteRangeCache}; pub use self::memory_sized_cache::MemorySizedCache; -use crate::{OwnedBytes, Storage}; +use crate::{OwnedBytes, StorageGetSlice}; /// Wraps the given directory with a slice cache that is actually global /// to quickwit. @@ -39,14 +39,14 @@ use crate::{OwnedBytes, Storage}; /// - it uses a global /// - it relies on the idea that all of the files we attempt to cache have universally unique names. /// It happens to be true today, but this might be very error prone in the future. -pub fn wrap_storage_with_cache( +pub fn wrap_storage_with_cache( long_term_cache: Arc, - storage: Arc, -) -> Arc { - Arc::new(StorageWithCache { + storage: T, +) -> StorageWithCache { + StorageWithCache { storage, cache: long_term_cache, - }) + } } /// The `StorageCache` trait is the abstraction used to describe the caching logic diff --git a/quickwit/quickwit-storage/src/cache/storage_with_cache.rs b/quickwit/quickwit-storage/src/cache/storage_with_cache.rs index bc9d4064deb..cf41625ac3b 100644 --- a/quickwit/quickwit-storage/src/cache/storage_with_cache.rs +++ b/quickwit/quickwit-storage/src/cache/storage_with_cache.rs @@ -23,22 +23,27 @@ use tokio::io::AsyncRead; use crate::cache::StorageCache; use crate::storage::SendableAsync; -use crate::{BulkDeleteError, ListObjectsStream, OwnedBytes, Storage, StorageResult}; +use crate::{ + BulkDeleteError, ListObjectsStream, OwnedBytes, Storage, StorageGetSlice, StorageResult, +}; /// Use with care, StorageWithCache is read-only. -pub struct StorageWithCache { - pub storage: Arc, +#[derive(Clone)] +pub struct StorageWithCache { + /// Wrapped storage. + pub storage: T, + /// Slice cache consulted before the wrapped storage. pub cache: Arc, } -impl fmt::Debug for StorageWithCache { +impl fmt::Debug for StorageWithCache { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("StorageWithCache").finish() } } #[async_trait] -impl Storage for StorageWithCache { +impl Storage for StorageWithCache { async fn check_connectivity(&self) -> anyhow::Result<()> { self.storage.check_connectivity().await } @@ -56,15 +61,7 @@ impl Storage for StorageWithCache { } async fn get_slice(&self, path: &Path, byte_range: Range) -> StorageResult { - if let Some(bytes) = self.cache.get(path, byte_range.clone()).await { - Ok(bytes) - } else { - let bytes = self.storage.get_slice(path, byte_range.clone()).await?; - self.cache - .put(path.to_owned(), byte_range, bytes.clone()) - .await; - Ok(bytes) - } + self.get_slice_unboxed(path, byte_range).await } async fn get_slice_stream( @@ -113,6 +110,26 @@ impl Storage for StorageWithCache { } } +impl StorageGetSlice for StorageWithCache { + async fn get_slice_unboxed( + &self, + path: &Path, + byte_range: Range, + ) -> StorageResult { + if let Some(bytes) = self.cache.get(path, byte_range.clone()).await { + return Ok(bytes); + } + let bytes = self + .storage + .get_slice_unboxed(path, byte_range.clone()) + .await?; + self.cache + .put(path.to_owned(), byte_range, bytes.clone()) + .await; + Ok(bytes) + } +} + #[cfg(test)] mod tests { use std::collections::HashMap; diff --git a/quickwit/quickwit-storage/src/counting_storage.rs b/quickwit/quickwit-storage/src/counting_storage.rs index 14aba79573b..fe880ab5f24 100644 --- a/quickwit/quickwit-storage/src/counting_storage.rs +++ b/quickwit/quickwit-storage/src/counting_storage.rs @@ -23,7 +23,9 @@ use tantivy::directory::OwnedBytes; use tokio::io::AsyncRead; use crate::storage::SendableAsync; -use crate::{BulkDeleteError, ListObjectsStream, PutPayload, Storage, StorageResult}; +use crate::{ + BulkDeleteError, ListObjectsStream, PutPayload, Storage, StorageGetSlice, StorageResult, +}; /// Per-request download counters tracked by [`CountingStorage`]. /// @@ -56,41 +58,39 @@ impl DownloadCounters { /// /// Wrap a base `Storage` with this proxy at the entry point of a request to /// observe the per-request download volume. Cached layers (split cache, footer -/// cache, hotcache, byte-range cache) live BELOW this wrapper, so reads served +/// cache, hotcache, byte-range cache) live above this wrapper, so reads served /// from cache do NOT contribute to the counters — that is the desired behavior /// for a "downloaded from object storage" measurement. /// /// Write methods (`put`, `delete`, `bulk_delete`) and metadata methods /// (`exists`, `file_num_bytes`, `check_connectivity`) are passed through /// without counting; we only record reads that materialise data. -#[derive(Debug)] -pub struct CountingStorage { - inner: Arc, +#[derive(Clone, Debug)] +pub struct CountingStorage { + storage: T, counters: Arc, } -impl CountingStorage { +impl CountingStorage { /// Wrap a storage object to count for download request and downloaded bytes. - pub fn instrument_storage( - inner: Arc, - ) -> (Arc, Arc) { + pub fn instrument_storage(storage: T) -> (Self, Arc) { let counters = Arc::new(DownloadCounters::default()); let instrumented_storage = Self { - inner, + storage, counters: counters.clone(), }; - (Arc::new(instrumented_storage), counters) + (instrumented_storage, counters) } } #[async_trait] -impl Storage for CountingStorage { +impl Storage for CountingStorage { async fn check_connectivity(&self) -> anyhow::Result<()> { - self.inner.check_connectivity().await + self.storage.check_connectivity().await } async fn put(&self, path: &Path, payload: Box) -> StorageResult<()> { - self.inner.put(path, payload).await + self.storage.put(path, payload).await } async fn copy_to(&self, path: &Path, output: &mut dyn SendableAsync) -> StorageResult<()> { @@ -99,19 +99,17 @@ impl Storage for CountingStorage { // hot path of leaf search (only `get_slice` is), so this approximation // is acceptable. self.counters.requests.fetch_add(1, Ordering::Relaxed); - self.inner.copy_to(path, output).await + self.storage.copy_to(path, output).await } async fn copy_to_file(&self, path: &Path, output_path: &Path) -> StorageResult { - let num_bytes = self.inner.copy_to_file(path, output_path).await?; + let num_bytes = self.storage.copy_to_file(path, output_path).await?; self.counters.record_read(num_bytes); Ok(num_bytes) } async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { - let bytes = self.inner.get_slice(path, range).await?; - self.counters.record_read(bytes.len() as u64); - Ok(bytes) + self.get_slice_unboxed(path, range).await } async fn get_slice_stream( @@ -123,39 +121,51 @@ impl Storage for CountingStorage { // The stream may yield fewer bytes if the caller drops it early, but // that is rare and the over-count is bounded by the requested range. let range_len = range.len() as u64; - let stream = self.inner.get_slice_stream(path, range).await?; + let stream = self.storage.get_slice_stream(path, range).await?; self.counters.record_read(range_len); Ok(stream) } async fn get_all(&self, path: &Path) -> StorageResult { - let bytes = self.inner.get_all(path).await?; + let bytes = self.storage.get_all(path).await?; self.counters.record_read(bytes.len() as u64); Ok(bytes) } async fn delete(&self, path: &Path) -> StorageResult<()> { - self.inner.delete(path).await + self.storage.delete(path).await } async fn bulk_delete<'a>(&self, paths: &[&'a Path]) -> Result<(), BulkDeleteError> { - self.inner.bulk_delete(paths).await + self.storage.bulk_delete(paths).await } fn list(&self, prefix: &Path) -> ListObjectsStream { - self.inner.list(prefix) + self.storage.list(prefix) } async fn exists(&self, path: &Path) -> StorageResult { - self.inner.exists(path).await + self.storage.exists(path).await } async fn file_num_bytes(&self, path: &Path) -> StorageResult { - self.inner.file_num_bytes(path).await + self.storage.file_num_bytes(path).await } fn uri(&self) -> &Uri { - self.inner.uri() + self.storage.uri() + } +} + +impl StorageGetSlice for CountingStorage { + async fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> StorageResult { + let bytes = self.storage.get_slice_unboxed(path, range).await?; + self.counters.record_read(bytes.len() as u64); + Ok(bytes) } } @@ -166,10 +176,10 @@ mod tests { #[tokio::test] async fn test_counting_storage_counts_get_slice() { - let inner = RamStorageBuilder::default() + let ram_storage = RamStorageBuilder::default() .put("foo", b"hello world") .build(); - let (storage, counters) = CountingStorage::instrument_storage(Arc::new(inner)); + let (storage, counters) = CountingStorage::instrument_storage(Arc::new(ram_storage)); let bytes = storage.get_slice(Path::new("foo"), 0..5).await.unwrap(); assert_eq!(bytes.as_slice(), b"hello"); @@ -184,8 +194,8 @@ mod tests { #[tokio::test] async fn test_counting_storage_counts_get_all() { - let inner = RamStorageBuilder::default().put("foo", b"hello").build(); - let (storage, counters) = CountingStorage::instrument_storage(Arc::new(inner)); + let ram_storage = RamStorageBuilder::default().put("foo", b"hello").build(); + let (storage, counters) = CountingStorage::instrument_storage(Arc::new(ram_storage)); let bytes = storage.get_all(Path::new("foo")).await.unwrap(); assert_eq!(bytes.as_slice(), b"hello"); @@ -197,8 +207,8 @@ mod tests { #[tokio::test] async fn test_counting_storage_does_not_count_metadata() { - let inner = RamStorageBuilder::default().put("foo", b"hello").build(); - let (storage, counters) = CountingStorage::instrument_storage(Arc::new(inner)); + let ram_storage = RamStorageBuilder::default().put("foo", b"hello").build(); + let (storage, counters) = CountingStorage::instrument_storage(Arc::new(ram_storage)); assert!(storage.exists(Path::new("foo")).await.unwrap()); assert_eq!(storage.file_num_bytes(Path::new("foo")).await.unwrap(), 5); diff --git a/quickwit/quickwit-storage/src/debouncer.rs b/quickwit/quickwit-storage/src/debouncer.rs index 04a4338c037..7d72de25c58 100644 --- a/quickwit/quickwit-storage/src/debouncer.rs +++ b/quickwit/quickwit-storage/src/debouncer.rs @@ -13,32 +13,30 @@ // limitations under the License. use std::fmt; +use std::future::Future; use std::hash::Hash; use std::ops::Range; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, Weak}; use ahash::HashMap; use async_trait::async_trait; -use futures::future::{BoxFuture, Shared, WeakShared}; -use futures::{Future, FutureExt}; use quickwit_common::uri::Uri; use tantivy::directory::OwnedBytes; use tokio::io::AsyncRead; +use tokio::sync::OnceCell; use crate::storage::SendableAsync; -use crate::{BulkDeleteError, ListObjectsStream, Storage, StorageResult}; +use crate::{BulkDeleteError, ListObjectsStream, Storage, StorageGetSlice, StorageResult}; /// The AsyncDebouncer debounces inflight Futures, so that concurrent async request to the same data /// source can be deduplicated. /// -/// Since we pass the Future potentially to multiple consumer, everything needs to be cloneable. The -/// data and the future. This is reflected on the generic type bounds for the value V: Clone. -/// -/// Since most Futures return an Result, this also encompasses the error. +/// Concurrent consumers receive clones of the initialized value. Since most futures return a +/// `Result`, the clone requirement also applies to the error. pub struct AsyncDebouncer { - cache: Mutex>>>, + cache: Mutex>>>, /// Number of inserts performed since the last full garbage-collection scan. /// /// Used to amortize the cost of reclaiming entries left behind by cancelled futures (see @@ -62,7 +60,7 @@ impl Default for AsyncDebouncer { /// so we only need a periodic sweep to reclaim entries whose key is never accessed again. const CLEANUP_INTERVAL: usize = 1_024; -impl AsyncDebouncer { +impl AsyncDebouncer { /// Returns the number of entries in the debouncing cache. /// /// This is always greater than the number of inflight futures, and smaller @@ -72,57 +70,41 @@ impl AsyncDebouncer { } /// Returns the inflight future for `key`, deduplicating concurrent calls: if a future for - /// `key` is already inflight, all callers await that same shared future; otherwise - /// `build_a_future_fast` is invoked to create one. - /// - /// # Hidden contract + /// `key` is already inflight, all callers await that same operation; otherwise + /// `build_future` is invoked to create one. /// - /// `build_a_future_fast` is invoked **while the internal cache lock is held**. It must - /// therefore: - /// - be cheap — it only *constructs* the future, it must not perform blocking work (the future - /// itself is awaited later, outside the lock); - /// - never re-enter this `AsyncDebouncer` (e.g. call `get_or_create` on the same instance), - /// which would deadlock, since the cache lock is a non-reentrant [`std::sync::Mutex`]. - /// - /// Holding the lock across both the lookup and the insert is deliberate: it makes the - /// lookup-then-insert atomic, so two concurrent callers for the same key cannot both build - /// and race to insert. The lock is always released before the future is awaited. - fn get_or_create(&self, key: K, build_a_future_fast: T) -> Shared> + /// The lookup and insertion happen synchronously while the cache lock is held, but the future + /// builder is invoked later, outside the lock. [`OnceCell`] ensures only one caller initializes + /// a key. If that caller is cancelled, another waiter can take over initialization. + fn get_or_create<'a, T, F>( + &'a self, + key: K, + build_future: T, + ) -> impl Future + Send + 'a where - T: FnOnce() -> F, - F: Future + Send + 'static, + T: FnOnce() -> F + Send + 'a, + F: Future + Send + 'a, { - // `created` distinguishes the caller that actually built the entry (and is therefore - // responsible for removing it once the future resolves) from callers that merely joined - // an already-inflight future. let mut debouncing_cache_guard = self.cache.lock().unwrap(); - // A stale entry (left behind by a cancelled future) fails to upgrade and is simply - // overwritten by the insert below. - if let Some(fut) = debouncing_cache_guard - .get(&key) - .and_then(WeakShared::upgrade) - { - return fut; - } - - // Note that we are call build_a_future_fast WHILE holding the cache lock. - // For this reason, it is crucial to only call this function with a cheap and simple future - // builder. - let fut = Box::pin(build_a_future_fast()) as BoxFuture<'static, V>; - let fut = fut.shared(); - let weak_fut = fut.clone().downgrade().unwrap(); - debouncing_cache_guard.insert(key.clone(), weak_fut); - - // Amortized garbage collection. Running a full scan on every call would make - // each call O(n); instead we scan once every `CLEANUP_INTERVAL` inserts. - let num_inserts = self.inserts_since_cleanup.fetch_add(1, Ordering::Relaxed); - if num_inserts >= CLEANUP_INTERVAL { - self.inserts_since_cleanup.store(0, Ordering::Relaxed); - debouncing_cache_guard.retain(|_, weak_future| weak_future.upgrade().is_some()); - } + let cell = if let Some(cell) = debouncing_cache_guard.get(&key).and_then(Weak::upgrade) { + cell + } else { + let cell = Arc::new(OnceCell::new()); + debouncing_cache_guard.insert(key, Arc::downgrade(&cell)); + + // Amortized garbage collection. Running a full scan on every call would make each + // call O(n); instead we scan once every `CLEANUP_INTERVAL` inserts. + let num_inserts = self.inserts_since_cleanup.fetch_add(1, Ordering::Relaxed); + if num_inserts >= CLEANUP_INTERVAL { + self.inserts_since_cleanup.store(0, Ordering::Relaxed); + debouncing_cache_guard.retain(|_, weak_cell| weak_cell.upgrade().is_some()); + } + cell + }; + drop(debouncing_cache_guard); - fut + async move { cell.get_or_init(build_future).await.clone() } } } @@ -140,10 +122,8 @@ type DebouncerKey = (PathBuf, Range); /// /// ==> R3 would return the cached result, although the resource has been deleted. pub(crate) struct DebouncedStorage { - // wrap both in Arc, because the Future is stored in the cache, which has 'static lifetime - // associated - underlying: Arc, - slice_debouncer: Arc>>, + storage: T, + slice_debouncer: AsyncDebouncer>, } impl fmt::Debug for DebouncedStorage { @@ -152,19 +132,19 @@ impl fmt::Debug for DebouncedStorage { } } -impl DebouncedStorage { - pub(crate) fn new(underlying: T) -> Self { +impl DebouncedStorage { + pub(crate) fn new(storage: T) -> Self { Self { - underlying: Arc::new(underlying), - slice_debouncer: Arc::new(AsyncDebouncer::default()), + storage, + slice_debouncer: AsyncDebouncer::default(), } } } #[async_trait] -impl Storage for DebouncedStorage { +impl Storage for DebouncedStorage { async fn check_connectivity(&self) -> anyhow::Result<()> { - self.underlying.check_connectivity().await + self.storage.check_connectivity().await } async fn put( @@ -172,21 +152,15 @@ impl Storage for DebouncedStorage { path: &Path, payload: Box, ) -> crate::StorageResult<()> { - self.underlying.put(path, payload).await + self.storage.put(path, payload).await } async fn copy_to(&self, path: &Path, output: &mut dyn SendableAsync) -> StorageResult<()> { - self.underlying.copy_to(path, output).await + self.storage.copy_to(path, output).await } async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { - let (debouncer, underlying) = (self.slice_debouncer.clone(), self.underlying.clone()); - let key = (path.to_owned(), range); - debouncer - .get_or_create(key.clone(), || async move { - underlying.get_slice(&key.0, key.1).await - }) - .await + self.get_slice_unboxed(path, range).await } async fn get_slice_stream( @@ -195,38 +169,49 @@ impl Storage for DebouncedStorage { range: Range, ) -> StorageResult> { // Getting a stream bypasses the debouncer - self.underlying.get_slice_stream(path, range).await + self.storage.get_slice_stream(path, range).await } async fn delete(&self, path: &Path) -> StorageResult<()> { - self.underlying.delete(path).await + self.storage.delete(path).await } async fn bulk_delete<'a>(&self, paths: &[&'a Path]) -> Result<(), BulkDeleteError> { - self.underlying.bulk_delete(paths).await + self.storage.bulk_delete(paths).await } fn list(&self, prefix: &Path) -> ListObjectsStream { - self.underlying.list(prefix) + self.storage.list(prefix) } async fn get_all(&self, path: &Path) -> StorageResult { - let (debouncer, underlying) = (self.slice_debouncer.clone(), self.underlying.clone()); let key = (path.to_owned(), 0..usize::MAX); - debouncer - .get_or_create( - key.clone(), - || async move { underlying.get_all(&key.0).await }, - ) + self.slice_debouncer + .get_or_create(key.clone(), || self.storage.get_all(&key.0)) .await } fn uri(&self) -> &Uri { - self.underlying.uri() + self.storage.uri() } async fn file_num_bytes(&self, path: &Path) -> StorageResult { - self.underlying.file_num_bytes(path).await + self.storage.file_num_bytes(path).await + } +} + +impl StorageGetSlice for DebouncedStorage { + async fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> StorageResult { + let key = (path.to_owned(), range); + self.slice_debouncer + .get_or_create(key.clone(), || { + self.storage.get_slice_unboxed(&key.0, key.1) + }) + .await } } @@ -290,7 +275,7 @@ mod tests { let test_filepath1 = test_filepath1.clone(); async move { COUNT.fetch_add(1, Ordering::SeqCst); - let contents = Box::pin(fs::read_to_string(test_filepath1.as_ref().clone())) + let contents = fs::read_to_string(test_filepath1.as_ref().clone()) .await // to string, so that the error is cloneable .map_err(|err| err.to_string())?; @@ -413,8 +398,7 @@ mod tests { async move { debouncer_clone.get_or_create("key0".to_owned(), load).await }, ); tokio::time::sleep(Duration::from_millis(10)).await; - // This will cause the Future to be cancelled, so it will not be polled anymore. - // That also means the remove in the cache is not called, which is awaiting the future + // This cancels initialization and leaves a stale weak entry until the next lookup. handle.abort(); tokio::time::sleep(Duration::from_secs(1)).await; @@ -439,9 +423,47 @@ mod tests { } } + #[tokio::test] + async fn test_waiter_takes_over_cancelled_initialization() { + let debouncer = Arc::new(AsyncDebouncer::default()); + let initialization_started = Arc::new(tokio::sync::Notify::new()); + + let (first_constructed_tx, first_constructed_rx) = tokio::sync::oneshot::channel(); + let first_debouncer = debouncer.clone(); + let first_started = initialization_started.clone(); + let first_handle = task::spawn(async move { + let future = first_debouncer.get_or_create("key".to_string(), || async move { + first_started.notify_one(); + std::future::pending::().await + }); + let _ = first_constructed_tx.send(()); + future.await + }); + first_constructed_rx.await.unwrap(); + initialization_started.notified().await; + + let (waiter_constructed_tx, waiter_constructed_rx) = tokio::sync::oneshot::channel(); + let waiter_debouncer = debouncer.clone(); + let waiter_handle = task::spawn(async move { + let future = waiter_debouncer + .get_or_create("key".to_string(), || async { "completed".to_string() }); + let _ = waiter_constructed_tx.send(()); + future.await + }); + waiter_constructed_rx.await.unwrap(); + + first_handle.abort(); + assert!(first_handle.await.unwrap_err().is_cancelled()); + let value = tokio::time::timeout(Duration::from_secs(1), waiter_handle) + .await + .unwrap() + .unwrap(); + assert_eq!(value, "completed"); + } + async fn load_via_fn(path: PathBuf, cnt: &AtomicU32) -> Result { cnt.fetch_add(1, Ordering::SeqCst); - let contents = Box::pin(fs::read_to_string(path)) + let contents = fs::read_to_string(path) .await .map_err(|err| err.to_string())?; // sleep so the requests can be reproducible debounced diff --git a/quickwit/quickwit-storage/src/lib.rs b/quickwit/quickwit-storage/src/lib.rs index 73ef173eee2..b53b5a38ea6 100644 --- a/quickwit/quickwit-storage/src/lib.rs +++ b/quickwit/quickwit-storage/src/lib.rs @@ -36,7 +36,7 @@ pub use debouncer::AsyncDebouncer; pub(crate) use debouncer::DebouncedStorage; pub use self::payload::PutPayload; -pub use self::storage::{ListObjectsStream, ObjectMetadata, Storage}; +pub use self::storage::{ListObjectsStream, ObjectMetadata, Storage, StorageGetSlice}; mod bundle_storage; mod error; @@ -67,7 +67,7 @@ pub use self::bundle_storage::{ pub use self::cache::MockStorageCache; pub use self::cache::{ ByteRangeCache, FileByteRangeCache, MemorySizedCache, QuickwitCache, StorageCache, - wrap_storage_with_cache, + StorageWithCache, wrap_storage_with_cache, }; pub use self::counting_storage::{CountingStorage, DownloadCounters}; pub use self::local_file_storage::{LocalFileStorage, LocalFileStorageFactory}; @@ -84,9 +84,7 @@ pub use self::ram_storage::{RamStorage, RamStorageBuilder}; pub use self::split::{SplitPayload, SplitPayloadBuilder}; #[cfg(any(test, feature = "testsuite"))] pub use self::storage::MockStorage; -#[cfg(any(test, feature = "testsuite"))] -pub use self::storage_factory::MockStorageFactory; -pub use self::storage_factory::{StorageFactory, UnsupportedStorage}; +pub use self::storage_factory::{ResolvedStorage, StorageFactory, UnsupportedStorage}; pub use self::storage_resolver::StorageResolver; #[cfg(feature = "integration-testsuite")] pub use self::test_suite::{ diff --git a/quickwit/quickwit-storage/src/local_file_storage.rs b/quickwit/quickwit-storage/src/local_file_storage.rs index c0e36b5ba4f..2e6b9e9cd6d 100644 --- a/quickwit/quickwit-storage/src/local_file_storage.rs +++ b/quickwit/quickwit-storage/src/local_file_storage.rs @@ -14,17 +14,16 @@ use std::collections::{BTreeSet, HashMap}; use std::fmt; +use std::future::Future; use std::io::{ErrorKind, SeekFrom}; use std::ops::Range; use std::path::{Component, Path, PathBuf}; -use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt; use futures::future::{BoxFuture, FutureExt}; use quickwit_common::ignore_error_kind; use quickwit_common::uri::Uri; -use quickwit_config::StorageBackend; use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; use tracing::warn; @@ -32,7 +31,7 @@ use crate::metrics::object_storage_get_slice_in_flight_guards; use crate::storage::SendableAsync; use crate::{ BulkDeleteError, DebouncedStorage, DeleteFailure, OwnedBytes, Storage, StorageError, - StorageErrorKind, StorageFactory, StorageResolverError, StorageResult, + StorageErrorKind, StorageGetSlice, StorageResolverError, StorageResult, }; /// File system compatible storage implementation. @@ -91,6 +90,30 @@ impl LocalFileStorage { ignore_error_kind!(ErrorKind::NotFound, tokio::fs::remove_file(full_path).await)?; Ok(()) } + + #[tracing::instrument(name = "storage.local_file.get_slice", skip(self), level = "debug")] + async fn get_slice_impl(&self, path: &Path, range: Range) -> StorageResult { + let full_path = self.full_path(path)?; + tokio::task::spawn_blocking(move || { + use std::io::{Read, Seek}; + // We run these IO operations in a blocking task so there is no scheduling delay + // between them, as there would be when using a Tokio file. + let mut file = std::fs::File::open(full_path)?; + file.seek(SeekFrom::Start(range.start as u64))?; + let _in_flight_guards = object_storage_get_slice_in_flight_guards(range.len()); + let mut content_bytes: Vec = Vec::with_capacity(range.len()); + #[allow(clippy::uninit_vec)] + unsafe { + content_bytes.set_len(range.len()); + } + file.read_exact(&mut content_bytes)?; + Ok(OwnedBytes::new(content_bytes)) + }) + .await + .map_err(|_| { + StorageErrorKind::Internal.with_error(anyhow::anyhow!("reading file panicked")) + })? + } } /// Ensure that the path given does not include any ".." for security reasons. @@ -212,28 +235,8 @@ impl Storage for LocalFileStorage { Ok(()) } - #[tracing::instrument(name = "storage.local_file.get_slice", skip(self), level = "debug")] async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { - let full_path = self.full_path(path)?; - tokio::task::spawn_blocking(move || { - use std::io::{Read, Seek}; - // we run these io in a spawn_blocking so there is no scheduling delay between each - // step, as there would be if using tokio async File. - let mut file = std::fs::File::open(full_path)?; - file.seek(SeekFrom::Start(range.start as u64))?; - let _in_flight_guards = object_storage_get_slice_in_flight_guards(range.len()); - let mut content_bytes: Vec = Vec::with_capacity(range.len()); - #[allow(clippy::uninit_vec)] - unsafe { - content_bytes.set_len(range.len()); - } - file.read_exact(&mut content_bytes)?; - Ok(OwnedBytes::new(content_bytes)) - }) - .await - .map_err(|_| { - StorageErrorKind::Internal.with_error(anyhow::anyhow!("reading file panicked")) - })? + self.get_slice_impl(path, range).await } #[tracing::instrument( @@ -364,19 +367,27 @@ impl Storage for LocalFileStorage { } } +impl StorageGetSlice for LocalFileStorage { + fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> impl Future> + Send { + self.get_slice_impl(path, range) + } +} + /// A File storage resolver #[derive(Clone, Debug, Default)] pub struct LocalFileStorageFactory; -#[async_trait] -impl StorageFactory for LocalFileStorageFactory { - fn backend(&self) -> StorageBackend { - StorageBackend::File - } - - async fn resolve(&self, uri: &Uri) -> Result, StorageResolverError> { +impl LocalFileStorageFactory { + pub(crate) fn resolve( + &self, + uri: &Uri, + ) -> Result, StorageResolverError> { let storage = LocalFileStorage::from_uri(uri)?; - Ok(Arc::new(DebouncedStorage::new(storage))) + Ok(DebouncedStorage::new(storage)) } } @@ -422,19 +433,17 @@ mod tests { let index_uri = Uri::from_str(&format!("file://{}/foo/bar", temp_dir.path().display())).unwrap(); let local_file_storage_factory = LocalFileStorageFactory; - let local_file_storage = local_file_storage_factory.resolve(&index_uri).await?; + let local_file_storage = local_file_storage_factory.resolve(&index_uri)?; assert_eq!(local_file_storage.uri(), &index_uri); let err = local_file_storage_factory .resolve(&Uri::for_test("s3://foo/bar")) - .await .err() .unwrap(); assert!(matches!(err, StorageResolverError::InvalidUri { .. })); let err = local_file_storage_factory .resolve(&Uri::for_test("s3://")) - .await .err() .unwrap(); assert!(matches!(err, StorageResolverError::InvalidUri { .. })); diff --git a/quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs b/quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs index cbbf6ea8ded..25bb86492f1 100644 --- a/quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs +++ b/quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs @@ -13,10 +13,11 @@ // limitations under the License. use std::collections::HashMap; +use std::future::Future; use std::num::NonZeroU32; use std::ops::Range; use std::path::{Path, PathBuf}; -use std::sync::{Arc, LazyLock}; +use std::sync::LazyLock; use std::{fmt, io}; use async_trait::async_trait; @@ -34,7 +35,7 @@ use md5::Digest; use quickwit_common::retry::{RetryParams, Retryable, retry}; use quickwit_common::uri::Uri; use quickwit_common::{chunk_range, ignore_error_kind, into_u64_range}; -use quickwit_config::{AzureStorageConfig, StorageBackend}; +use quickwit_config::AzureStorageConfig; use quickwit_metrics::HistogramTimer; use regex::Regex; use tantivy::directory::OwnedBytes; @@ -50,7 +51,7 @@ use crate::stable_deref_bytes::into_owned_bytes; use crate::storage::SendableAsync; use crate::{ BulkDeleteError, DeleteFailure, MultiPartPolicy, PutPayload, Storage, StorageError, - StorageErrorKind, StorageFactory, StorageResolverError, StorageResult, + StorageErrorKind, StorageGetSlice, StorageResolverError, StorageResult, }; /// Azure object storage resolver. @@ -65,15 +66,13 @@ impl AzureBlobStorageFactory { } } -#[async_trait] -impl StorageFactory for AzureBlobStorageFactory { - fn backend(&self) -> StorageBackend { - StorageBackend::Azure - } - - async fn resolve(&self, uri: &Uri) -> Result, StorageResolverError> { +impl AzureBlobStorageFactory { + pub(crate) fn resolve( + &self, + uri: &Uri, + ) -> Result, StorageResolverError> { let storage = AzureBlobStorage::from_uri(&self.storage_config, uri)?; - Ok(Arc::new(DebouncedStorage::new(storage))) + Ok(DebouncedStorage::new(storage)) } } @@ -338,6 +337,23 @@ impl AzureBlobStorage { } } +impl AzureBlobStorage { + #[instrument(name = "storage.azure.get_slice", level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))] + async fn get_slice_impl(&self, path: &Path, range: Range) -> StorageResult { + self.get_to_bytes(path, Some(range.clone())) + .await + .map(into_owned_bytes) + .map_err(|err| { + err.add_context(format!( + "failed to fetch slice {:?} for object: {}/{}", + range, + self.uri, + path.display(), + )) + }) + } +} + #[async_trait] impl Storage for AzureBlobStorage { async fn check_connectivity(&self) -> anyhow::Result<()> { @@ -455,19 +471,8 @@ impl Storage for AzureBlobStorage { } } - #[instrument(name = "storage.azure.get_slice", level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))] async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { - self.get_to_bytes(path, Some(range.clone())) - .await - .map(into_owned_bytes) - .map_err(|err| { - err.add_context(format!( - "failed to fetch slice {:?} for object: {}/{}", - range, - self.uri, - path.display(), - )) - }) + self.get_slice_impl(path, range).await } #[instrument(name = "storage.azure.get_slice_stream", level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))] @@ -547,6 +552,16 @@ impl Storage for AzureBlobStorage { } } +impl StorageGetSlice for AzureBlobStorage { + fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> impl Future> + Send { + self.get_slice_impl(path, range) + } +} + /// Copy range of payload into `Bytes` and return the computed md5. async fn extract_range_data_and_hash( payload: Box, diff --git a/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs b/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs index 3ec4afdef5b..0c9ed141908 100644 --- a/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs +++ b/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs @@ -14,6 +14,7 @@ use std::borrow::Cow; use std::collections::HashMap; +use std::future::Future; use std::ops::Range; use std::path::{Path, PathBuf}; use std::pin::Pin; @@ -64,7 +65,7 @@ use crate::stable_deref_bytes::into_owned_bytes; use crate::storage::SendableAsync; use crate::{ BulkDeleteError, DeleteFailure, ListObjectsStream, ObjectMetadata, OwnedBytes, Storage, - StorageError, StorageErrorKind, StorageResolverError, StorageResult, + StorageError, StorageErrorKind, StorageGetSlice, StorageResolverError, StorageResult, }; /// Semaphore to limit the number of concurrent requests to the object store. Some object stores @@ -888,6 +889,24 @@ fn byte_stream_to_storage_error(error: aws_sdk_s3::primitives::ByteStreamError) kind.with_error(error) } +impl S3CompatibleObjectStorage { + #[instrument(name = "storage.s3.get_slice", level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))] + async fn get_slice_impl(&self, path: &Path, range: Range) -> StorageResult { + let _permit = REQUEST_SEMAPHORE.acquire().await; + self.get_to_bytes(path, Some(range.clone())) + .await + .map(into_owned_bytes) + .map_err(|err| { + err.add_context(format!( + "failed to fetch slice {:?} for object: {}/{}", + range, + self.uri, + path.display(), + )) + }) + } +} + #[async_trait] impl Storage for S3CompatibleObjectStorage { async fn check_connectivity(&self) -> anyhow::Result<()> { @@ -1068,20 +1087,8 @@ impl Storage for S3CompatibleObjectStorage { .boxed() } - #[instrument(name = "storage.s3.get_slice", level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))] async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { - let _permit = REQUEST_SEMAPHORE.acquire().await; - self.get_to_bytes(path, Some(range.clone())) - .await - .map(into_owned_bytes) - .map_err(|err| { - err.add_context(format!( - "failed to fetch slice {:?} for object: {}/{}", - range, - self.uri, - path.display(), - )) - }) + self.get_slice_impl(path, range).await } #[instrument(name = "storage.s3.get_slice_stream", level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))] @@ -1147,6 +1154,16 @@ impl Storage for S3CompatibleObjectStorage { } } +impl StorageGetSlice for S3CompatibleObjectStorage { + fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> impl Future> + Send { + self.get_slice_impl(path, range) + } +} + /// Converts S3 object metadata to paths relative to the configured storage root. fn convert_list_objects( objects: &[aws_sdk_s3::types::Object], diff --git a/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage_resolver.rs b/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage_resolver.rs index 92b4406486b..3e1b5f57742 100644 --- a/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage_resolver.rs +++ b/quickwit/quickwit-storage/src/object_storage/s3_compatible_storage_resolver.rs @@ -12,18 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::sync::Arc; - -use async_trait::async_trait; use aws_sdk_s3::Client as S3Client; use quickwit_common::uri::Uri; -use quickwit_config::{S3StorageConfig, StorageBackend}; +use quickwit_config::S3StorageConfig; use tokio::sync::OnceCell; use super::s3_compatible_storage::create_s3_client; -use crate::{ - DebouncedStorage, S3CompatibleObjectStorage, Storage, StorageFactory, StorageResolverError, -}; +use crate::{DebouncedStorage, S3CompatibleObjectStorage, StorageResolverError}; /// S3 compatible object storage resolver. pub struct S3CompatibleObjectStorageFactory { @@ -46,13 +41,11 @@ impl S3CompatibleObjectStorageFactory { } } -#[async_trait] -impl StorageFactory for S3CompatibleObjectStorageFactory { - fn backend(&self) -> StorageBackend { - StorageBackend::S3 - } - - async fn resolve(&self, uri: &Uri) -> Result, StorageResolverError> { +impl S3CompatibleObjectStorageFactory { + pub(crate) async fn resolve( + &self, + uri: &Uri, + ) -> Result, StorageResolverError> { let s3_client = self .s3_client .get_or_init(|| create_s3_client(&self.storage_config)) @@ -61,6 +54,6 @@ impl StorageFactory for S3CompatibleObjectStorageFactory { let storage = S3CompatibleObjectStorage::from_uri_and_client(&self.storage_config, uri, s3_client) .await?; - Ok(Arc::new(DebouncedStorage::new(storage))) + Ok(DebouncedStorage::new(storage)) } } diff --git a/quickwit/quickwit-storage/src/opendal_storage/base.rs b/quickwit/quickwit-storage/src/opendal_storage/base.rs index a933c136025..8efbf3db8fb 100644 --- a/quickwit/quickwit-storage/src/opendal_storage/base.rs +++ b/quickwit/quickwit-storage/src/opendal_storage/base.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::future::Future; use std::ops::Range; use std::path::Path; use std::{fmt, io}; @@ -30,7 +31,7 @@ use crate::stable_deref_bytes::into_owned_bytes; use crate::storage::SendableAsync; use crate::{ BulkDeleteError, MultiPartPolicy, OwnedBytes, PutPayload, Storage, StorageError, - StorageErrorKind, StorageResolverError, StorageResult, + StorageErrorKind, StorageGetSlice, StorageResolverError, StorageResult, }; /// OpenDAL based storage implementation. @@ -123,6 +124,24 @@ where } } +impl OpendalStorage { + #[instrument(name = "storage.gcs.get_slice", level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))] + async fn get_slice_impl(&self, path: &Path, range: Range) -> StorageResult { + let path = path.as_os_str().to_string_lossy(); + let size = range.len(); + let range = range.start as u64..range.end as u64; + // Unlike other object store implementations, in-flight requests are recorded before + // issuing the query to the object store. + let _inflight_guards = object_storage_get_slice_in_flight_guards(size); + crate::metrics::OBJECT_STORAGE_GET_TOTAL.inc(); + let _timer = HistogramTimer::new(&crate::metrics::OBJECT_STORAGE_GET_OBJECT_DURATION); + // `Buffer::to_bytes` is zero-copy when the underlying buffer is contiguous, and coalesces + // into a single `Bytes` otherwise. + let storage_content = self.op.read_with(&path).range(range).await?.to_bytes(); + Ok(into_owned_bytes(storage_content)) + } +} + #[async_trait] impl Storage for OpendalStorage { async fn check_connectivity(&self) -> anyhow::Result<()> { @@ -167,21 +186,8 @@ impl Storage for OpendalStorage { Ok(()) } - #[instrument(name = "storage.gcs.get_slice", level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))] async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { - let path = path.as_os_str().to_string_lossy(); - let size = range.len(); - let range = range.start as u64..range.end as u64; - // Unlike other object store implementations, in flight requests are - // recorded before issuing the query to the object store. - let _inflight_guards = object_storage_get_slice_in_flight_guards(size); - crate::metrics::OBJECT_STORAGE_GET_TOTAL.inc(); - let _timer = HistogramTimer::new(&crate::metrics::OBJECT_STORAGE_GET_OBJECT_DURATION); - // `Buffer::to_bytes` is zero-copy when the underlying buffer is contiguous, and coalesces - // into a single `Bytes` otherwise — avoiding the extra `Vec` round-trip `to_vec` would - // perform. - let storage_content = self.op.read_with(&path).range(range).await?.to_bytes(); - Ok(into_owned_bytes(storage_content)) + self.get_slice_impl(path, range).await } #[instrument(name = "storage.gcs.get_slice_stream", level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))] @@ -288,6 +294,16 @@ impl Storage for OpendalStorage { } } +impl StorageGetSlice for OpendalStorage { + fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> impl Future> + Send { + self.get_slice_impl(path, range) + } +} + impl From for StorageError { fn from(err: opendal::Error) -> Self { match err.kind() { diff --git a/quickwit/quickwit-storage/src/opendal_storage/google_cloud_storage.rs b/quickwit/quickwit-storage/src/opendal_storage/google_cloud_storage.rs index ccf41e5adb2..b0a93cf3cc5 100644 --- a/quickwit/quickwit-storage/src/opendal_storage/google_cloud_storage.rs +++ b/quickwit/quickwit-storage/src/opendal_storage/google_cloud_storage.rs @@ -13,17 +13,16 @@ // limitations under the License. use std::path::PathBuf; -use std::sync::{Arc, LazyLock}; +use std::sync::LazyLock; -use async_trait::async_trait; use quickwit_common::uri::Uri; -use quickwit_config::{GoogleCloudStorageConfig, StorageBackend}; +use quickwit_config::GoogleCloudStorageConfig; use regex::Regex; use tracing::info; use super::OpendalStorage; +use crate::StorageResolverError; use crate::debouncer::DebouncedStorage; -use crate::{Storage, StorageFactory, StorageResolverError}; /// Google cloud storage resolver. pub struct GoogleCloudStorageFactory { @@ -37,15 +36,13 @@ impl GoogleCloudStorageFactory { } } -#[async_trait] -impl StorageFactory for GoogleCloudStorageFactory { - fn backend(&self) -> StorageBackend { - StorageBackend::Google - } - - async fn resolve(&self, uri: &Uri) -> Result, StorageResolverError> { +impl GoogleCloudStorageFactory { + pub(crate) fn resolve( + &self, + uri: &Uri, + ) -> Result, StorageResolverError> { let storage = from_uri(&self.storage_config, uri)?; - Ok(Arc::new(DebouncedStorage::new(storage))) + Ok(DebouncedStorage::new(storage)) } } diff --git a/quickwit/quickwit-storage/src/opendal_storage/mod.rs b/quickwit/quickwit-storage/src/opendal_storage/mod.rs index a7d7bb57983..654ed606a32 100644 --- a/quickwit/quickwit-storage/src/opendal_storage/mod.rs +++ b/quickwit/quickwit-storage/src/opendal_storage/mod.rs @@ -13,7 +13,7 @@ // limitations under the License. mod base; -use base::OpendalStorage; +pub(crate) use base::OpendalStorage; mod google_cloud_storage; diff --git a/quickwit/quickwit-storage/src/prefix_storage.rs b/quickwit/quickwit-storage/src/prefix_storage.rs index 3242f24ea49..585f2df879e 100644 --- a/quickwit/quickwit-storage/src/prefix_storage.rs +++ b/quickwit/quickwit-storage/src/prefix_storage.rs @@ -15,7 +15,6 @@ use std::fmt; use std::ops::Range; use std::path::{Path, PathBuf}; -use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt; @@ -25,18 +24,19 @@ use tokio::io::AsyncRead; use crate::storage::SendableAsync; use crate::{ BulkDeleteError, ListObjectsStream, ObjectMetadata, OwnedBytes, Storage, StorageErrorKind, - StorageResult, + StorageGetSlice, StorageResult, }; /// This storage acts as a proxy to another storage that simply modifies each API call /// by preceding each path with a given a prefix. -struct PrefixStorage { - pub storage: Arc, +#[derive(Clone)] +pub(crate) struct PrefixStorage { + pub storage: T, pub prefix: PathBuf, uri: Uri, } -impl fmt::Debug for PrefixStorage { +impl fmt::Debug for PrefixStorage { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("PrefixStorage") .field("uri", &self.uri) @@ -46,7 +46,7 @@ impl fmt::Debug for PrefixStorage { } #[async_trait] -impl Storage for PrefixStorage { +impl Storage for PrefixStorage { async fn check_connectivity(&self) -> anyhow::Result<()> { self.storage.check_connectivity().await } @@ -72,7 +72,7 @@ impl Storage for PrefixStorage { path: &Path, range: Range, ) -> crate::StorageResult { - self.storage.get_slice(&self.prefix.join(path), range).await + self.get_slice_unboxed(path, range).await } async fn get_all(&self, path: &Path) -> crate::StorageResult { @@ -167,17 +167,28 @@ impl Storage for PrefixStorage { } } +impl StorageGetSlice for PrefixStorage { + async fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> StorageResult { + let prefixed_path = self.prefix.join(path); + self.storage.get_slice_unboxed(&prefixed_path, range).await + } +} + /// Creates a [`PrefixStorage`] using an underlying storage and a prefix. -pub(crate) fn add_prefix_to_storage( - storage: Arc, +pub(crate) fn add_prefix_to_storage( + storage: T, prefix: PathBuf, uri: Uri, -) -> Arc { - Arc::new(PrefixStorage { +) -> PrefixStorage { + PrefixStorage { storage, prefix, uri, - }) + } } fn strip_prefix_from_error(error: BulkDeleteError, prefix: &Path) -> BulkDeleteError { @@ -233,8 +244,8 @@ fn strip_prefix_from_error(error: BulkDeleteError, prefix: &Path) -> BulkDeleteE #[cfg(test)] mod tests { - use std::collections::HashMap; + use std::sync::Arc; use std::time::SystemTime; use futures::{TryStreamExt, stream}; diff --git a/quickwit/quickwit-storage/src/ram_storage.rs b/quickwit/quickwit-storage/src/ram_storage.rs index e81b14beab5..230985edb9a 100644 --- a/quickwit/quickwit-storage/src/ram_storage.rs +++ b/quickwit/quickwit-storage/src/ram_storage.rs @@ -14,6 +14,7 @@ use std::collections::HashMap; use std::fmt; +use std::future::Future; use std::io::Cursor; use std::ops::Range; use std::path::{Path, PathBuf}; @@ -21,14 +22,13 @@ use std::sync::Arc; use async_trait::async_trait; use quickwit_common::uri::{Protocol, Uri}; -use quickwit_config::StorageBackend; use tokio::io::{AsyncRead, AsyncWriteExt}; use tokio::sync::RwLock; use crate::prefix_storage::add_prefix_to_storage; use crate::storage::SendableAsync; use crate::{ - BulkDeleteError, OwnedBytes, Storage, StorageErrorKind, StorageFactory, StorageResolverError, + BulkDeleteError, OwnedBytes, Storage, StorageErrorKind, StorageGetSlice, StorageResolverError, StorageResult, }; @@ -73,6 +73,14 @@ impl RamStorage { self.files.read().await.get(path).cloned() } + async fn get_slice_impl(&self, path: &Path, range: Range) -> StorageResult { + let payload_bytes = self.get_data(path).await.ok_or_else(|| { + StorageErrorKind::NotFound + .with_error(anyhow::anyhow!("failed to find dest_path {:?}", path)) + })?; + Ok(payload_bytes.slice(range.start..range.end)) + } + /// Returns the list of files that are present in the RamStorage. pub async fn list_files(&self) -> Vec { self.files.read().await.keys().cloned().collect() @@ -106,11 +114,7 @@ impl Storage for RamStorage { } async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { - let payload_bytes = self.get_data(path).await.ok_or_else(|| { - StorageErrorKind::NotFound - .with_error(anyhow::anyhow!("failed to find dest_path {:?}", path)) - })?; - Ok(payload_bytes.slice(range.start..range.end)) + self.get_slice_impl(path, range).await } async fn get_slice_stream( @@ -157,6 +161,16 @@ impl Storage for RamStorage { } } +impl StorageGetSlice for RamStorage { + fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> impl Future> + Send { + self.get_slice_impl(path, range) + } +} + /// Builder to create a prepopulated [`RamStorage`]. This is mostly useful for tests. #[derive(Default)] pub struct RamStorageBuilder { @@ -182,7 +196,7 @@ impl RamStorageBuilder { /// Storage resolver for [`RamStorage`]. pub struct RamStorageFactory { - ram_storage: Arc, + ram_storage: Arc, } impl Default for RamStorageFactory { @@ -193,13 +207,18 @@ impl Default for RamStorageFactory { } } -#[async_trait] -impl StorageFactory for RamStorageFactory { - fn backend(&self) -> StorageBackend { - StorageBackend::Ram +impl RamStorageFactory { + /// Creates a RAM storage factory backed by the supplied storage instance. + pub fn new(ram_storage: RamStorage) -> Self { + Self { + ram_storage: Arc::new(ram_storage), + } } - async fn resolve(&self, uri: &Uri) -> Result, StorageResolverError> { + pub(crate) fn resolve( + &self, + uri: &Uri, + ) -> Result>, StorageResolverError> { match uri.filepath() { Some(prefix) if uri.protocol() == Protocol::Ram => Ok(add_prefix_to_storage( self.ram_storage.clone(), @@ -231,16 +250,16 @@ mod tests { async fn test_ram_storage_factory() { let ram_storage_factory = RamStorageFactory::default(); let ram_uri = Uri::for_test("s3:///foo"); - let err = ram_storage_factory.resolve(&ram_uri).await.err().unwrap(); + let err = ram_storage_factory.resolve(&ram_uri).err().unwrap(); assert!(matches!(err, StorageResolverError::InvalidUri { .. })); let data_uri = Uri::for_test("ram:///data"); - let data_storage = ram_storage_factory.resolve(&data_uri).await.ok().unwrap(); + let data_storage = ram_storage_factory.resolve(&data_uri).ok().unwrap(); let home_uri = Uri::for_test("ram:///home"); - let home_storage = ram_storage_factory.resolve(&home_uri).await.ok().unwrap(); + let home_storage = ram_storage_factory.resolve(&home_uri).ok().unwrap(); assert_ne!(data_storage.uri(), home_storage.uri()); - let data_storage_two = ram_storage_factory.resolve(&data_uri).await.ok().unwrap(); + let data_storage_two = ram_storage_factory.resolve(&data_uri).ok().unwrap(); assert_eq!(data_storage.uri(), data_storage_two.uri()); } diff --git a/quickwit/quickwit-storage/src/split_cache/download_task.rs b/quickwit/quickwit-storage/src/split_cache/download_task.rs index 1f58b691459..6ef1ccde440 100644 --- a/quickwit/quickwit-storage/src/split_cache/download_task.rs +++ b/quickwit/quickwit-storage/src/split_cache/download_task.rs @@ -21,7 +21,7 @@ use quickwit_common::split_file; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use crate::split_cache::split_table::{CandidateSplit, DownloadOpportunity}; -use crate::{SearchSplitCache, StorageResolver}; +use crate::{SearchSplitCache, Storage, StorageResolver}; async fn download_split( root_path: &Path, diff --git a/quickwit/quickwit-storage/src/split_cache/mod.rs b/quickwit/quickwit-storage/src/split_cache/mod.rs index 57569060697..fa6eed4f53f 100644 --- a/quickwit/quickwit-storage/src/split_cache/mod.rs +++ b/quickwit/quickwit-storage/src/split_cache/mod.rs @@ -35,7 +35,7 @@ use tracing::{error, info, instrument, warn}; use crate::file_descriptor_cache::{FileDescriptorCache, SplitFile}; use crate::split_cache::download_task::spawn_download_task; use crate::split_cache::split_table::SplitTable; -use crate::{Storage, StorageCache, wrap_storage_with_cache}; +use crate::{StorageCache, StorageGetSlice, StorageWithCache, wrap_storage_with_cache}; /// On disk Cache of splits for searchers. /// @@ -126,7 +126,10 @@ impl SearchSplitCache { } /// Wraps a storage with our split cache. - pub fn wrap_storage(self_arc: Arc, storage: Arc) -> Arc { + pub fn wrap_storage( + self_arc: Arc, + storage: T, + ) -> StorageWithCache { let cache = Arc::new(SplitCacheBackingStorage { split_cache: self_arc, storage_root_uri: storage.uri().clone(), diff --git a/quickwit/quickwit-storage/src/storage.rs b/quickwit/quickwit-storage/src/storage.rs index ca87c124a07..ff449e9a7ea 100644 --- a/quickwit/quickwit-storage/src/storage.rs +++ b/quickwit/quickwit-storage/src/storage.rs @@ -13,9 +13,11 @@ // limitations under the License. use std::fmt; +use std::future::Future; use std::io::{self}; use std::ops::Range; use std::path::{Path, PathBuf}; +use std::sync::Arc; use std::time::SystemTime; use async_trait::async_trait; @@ -77,10 +79,11 @@ pub trait Storage: fmt::Debug + Send + Sync + 'static { /// /// See also `copy_to_file`. /// - /// async_trait Expansion of - /// async fn copy_to(&self, path: &Path, output: &mut dyn SendableAsync) -> StorageResult<()>; + /// `async_trait` expansion of: + /// `async fn copy_to(&self, path: &Path, output: &mut dyn SendableAsync) -> StorageResult<()>;` /// - /// Just putting the async form is breaking mockall. + /// The expanded form is required because `mockall` cannot generate a matching implementation + /// for the async form when the method borrows multiple arguments. fn copy_to<'life0, 'life1, 'life2, 'async_trait>( &'life0 self, path: &'life1 Path, @@ -173,6 +176,113 @@ pub trait Storage: fmt::Debug + Send + Sync + 'static { fn uri(&self) -> &Uri; } +/// Statically dispatched counterpart of [`Storage::get_slice`]. +/// +/// `async_trait` keeps [`Storage`] object-safe by boxing the future returned by every async +/// method. Generic storage wrappers use this companion trait internally so composing several +/// layers does not allocate one boxed future per layer. Type erasure can still happen once at the +/// outermost boundary, where a single boxed future is unavoidable. +pub trait StorageGetSlice: Storage { + /// Downloads a slice without boxing the returned future. + fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> impl Future> + Send; +} + +impl StorageGetSlice for Arc { + fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> impl Future> + Send { + (**self).get_slice_unboxed(path, range) + } +} + +// A storage that was already erased cannot recover static dispatch. This adapter lets generic +// code accept an existing `Arc` while returning its one already-boxed future directly. +impl StorageGetSlice for Arc { + async fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> StorageResult { + Storage::get_slice(self.as_ref(), path, range).await + } +} + +#[cfg(any(test, feature = "testsuite"))] +impl StorageGetSlice for MockStorage { + async fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> StorageResult { + Storage::get_slice(self, path, range).await + } +} + +#[async_trait] +impl Storage for Arc { + async fn check_connectivity(&self) -> anyhow::Result<()> { + (**self).check_connectivity().await + } + + async fn put(&self, path: &Path, payload: Box) -> StorageResult<()> { + (**self).put(path, payload).await + } + + async fn copy_to(&self, path: &Path, output: &mut dyn SendableAsync) -> StorageResult<()> { + (**self).copy_to(path, output).await + } + + async fn copy_to_file(&self, path: &Path, output_path: &Path) -> StorageResult { + (**self).copy_to_file(path, output_path).await + } + + async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { + (**self).get_slice(path, range).await + } + + async fn get_slice_stream( + &self, + path: &Path, + range: Range, + ) -> StorageResult> { + (**self).get_slice_stream(path, range).await + } + + async fn get_all(&self, path: &Path) -> StorageResult { + (**self).get_all(path).await + } + + async fn delete(&self, path: &Path) -> StorageResult<()> { + (**self).delete(path).await + } + + async fn bulk_delete<'a>(&self, paths: &[&'a Path]) -> Result<(), BulkDeleteError> { + (**self).bulk_delete(paths).await + } + + fn list(&self, prefix: &Path) -> ListObjectsStream { + (**self).list(prefix) + } + + async fn exists(&self, path: &Path) -> StorageResult { + (**self).exists(path).await + } + + async fn file_num_bytes(&self, path: &Path) -> StorageResult { + (**self).file_num_bytes(path).await + } + + fn uri(&self) -> &Uri { + (**self).uri() + } +} + async fn default_copy_to_file( storage: &S, path: &Path, diff --git a/quickwit/quickwit-storage/src/storage_factory.rs b/quickwit/quickwit-storage/src/storage_factory.rs index 61b2bcfaf66..0fdd46ef996 100644 --- a/quickwit/quickwit-storage/src/storage_factory.rs +++ b/quickwit/quickwit-storage/src/storage_factory.rs @@ -12,24 +12,224 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::ops::Range; +use std::path::Path; use std::sync::Arc; use async_trait::async_trait; use quickwit_common::uri::Uri; use quickwit_config::StorageBackend; +use tokio::io::AsyncRead; -use crate::{Storage, StorageResolverError}; +use crate::local_file_storage::{LocalFileStorage, LocalFileStorageFactory}; +use crate::prefix_storage::PrefixStorage; +use crate::ram_storage::{RamStorage, RamStorageFactory}; +use crate::storage::SendableAsync; +#[cfg(feature = "azure")] +use crate::{AzureBlobStorage, AzureBlobStorageFactory}; +use crate::{ + BulkDeleteError, DebouncedStorage, ListObjectsStream, OwnedBytes, PutPayload, + S3CompatibleObjectStorage, S3CompatibleObjectStorageFactory, Storage, StorageGetSlice, + StorageResolverError, StorageResult, +}; +#[cfg(feature = "gcs")] +use crate::{GoogleCloudStorageFactory, opendal_storage::OpendalStorage}; + +/// A registered factory for one of Quickwit's supported storage backends. +/// +/// Keeping the heterogeneous registry as an enum lets each factory return its concrete storage +/// type. Type erasure is deferred until the complete storage stack is shared with a caller. +pub enum StorageFactory { + /// Local filesystem storage. + File(LocalFileStorageFactory), + /// In-memory storage. + Ram(RamStorageFactory), + /// S3-compatible object storage. + S3(S3CompatibleObjectStorageFactory), + /// Azure Blob Storage. + #[cfg(feature = "azure")] + Azure(AzureBlobStorageFactory), + /// Google Cloud Storage. + #[cfg(feature = "gcs")] + Google(GoogleCloudStorageFactory), + /// A backend unavailable in this build. + Unsupported(UnsupportedStorage), +} + +impl StorageFactory { + /// Returns the backend targeted by this factory. + pub fn backend(&self) -> StorageBackend { + match self { + Self::File(_) => StorageBackend::File, + Self::Ram(_) => StorageBackend::Ram, + Self::S3(_) => StorageBackend::S3, + #[cfg(feature = "azure")] + Self::Azure(_) => StorageBackend::Azure, + #[cfg(feature = "gcs")] + Self::Google(_) => StorageBackend::Google, + Self::Unsupported(factory) => factory.backend, + } + } + + pub(crate) async fn resolve(&self, uri: &Uri) -> Result { + let storage = match self { + Self::File(factory) => ResolvedStorageInner::File(factory.resolve(uri)?), + Self::Ram(factory) => ResolvedStorageInner::Ram(factory.resolve(uri)?), + Self::S3(factory) => ResolvedStorageInner::S3(factory.resolve(uri).await?), + #[cfg(feature = "azure")] + Self::Azure(factory) => ResolvedStorageInner::Azure(factory.resolve(uri)?), + #[cfg(feature = "gcs")] + Self::Google(factory) => ResolvedStorageInner::Google(factory.resolve(uri)?), + Self::Unsupported(factory) => { + return Err(StorageResolverError::UnsupportedBackend( + factory.message.to_string(), + )); + } + }; + Ok(ResolvedStorage { storage }) + } +} + +impl From for StorageFactory { + fn from(factory: LocalFileStorageFactory) -> Self { + Self::File(factory) + } +} + +impl From for StorageFactory { + fn from(factory: RamStorageFactory) -> Self { + Self::Ram(factory) + } +} + +impl From for StorageFactory { + fn from(factory: S3CompatibleObjectStorageFactory) -> Self { + Self::S3(factory) + } +} + +#[cfg(feature = "azure")] +impl From for StorageFactory { + fn from(factory: AzureBlobStorageFactory) -> Self { + Self::Azure(factory) + } +} + +#[cfg(feature = "gcs")] +impl From for StorageFactory { + fn from(factory: GoogleCloudStorageFactory) -> Self { + Self::Google(factory) + } +} + +impl From for StorageFactory { + fn from(factory: UnsupportedStorage) -> Self { + Self::Unsupported(factory) + } +} + +/// A concrete storage returned by [`crate::StorageResolver`]. +/// +/// The backend variants are intentionally private. Callers operate on this type through the +/// [`Storage`] implementation and can keep composing generic storage layers before introducing an +/// erased trait object at their outermost consumer boundary. +#[derive(Debug)] +pub struct ResolvedStorage { + storage: ResolvedStorageInner, +} + +#[derive(Debug)] +enum ResolvedStorageInner { + File(DebouncedStorage), + Ram(PrefixStorage>), + S3(DebouncedStorage), + #[cfg(feature = "azure")] + Azure(DebouncedStorage), + #[cfg(feature = "gcs")] + Google(DebouncedStorage), +} + +macro_rules! dispatch_storage { + ($self:expr, $storage:ident => $expression:expr) => { + match &$self.storage { + ResolvedStorageInner::File($storage) => $expression, + ResolvedStorageInner::Ram($storage) => $expression, + ResolvedStorageInner::S3($storage) => $expression, + #[cfg(feature = "azure")] + ResolvedStorageInner::Azure($storage) => $expression, + #[cfg(feature = "gcs")] + ResolvedStorageInner::Google($storage) => $expression, + } + }; +} -/// A storage factory builds a [`Storage`] object for a target [`StorageBackend`] from a -/// [`Uri`]. -#[cfg_attr(any(test, feature = "testsuite"), mockall::automock)] #[async_trait] -pub trait StorageFactory: Send + Sync + 'static { - /// Returns the storage backend targeted by the factory. - fn backend(&self) -> StorageBackend; +impl Storage for ResolvedStorage { + async fn check_connectivity(&self) -> anyhow::Result<()> { + dispatch_storage!(self, storage => storage.check_connectivity().await) + } - /// Returns the appropriate [`Storage`] object for the URI. - async fn resolve(&self, uri: &Uri) -> Result, StorageResolverError>; + async fn put(&self, path: &Path, payload: Box) -> StorageResult<()> { + dispatch_storage!(self, storage => storage.put(path, payload).await) + } + + async fn copy_to(&self, path: &Path, output: &mut dyn SendableAsync) -> StorageResult<()> { + dispatch_storage!(self, storage => storage.copy_to(path, output).await) + } + + async fn copy_to_file(&self, path: &Path, output_path: &Path) -> StorageResult { + dispatch_storage!(self, storage => storage.copy_to_file(path, output_path).await) + } + + async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { + dispatch_storage!(self, storage => storage.get_slice(path, range).await) + } + + async fn get_slice_stream( + &self, + path: &Path, + range: Range, + ) -> StorageResult> { + dispatch_storage!(self, storage => storage.get_slice_stream(path, range).await) + } + + async fn get_all(&self, path: &Path) -> StorageResult { + dispatch_storage!(self, storage => storage.get_all(path).await) + } + + async fn delete(&self, path: &Path) -> StorageResult<()> { + dispatch_storage!(self, storage => storage.delete(path).await) + } + + async fn bulk_delete<'a>(&self, paths: &[&'a Path]) -> Result<(), BulkDeleteError> { + dispatch_storage!(self, storage => storage.bulk_delete(paths).await) + } + + fn list(&self, prefix: &Path) -> ListObjectsStream { + dispatch_storage!(self, storage => storage.list(prefix)) + } + + async fn exists(&self, path: &Path) -> StorageResult { + dispatch_storage!(self, storage => storage.exists(path).await) + } + + async fn file_num_bytes(&self, path: &Path) -> StorageResult { + dispatch_storage!(self, storage => storage.file_num_bytes(path).await) + } + + fn uri(&self) -> &Uri { + dispatch_storage!(self, storage => storage.uri()) + } +} + +impl StorageGetSlice for ResolvedStorage { + async fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> StorageResult { + dispatch_storage!(self, storage => storage.get_slice_unboxed(path, range).await) + } } /// A storage factory for handling unsupported or unavailable storage backends. @@ -45,16 +245,3 @@ impl UnsupportedStorage { Self { backend, message } } } - -#[async_trait] -impl StorageFactory for UnsupportedStorage { - fn backend(&self) -> StorageBackend { - self.backend - } - - async fn resolve(&self, _uri: &Uri) -> Result, StorageResolverError> { - Err(StorageResolverError::UnsupportedBackend( - self.message.to_string(), - )) - } -} diff --git a/quickwit/quickwit-storage/src/storage_resolver.rs b/quickwit/quickwit-storage/src/storage_resolver.rs index 3853f763c94..1711245a8cf 100644 --- a/quickwit/quickwit-storage/src/storage_resolver.rs +++ b/quickwit/quickwit-storage/src/storage_resolver.rs @@ -27,14 +27,16 @@ use crate::AzureBlobStorageFactory; use crate::GoogleCloudStorageFactory; use crate::local_file_storage::LocalFileStorageFactory; use crate::ram_storage::RamStorageFactory; -use crate::{S3CompatibleObjectStorageFactory, Storage, StorageFactory, StorageResolverError}; +use crate::{ + ResolvedStorage, S3CompatibleObjectStorageFactory, StorageFactory, StorageResolverError, +}; /// Returns the [`Storage`] instance associated with the protocol of a URI. The actual creation of /// storage objects is delegated to pre-registered [`StorageFactory`]. The resolver is only /// responsible for dispatching to the appropriate factory. #[derive(Clone)] pub struct StorageResolver { - per_backend_factories: Arc>>, + per_backend_factories: Arc>, } impl fmt::Debug for StorageResolver { @@ -50,7 +52,7 @@ impl StorageResolver { } /// Resolves the given URI. - pub async fn resolve(&self, uri: &Uri) -> Result, StorageResolverError> { + pub async fn resolve(&self, uri: &Uri) -> Result, StorageResolverError> { let backend = match uri.protocol() { Protocol::Azure => StorageBackend::Azure, Protocol::File => StorageBackend::File, @@ -70,7 +72,7 @@ impl StorageResolver { StorageResolverError::UnsupportedBackend(message) })?; let storage = storage_factory.resolve(uri).await?; - Ok(storage) + Ok(Arc::new(storage)) } /// Creates and returns a default [`StorageResolver`] with the default storage configuration for @@ -148,14 +150,15 @@ impl StorageResolver { #[derive(Default)] pub struct StorageResolverBuilder { - per_backend_factories: HashMap>, + per_backend_factories: HashMap, } impl StorageResolverBuilder { /// Registers a [`StorageFactory`]. - pub fn register(mut self, storage_factory: S) -> Self { + pub fn register>(mut self, storage_factory: S) -> Self { + let storage_factory = storage_factory.into(); self.per_backend_factories - .insert(storage_factory.backend(), Box::new(storage_factory)); + .insert(storage_factory.backend(), storage_factory); self } @@ -173,26 +176,16 @@ mod tests { use std::path::Path; use super::*; - use crate::{MockStorageFactory, RamStorage}; + use crate::{RamStorage, Storage}; #[tokio::test] async fn test_storage_resolver_simple() -> anyhow::Result<()> { - let mut file_storage_factory = MockStorageFactory::new(); - file_storage_factory - .expect_backend() - .returning(|| StorageBackend::File); - - let mut ram_storage_factory = MockStorageFactory::new(); - ram_storage_factory - .expect_backend() - .returning(|| StorageBackend::Ram); - ram_storage_factory.expect_resolve().returning(|_uri| { - Ok(Arc::new( - RamStorage::builder() - .put("hello", b"hello_content_second") - .build(), - )) - }); + let file_storage_factory = LocalFileStorageFactory; + let ram_storage_factory = RamStorageFactory::new( + RamStorage::builder() + .put("/hello", b"hello_content_second") + .build(), + ); let storage_resolver = StorageResolver::builder() .register(file_storage_factory) .register(ram_storage_factory) @@ -206,25 +199,16 @@ mod tests { #[tokio::test] async fn test_storage_resolver_override() -> anyhow::Result<()> { - let mut first_ram_storage_factory = MockStorageFactory::new(); - first_ram_storage_factory - .expect_backend() - .returning(|| StorageBackend::Ram); - - let mut second_ram_storage_factory = MockStorageFactory::new(); - second_ram_storage_factory - .expect_backend() - .returning(|| StorageBackend::Ram); - second_ram_storage_factory - .expect_resolve() - .returning(|uri| { - assert_eq!(uri.as_str(), "ram:///home"); - Ok(Arc::new( - RamStorage::builder() - .put("hello", b"hello_content_second") - .build(), - )) - }); + let first_ram_storage_factory = RamStorageFactory::new( + RamStorage::builder() + .put("/home/hello", b"hello_content_first") + .build(), + ); + let second_ram_storage_factory = RamStorageFactory::new( + RamStorage::builder() + .put("/home/hello", b"hello_content_second") + .build(), + ); let storage_resolver = StorageResolver::builder() .register(first_ram_storage_factory) .register(second_ram_storage_factory) diff --git a/quickwit/quickwit-storage/src/timeout_and_retry_storage.rs b/quickwit/quickwit-storage/src/timeout_and_retry_storage.rs index 1e0ed5d9d2a..dae8ec437fa 100644 --- a/quickwit/quickwit-storage/src/timeout_and_retry_storage.rs +++ b/quickwit/quickwit-storage/src/timeout_and_retry_storage.rs @@ -14,7 +14,6 @@ use std::ops::Range; use std::path::Path; -use std::sync::Arc; use async_trait::async_trait; use quickwit_common::uri::Uri; @@ -25,98 +24,56 @@ use tokio::io::AsyncRead; use crate::storage::SendableAsync; use crate::{ - BulkDeleteError, ListObjectsStream, PutPayload, Storage, StorageErrorKind, StorageResult, + BulkDeleteError, ListObjectsStream, PutPayload, Storage, StorageErrorKind, StorageGetSlice, + StorageResult, }; /// Storage proxy that implements a retry operation if the underlying storage /// takes too long. /// /// This is useful in order to ensure a low latency on S3. -/// Retrying agressively is recommended for S3. +/// Retrying aggressively is recommended for S3. /// /// #[derive(Clone, Debug)] -pub struct TimeoutAndRetryStorage { - underlying: Arc, +pub struct TimeoutAndRetryStorage { + storage: T, storage_timeout_policy: StorageTimeoutPolicy, } -impl TimeoutAndRetryStorage { +impl TimeoutAndRetryStorage { /// Creates a new `TimeoutAndRetryStorage`. /// /// See [StorageTimeoutPolicy] for more information. - pub fn new(storage: Arc, storage_timeout_policy: StorageTimeoutPolicy) -> Self { - TimeoutAndRetryStorage { - underlying: storage, + pub fn new(storage: T, storage_timeout_policy: StorageTimeoutPolicy) -> Self { + Self { + storage, storage_timeout_policy, } } } #[async_trait] -impl Storage for TimeoutAndRetryStorage { +impl Storage for TimeoutAndRetryStorage { async fn check_connectivity(&self) -> anyhow::Result<()> { - self.underlying.check_connectivity().await + self.storage.check_connectivity().await } async fn put(&self, path: &Path, payload: Box) -> StorageResult<()> { - self.underlying.put(path, payload).await + self.storage.put(path, payload).await } - fn copy_to<'life0, 'life1, 'life2, 'async_trait>( - &'life0 self, - path: &'life1 Path, - output: &'life2 mut dyn SendableAsync, - ) -> ::core::pin::Pin< - Box< - dyn ::core::future::Future> - + ::core::marker::Send - + 'async_trait, - >, - > - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - Self: 'async_trait, - { - self.underlying.copy_to(path, output) + async fn copy_to(&self, path: &Path, output: &mut dyn SendableAsync) -> StorageResult<()> { + self.storage.copy_to(path, output).await } async fn copy_to_file(&self, path: &Path, output_path: &Path) -> StorageResult { - self.underlying.copy_to_file(path, output_path).await + self.storage.copy_to_file(path, output_path).await } /// Downloads a slice of a file from the storage, and returns an in memory buffer async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { - let num_bytes = range.len(); - for (attempt_id, timeout_duration) in self - .storage_timeout_policy - .compute_timeout(num_bytes) - .enumerate() - { - let get_slice_fut = self.underlying.get_slice(path, range.clone()); - // TODO test avoid aborting timed out requests. #5468 - match tokio::time::timeout(timeout_duration, get_slice_fut).await { - Ok(result) => { - match attempt_id { - 0 => crate::metrics::GET_SLICE_TIMEOUT_SUCCESS_AFTER_0_TIMEOUT.inc(), - 1 => crate::metrics::GET_SLICE_TIMEOUT_SUCCESS_AFTER_1_TIMEOUT.inc(), - _ => crate::metrics::GET_SLICE_TIMEOUT_SUCCESS_AFTER_2_PLUS_TIMEOUT.inc(), - } - return result; - } - Err(_elapsed) => { - rate_limited_info!(limit_per_min=60, num_bytes=num_bytes, path=%path.display(), timeout_secs=timeout_duration.as_secs_f32(), "get timeout elapsed"); - continue; - } - } - } - rate_limited_warn!(limit_per_min=60, num_bytes=num_bytes, path=%path.display(), "all get_slice attempts timeouted"); - crate::metrics::GET_SLICE_TIMEOUT_ALL_TIMEOUTS.inc(); - return Err( - StorageErrorKind::Timeout.with_error(anyhow::anyhow!("internal timeout on get_slice")) - ); + self.get_slice_unboxed(path, range).await } async fn get_slice_stream( @@ -124,42 +81,76 @@ impl Storage for TimeoutAndRetryStorage { path: &Path, range: Range, ) -> StorageResult> { - self.underlying.get_slice_stream(path, range).await + self.storage.get_slice_stream(path, range).await } async fn get_all(&self, path: &Path) -> StorageResult { - self.underlying.get_all(path).await + self.storage.get_all(path).await } async fn delete(&self, path: &Path) -> StorageResult<()> { - self.underlying.delete(path).await + self.storage.delete(path).await } async fn bulk_delete<'a>(&self, paths: &[&'a Path]) -> Result<(), BulkDeleteError> { - self.underlying.bulk_delete(paths).await + self.storage.bulk_delete(paths).await } fn list(&self, prefix: &Path) -> ListObjectsStream { - self.underlying.list(prefix) + self.storage.list(prefix) } async fn exists(&self, path: &Path) -> StorageResult { - self.underlying.exists(path).await + self.storage.exists(path).await } async fn file_num_bytes(&self, path: &Path) -> StorageResult { - self.underlying.file_num_bytes(path).await + self.storage.file_num_bytes(path).await } fn uri(&self) -> &Uri { - self.underlying.uri() + self.storage.uri() + } +} + +impl StorageGetSlice for TimeoutAndRetryStorage { + async fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> StorageResult { + let num_bytes = range.len(); + for (attempt_id, timeout_duration) in self + .storage_timeout_policy + .compute_timeout(num_bytes) + .enumerate() + { + let get_slice_fut = self.storage.get_slice_unboxed(path, range.clone()); + // TODO test avoid aborting timed out requests. #5468 + match tokio::time::timeout(timeout_duration, get_slice_fut).await { + Ok(result) => { + match attempt_id { + 0 => crate::metrics::GET_SLICE_TIMEOUT_SUCCESS_AFTER_0_TIMEOUT.inc(), + 1 => crate::metrics::GET_SLICE_TIMEOUT_SUCCESS_AFTER_1_TIMEOUT.inc(), + _ => crate::metrics::GET_SLICE_TIMEOUT_SUCCESS_AFTER_2_PLUS_TIMEOUT.inc(), + } + return result; + } + Err(_elapsed) => { + rate_limited_info!(limit_per_min=60, num_bytes=num_bytes, path=%path.display(), timeout_secs=timeout_duration.as_secs_f32(), "get timeout elapsed"); + } + } + } + rate_limited_warn!(limit_per_min=60, num_bytes=num_bytes, path=%path.display(), "all get_slice attempts timeouted"); + crate::metrics::GET_SLICE_TIMEOUT_ALL_TIMEOUTS.inc(); + Err(StorageErrorKind::Timeout.with_error(anyhow::anyhow!("internal timeout on get_slice"))) } } #[cfg(test)] mod tests { - use std::sync::Mutex; + use std::sync::{Arc, Mutex}; use std::time::Duration; use tokio::time::Instant; @@ -192,23 +183,11 @@ mod tests { async fn put(&self, _path: &Path, _payload: Box) -> StorageResult<()> { todo!(); } - fn copy_to<'life0, 'life1, 'life2, 'async_trait>( - &'life0 self, - _path: &'life1 Path, - _output: &'life2 mut dyn SendableAsync, - ) -> ::core::pin::Pin< - Box< - dyn ::core::future::Future> - + ::core::marker::Send - + 'async_trait, - >, - > - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - Self: 'async_trait, - { + async fn copy_to( + &self, + _path: &Path, + _output: &mut dyn SendableAsync, + ) -> StorageResult<()> { todo!(); } @@ -247,6 +226,16 @@ mod tests { } } + impl StorageGetSlice for StorageWithDelay { + async fn get_slice_unboxed( + &self, + path: &Path, + range: Range, + ) -> StorageResult { + Storage::get_slice(self, path, range).await + } + } + #[tokio::test] async fn test_timeout_and_retry_storage() { tokio::time::pause(); @@ -263,8 +252,7 @@ mod tests { let now = Instant::now(); let storage_with_delay = StorageWithDelay::new(vec![Duration::from_secs(5), Duration::from_secs(3)]); - let storage = - TimeoutAndRetryStorage::new(Arc::new(storage_with_delay), timeout_policy.clone()); + let storage = TimeoutAndRetryStorage::new(Arc::new(storage_with_delay), timeout_policy); assert_eq!( storage.get_slice(path, 10..100).await.unwrap_err().kind, StorageErrorKind::Timeout