Incorporate the request offset and size into the chunk cache key so cache chunks aren't shared between requests of differing byte range.

Revert the previous fix which suffers from memory over consumption when bombarded by requests for byte ranges from the same very large file - we would need a way to ensure a single download of that file before servicing concurrent requests against it.

Author: maxstack

src/app.rs

@@ -182,33 +182,41 @@ async fn schema() -> &'static str {
 
 /// Download an object from S3
 ///
+/// Requests a byte range if `offset` or `size` is specified in the request.
+///
 /// # Arguments
 ///
 /// * `client`: S3 client object
-/// * `bucket`: Name of the bucket
-/// * `key`: Name of the object in the bucket
-/// * `range`: Optional byte range
+/// * `request_data`: RequestData object for the request
 /// * `resource_manager`: ResourceManager object
 /// * `mem_permits`: Memory permits for the request
-#[tracing::instrument(level = "DEBUG", skip(client, bucket, key, range, resource_manager))]
+#[tracing::instrument(level = "DEBUG", skip(client, request_data, resource_manager))]
 async fn download_s3_object<'a>(
     client: &s3_client::S3Client,
-    bucket: &str,
-    key: &str,
-    range: Option<String>,
+    request_data: &models::RequestData,
     resource_manager: &'a ResourceManager,
     mut mem_permits: Option<SemaphorePermit<'a>>,
 ) -> Result<Bytes, ActiveStorageError> {
+    // Convert request data to byte range for S3 request
+    let range = s3_client::get_range(request_data.offset, request_data.size);
     // Acquire connection permit to be freed via drop when this function returns
     let _conn_permits = resource_manager.s3_connection().await?;
 
     client
-        .download_object(bucket, key, range, resource_manager, &mut mem_permits)
+        .download_object(
+            &request_data.bucket,
+            &request_data.object,
+            range,
+            resource_manager,
+            &mut mem_permits,
+        )
         .await
 }
 
 /// Download and cache an object from S3
 ///
+/// Requests a byte range if `offset` or `size` is specified in the request.
+///
 /// # Arguments
 ///
 /// * `client`: S3 client object
@@ -232,12 +240,14 @@ async fn download_and_cache_s3_object<'a>(
     // which may feasibly indicate a change to the upstream object
     // lead to a new cache key.
     let key = format!(
-        "{}-{}-{}-{}-{:?}-{:?}",
+        "{}-{}-{}-{}-{:?}-{:?}-{:?}-{:?}",
         request_data.source.as_str(),
         request_data.bucket,
         request_data.object,
         request_data.dtype,
         request_data.byte_order,
+        request_data.offset,
+        request_data.size,
         request_data.compression,
     );
 
@@ -279,35 +289,19 @@ async fn download_and_cache_s3_object<'a>(
             .instrument(tracing::Span::current())
             .await?;
         if let Some(bytes) = cache_value {
-            return Ok(s3_client::apply_range(
-                bytes,
-                request_data.offset,
-                request_data.size,
-            ));
+            return Ok(bytes);
         }
     }
 
-    let data = download_s3_object(
-        client,
-        &request_data.bucket,
-        &request_data.object,
-        None,
-        resource_manager,
-        mem_permits,
-    )
-    .await?;
+    let data = download_s3_object(client, request_data, resource_manager, mem_permits).await?;
 
     // Write data to cache
     chunk_cache.set(&key, &data).await?;
 
     // Increment the prometheus metric for cache misses
     LOCAL_CACHE_MISSES.with_label_values(&["disk"]).inc();
 
-    Ok(s3_client::apply_range(
-        data,
-        request_data.offset,
-        request_data.size,
-    ))
+    Ok(data)
 }
 
 /// Handler for Active Storage operations
@@ -351,9 +345,7 @@ async fn operation_handler<T: operation::Operation>(
 
     let data = match (&state.args.use_chunk_cache, &state.chunk_cache) {
         (false, _) => {
-            // Convert request data offset and size to byte range for S3 request
-            let range = s3_client::get_range(request_data.offset, request_data.size);
-            download_s3_object(&s3_client, &request_data.bucket, &request_data.object, range, &state.resource_manager, _mem_permits)
+            download_s3_object(&s3_client, &request_data, &state.resource_manager, _mem_permits)
                 .instrument(tracing::Span::current())
                 .await?
         }
@@ -405,7 +397,10 @@ fn operation<T: operation::Operation>(
         assert_eq!(ptr, data.as_ptr());
     }
     // Convert to a mutable vector to allow in-place byte order conversion.
+    let ptr = data.as_ptr();
     let vec: Vec<u8> = data.into();
+    // Assert that we're using zero-copy.
+    assert_eq!(ptr, vec.as_ptr());
     debug_span!("operation").in_scope(|| T::execute(&request_data, vec))
 }
 

src/chunk_cache.rs

@@ -104,7 +104,7 @@ impl ChunkCache {
     /// * `key`: Unique key identifying the chunk
     /// * `value`: Chunk `Bytes` to be cached
     pub async fn set(&self, key: &str, value: &Bytes) -> Result<(), ActiveStorageError> {
-        match self.sender.send(ChunkCacheEntry::new(key, &value)).await {
+        match self.sender.send(ChunkCacheEntry::new(key, value)).await {
             Ok(_) => Ok(()),
             Err(e) => Err(ActiveStorageError::ChunkCacheError {
                 error: format!("{}", e),

src/s3_client.rs

@@ -273,26 +273,6 @@ pub fn get_range(offset: Option<usize>, size: Option<usize>) -> Option<String> {
     }
 }
 
-/// Apply an optional byte range to a `Bytes` object
-/// to yield the same bytes as would be returned by a
-/// download request with HTTP Range header.
-///
-/// # Arguments
-///
-/// * `offset`: Optional offset of data in bytes
-/// * `size`: Optional size of data in bytes
-pub fn apply_range(bytes: Bytes, offset: Option<usize>, size: Option<usize>) -> Bytes {
-    match (offset, size) {
-        (offset, Some(size)) => {
-            let offset = offset.unwrap_or(0);
-            let end = offset + size;
-            bytes.slice(offset..end)
-        }
-        (Some(offset), None) => bytes.slice(offset..),
-        _ => bytes,
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;