POC: Remove deeply nested stack of boxed Storage::get_slice futures

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"

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;

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,

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<dyn Storage>` 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<Arc<dyn Storage>>,
+    storage: OnceCell<Arc<ResolvedStorage>>,
 }
 
 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<dyn Storage>> {
+    async fn storage(&self) -> ObjectStoreResult<&Arc<ResolvedStorage>> {
         self.storage
             .get_or_try_init(|| async {
                 self.storage_resolver

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<T> {
+    storage_directory: StorageDirectory<T>,
     path: PathBuf,
 }
 
-impl HasLen for StorageDirectoryFileHandle {
+impl<T> HasLen for StorageDirectoryFileHandle<T> {
     fn len(&self) -> usize {
         unimplemented!()
     }
 }
 
-impl fmt::Debug for StorageDirectoryFileHandle {
+impl<T: Storage> fmt::Debug for StorageDirectoryFileHandle<T> {
     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<T: StorageGetSlice + Clone> FileHandle for StorageDirectoryFileHandle<T> {
     fn read_bytes(&self, _byte_range: Range<usize>) -> io::Result<OwnedBytes> {
         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<dyn Storage>,
+pub struct StorageDirectory<T> {
+    storage: T,
 }
 
-impl Debug for StorageDirectory {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "StorageDirectory({:?})", self.uri())
+impl<T: Storage> Debug for StorageDirectory<T> {
+    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+        formatter
+            .debug_tuple("StorageDirectory")
+            .field(&self.uri())
+            .finish()
     }
 }
 
-impl StorageDirectory {
+impl<T: Storage> StorageDirectory<T> {
     /// Creates a new StorageDirectory, backed by the given `storage`.
-    pub fn new(storage: Arc<dyn Storage>) -> 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<usize>) -> io::Result<OwnedBytes> {
-        let payload: OwnedBytes = self.storage.get_slice(path, range).await?;
+    pub async fn get_slice(&self, path: &Path, range: Range<usize>) -> io::Result<OwnedBytes>
+    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<T: StorageGetSlice + Clone> Directory for StorageDirectory<T> {
     fn get_file_handle(&self, path: &Path) -> Result<Arc<dyn FileHandle>, OpenReadError> {
         Ok(Arc::new(StorageDirectoryFileHandle {
             storage_directory: self.clone(),

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<dyn Storage> = 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<dyn Storage> = self
             .storage_resolver
             .resolve(&index_config.index_uri)
             .await?;

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};

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;

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<T: StorageGetSlice + Clone>(
     searcher_context: Arc<SearcherContext>,
     mut global_doc_addrs: Vec<GlobalDocAddress>,
-    index_storage: Arc<dyn Storage>,
+    storage: T,
     splits: &[SplitIdAndFooterOffsets],
     doc_mapper: Arc<DocMapper>,
     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<T: StorageGetSlice + Clone>(
     searcher_context: Arc<SearcherContext>,
     partial_hits: Vec<PartialHit>,
-    index_storage: Arc<dyn Storage>,
+    storage: T,
     splits: &[SplitIdAndFooterOffsets],
     doc_mapper: Arc<DocMapper>,
     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<T: StorageGetSlice + Clone>(
     searcher_context: Arc<SearcherContext>,
     mut global_doc_addrs: Vec<GlobalDocAddress>,
-    index_storage: Arc<dyn Storage>,
+    storage: T,
     split: &SplitIdAndFooterOffsets,
     doc_mapper: Arc<DocMapper>,
     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,