Where the requested chunk has an offset/size the cache stores the downloaded byte range under a key that doesn't account for this range.

When another request is received we get a cache hit even though the requested byte range may differ.

Fix this by downloading and caching the entire chunk, applying the byte range post download or cache hit.

If caching is disabled honour the byte range in the S3 client download.

Author: maxstack

src/app.rs

@@ -182,45 +182,39 @@ 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
-/// * `request_data`: RequestData object for the request
+/// * `bucket`: Name of the bucket
+/// * `key`: Name of the object in the bucket
+/// * `range`: Optional byte range
 /// * `resource_manager`: ResourceManager object
-#[tracing::instrument(level = "DEBUG", skip(client, request_data, resource_manager))]
+/// * `mem_permits`: Memory permits for the request
+#[tracing::instrument(level = "DEBUG", skip(client, bucket, key, range, resource_manager))]
 async fn download_s3_object<'a>(
     client: &s3_client::S3Client,
-    request_data: &models::RequestData,
+    bucket: &str,
+    key: &str,
+    range: Option<String>,
     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(
-            &request_data.bucket,
-            &request_data.object,
-            range,
-            resource_manager,
-            &mut mem_permits,
-        )
+        .download_object(bucket, key, 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
 /// * `request_data`: RequestData object for the request
 /// * `resource_manager`: ResourceManager object
+/// * `mem_permits`: Memory permits for the request
 /// * `chunk_cache`: ChunkCache object
 #[tracing::instrument(
     level = "DEBUG",
@@ -285,19 +279,35 @@ async fn download_and_cache_s3_object<'a>(
             .instrument(tracing::Span::current())
             .await?;
         if let Some(bytes) = cache_value {
-            return Ok(bytes);
+            return Ok(s3_client::apply_range(
+                bytes,
+                request_data.offset,
+                request_data.size,
+            ));
         }
     }
 
-    let data = download_s3_object(client, request_data, resource_manager, mem_permits).await?;
+    let data = download_s3_object(
+        client,
+        &request_data.bucket,
+        &request_data.object,
+        None,
+        resource_manager,
+        mem_permits,
+    )
+    .await?;
 
     // Write data to cache
-    chunk_cache.set(&key, data.clone()).await?;
+    chunk_cache.set(&key, &data).await?;
 
     // Increment the prometheus metric for cache misses
     LOCAL_CACHE_MISSES.with_label_values(&["disk"]).inc();
 
-    Ok(data)
+    Ok(s3_client::apply_range(
+        data,
+        request_data.offset,
+        request_data.size,
+    ))
 }
 
 /// Handler for Active Storage operations
@@ -341,7 +351,9 @@ async fn operation_handler<T: operation::Operation>(
 
     let data = match (&state.args.use_chunk_cache, &state.chunk_cache) {
         (false, _) => {
-            download_s3_object(&s3_client, &request_data, &state.resource_manager, _mem_permits)
+            // 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)
                 .instrument(tracing::Span::current())
                 .await?
         }
@@ -393,10 +405,7 @@ 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

@@ -22,11 +22,11 @@ struct ChunkCacheEntry {
 
 impl ChunkCacheEntry {
     /// Return a ChunkCacheEntry object
-    fn new(key: &str, value: Bytes) -> Self {
+    fn new(key: &str, value: &Bytes) -> Self {
         let key = key.to_owned();
         // Make sure we own the `Bytes` so we don't see unexpected, but not incorrect,
         // behaviour caused by the zero copy of `Bytes`. i.e. let us choose when to copy.
-        let value = Bytes::copy_from_slice(&value);
+        let value = Bytes::copy_from_slice(value);
         Self { key, value }
     }
 }
@@ -103,8 +103,8 @@ 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 {
+    pub async fn set(&self, key: &str, value: &Bytes) -> Result<(), ActiveStorageError> {
+        match self.sender.send(ChunkCacheEntry::new(key, &value)).await {
             Ok(_) => Ok(()),
             Err(e) => Err(ActiveStorageError::ChunkCacheError {
                 error: format!("{}", e),

src/s3_client.rs

@@ -273,6 +273,26 @@ 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::*;