Provide a S3 download handler and a cached S3 download handler with the operation handler deciding which to use based on app settings.

Ideally move the decision even further up registering one S3 download handler in the shared state, the actual one used being based on app settings. This would allow a http handler to be registered in the same way. Proves "fun" with the lifetime of objects.

Author: maxstack

src/app.rs

@@ -58,7 +58,7 @@ struct AppState {
     resource_manager: ResourceManager,
 
     /// Object chunk cache
-    chunk_cache: ChunkCache,
+    chunk_cache: Option<ChunkCache>,
 }
 
 impl AppState {
@@ -67,7 +67,11 @@ impl AppState {
         let task_limit = args.thread_limit.or_else(|| Some(num_cpus::get() - 1));
         let resource_manager =
             ResourceManager::new(args.s3_connection_limit, args.memory_limit, task_limit);
-        let chunk_cache = ChunkCache::new(args);
+        let chunk_cache = if args.use_chunk_cache {
+            Some(ChunkCache::new(args))
+        } else {
+            None
+        };
 
         Self {
             args: args.clone(),
@@ -165,7 +169,43 @@ async fn schema() -> &'static str {
     "Hello, world!"
 }
 
-/// Download and optionally cache an object from S3
+/// Download an object from S3
+///
+/// Requests a byte range if `offset` or `size` is specified in the request.
+///
+/// # Arguments
+///
+/// * `client`: S3 client object
+/// * `request_data`: RequestData object for the request
+/// * `resource_manager`: ResourceManager object
+async fn download_s3_object<'a>(
+    client: &s3_client::S3Client,
+    request_data: &models::RequestData,
+    resource_manager: &'a ResourceManager,
+) -> Result<Bytes, ActiveStorageError> {
+
+    // If we're given a size in the request data then use this to
+    // get an initial guess at the required memory resources.
+    let memory = request_data.size.unwrap_or(0);
+    let mut mem_permits = resource_manager.memory(memory).await?;
+
+    let range = s3_client::get_range(request_data.offset, request_data.size);
+    let _conn_permits = resource_manager.s3_connection().await?;
+
+    let data = client
+        .download_object(
+            &request_data.bucket,
+            &request_data.object,
+            range,
+            resource_manager,
+            &mut mem_permits,
+        )
+        .await;
+
+    data
+}
+
+/// Download and cache an object from S3
 ///
 /// Requests a byte range if `offset` or `size` is specified in the request.
 ///
@@ -175,7 +215,7 @@ async fn schema() -> &'static str {
 /// * `request_data`: RequestData object for the request
 /// * `resource_manager`: ResourceManager object
 /// * `chunk_cache`: ChunkCache object
-async fn download_object<'a>(
+async fn download_and_cache_s3_object<'a>(
     client: &s3_client::S3Client,
     request_data: &models::RequestData,
     resource_manager: &'a ResourceManager,
@@ -184,21 +224,16 @@ async fn download_object<'a>(
 
     let key = format!("{},{:?}", client, request_data);
 
-    // If we're using the chunk cache and have a hit it'll return Some(bytes), None when disabled.
     match chunk_cache.get(&key).await {
         Ok(value) => {
             if let Some(bytes) = value {
-                // TODO: remove debug
-                println!("Cache hit for key: {}", key);
                 return Ok(bytes);
             }
         },
         Err(e) => {
             return Err(e);
         }
     }
-    // TODO: remove debug
-    println!("Cache miss for key: {}", key);
 
     // If we're given a size in the request data then use this to
     // get an initial guess at the required memory resources.
@@ -264,14 +299,24 @@ async fn operation_handler<T: operation::Operation>(
         .instrument(tracing::Span::current())
         .await;
 
-    let data = download_object(
-        &s3_client,
-        &request_data,
-        &state.resource_manager,
-        &state.chunk_cache,
-    )
-    .instrument(tracing::Span::current())
-    .await?;
+    let data = if state.args.use_chunk_cache {
+        download_and_cache_s3_object(
+            &s3_client,
+            &request_data,
+            &state.resource_manager,
+            &state.chunk_cache.as_ref().unwrap(),
+        )
+        .instrument(tracing::Span::current())
+        .await?
+    } else {
+        download_s3_object(
+            &s3_client,
+            &request_data,
+            &state.resource_manager,
+        )
+        .instrument(tracing::Span::current())
+        .await?
+    };
 
     // All remaining work i s synchronous. If the use_rayon argument was specified, delegate to the
     // Rayon thread pool. Otherwise, execute as normal using Tokio.

src/chunk_cache.rs

@@ -8,17 +8,11 @@ use std::{collections::HashMap, ops::Add, path::PathBuf, time::{SystemTime, UNIX
 use tokio::fs;
 
 pub struct ChunkCache {
-    cache: Option<SimpleDiskCache>,
+    cache: SimpleDiskCache,
 }
 
 impl ChunkCache {
     pub fn new(args: &CommandLineArgs) -> Self {
-        if !args.use_chunk_cache {
-            return Self {
-                cache: None,
-            };
-        }
-
         // Path to the cache directory.
         let path = <Option<String> as Clone>::clone(&args.chunk_cache_path)
             .expect("The chunk cache path must be specified when the chunk cache is enabled");
@@ -33,44 +27,36 @@ impl ChunkCache {
         };
 
         Self {
-            cache: Some(SimpleDiskCache::new(
+            cache: SimpleDiskCache::new(
                 "chunk_cache",
                 &path,
                 lifespan,
                 60 * 60,
                 max_size_bytes
-            )),
+            ),
         }
     }
 
     pub async fn set(&self, key: &String, value: Bytes) -> Result<Option<Bytes>, ActiveStorageError> {
-        // Cache will be None with chunk caching disabled.
-        if let Some(cache) = &self.cache {
-            match cache.set(key, value).await {
-                Ok(_) => {
-                    return Ok(None);
-                },
-                Err(e) => {
-                    return Err(ActiveStorageError::ChunkCacheError{ error: format!("{:?}", e) });
-                }
+        match self.cache.set(key, value).await {
+            Ok(_) => {
+                Ok(None)
+            },
+            Err(e) => {
+                Err(ActiveStorageError::ChunkCacheError{ error: format!("{:?}", e) })
             }
         }
-        Ok(None)
     }
 
     pub async fn get(&self, key: &String) -> Result<Option<Bytes>, ActiveStorageError> {
-        // Cache will be None with chunk caching disabled.
-        if let Some(cache) = &self.cache {
-            match cache.get(key).await {
-                Ok(value) => {
-                    return Ok(value);
-                },
-                Err(e) => {
-                    return Err(ActiveStorageError::ChunkCacheError{ error: format!("{:?}", e) });
-                }
+        match self.cache.get(key).await {
+            Ok(value) => {
+                Ok(value)
+            },
+            Err(e) => {
+                Err(ActiveStorageError::ChunkCacheError{ error: format!("{:?}", e) })
             }
         }
-        Ok(None)
     }
 }