Range requests for cache API

crates/adapter/src/fastly/cache.rs

@@ -20,6 +20,8 @@ bitflags::bitflags! {
     pub struct CacheLookupOptionsMask: u32 {
         const _RESERVED = 1 << 0;
         const REQUEST_HEADERS = 1 << 1;
+        const SERVICE_ID = 1 << 2;
+        const ALWAYS_USE_REQUESTED_RANGE = 1 << 3;
     }
 }
 
@@ -113,6 +115,15 @@ mod cache {
                 Self::REQUEST_HEADERS,
                 value.contains(CacheLookupOptionsMask::REQUEST_HEADERS),
             );
+            flags.set(
+                Self::SERVICE_ID,
+                value.contains(CacheLookupOptionsMask::SERVICE_ID),
+            );
+            flags.set(
+                Self::ALWAYS_USE_REQUESTED_RANGE,
+                value.contains(CacheLookupOptionsMask::ALWAYS_USE_REQUESTED_RANGE),
+            );
+
             flags
         }
     }
@@ -121,6 +132,9 @@ mod cache {
         fn from(value: CacheLookupOptions) -> Self {
             Self {
                 request_headers: value.request_headers,
+                // service_id is not supported in Viceroy.
+                // We ignore the value, but pass through the flag, so Viceroy can still raise an
+                // error when it's used-but-unsupported.
             }
         }
     }

lib/compute-at-edge-abi/cache.witx

@@ -24,6 +24,7 @@
         $reserved
         $request_headers
         $service_id
+        $always_use_requested_range
     )
 )
 
@@ -43,6 +44,7 @@
         $request_headers
         $replace_strategy
         $service_id
+        $always_use_requested_range
     )
 )
 

lib/proptest-regressions/body.txt

@@ -0,0 +1,8 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 0008d74790fb20a629ca4cc6f2c016c4a944c17d3059ee654ef4a48595b1fd0f # shrinks to body = Channel(Receiver { chan: Rx { inner: Chan { tx: Tx { block_tail: 0x760140044890, tail_position: 3 }, semaphore: Semaphore { semaphore: Semaphore { permits: 0 }, bound: 2 }, rx_waker: AtomicWaker, tx_count: 0, rx_fields: "..." } } })
+cc c1b60ad59f679fb3a4fc4bc96c50acedfc321467a9c0792f66321938054c0cad # shrinks to (body, chunk_lengths) = (b"\xe8\xfdiT3\x1dyr", [0])

lib/proptest-regressions/collecting_body.txt

@@ -0,0 +1,7 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc c36f255c0841ebe6cdb08ce10d8e5a6f12777d08a3aa4b76dd9c65c3f0dfdd27 # shrinks to (body, start) = ([b"\0"], 0)

lib/src/cache.rs

@@ -16,7 +16,7 @@ use http::{HeaderMap, HeaderValue};
 mod store;
 mod variance;
 
-use store::{CacheData, CacheKeyObjects, ObjectMeta, Obligation};
+use store::{CacheData, CacheKeyObjects, GetBodyBuilder, ObjectMeta, Obligation};
 pub use variance::VaryRule;
 
 #[derive(Debug, thiserror::Error)]
@@ -135,15 +135,36 @@ pub struct CacheEntry {
     key: CacheKey,
     found: Option<Found>,
     go_get: Option<Obligation>,
+
+    /// Respect the range in body() even when the body length is not yet known.
+    ///
+    /// When a cached item is Found, the length of the cached item may or may not be known:
+    /// if no expected length was provided and the body is still streaming, the length is unknown.
+    ///
+    /// When always_use_requested_range is false, and the length is unknown,
+    /// body() returns the full body regardless of the requested range.
+    /// When always_use_requested_range is true, and the length is unknown,
+    /// body() blocks until the start of the range is available.
+    always_use_requested_range: bool,
 }
 
 impl CacheEntry {
+    /// Set the always_use_requested_range flag.
+    /// This applies to all subsequent lookups from this CacheEntry or future entries derived from it.
+    pub fn with_always_use_requested_range(self, always_use_requested_range: bool) -> Self {
+        Self {
+            always_use_requested_range,
+            ..self
+        }
+    }
+
     /// Return a stub entry to hold in CacheBusy.
     pub fn stub(&self) -> CacheEntry {
         Self {
             key: self.key.clone(),
             found: None,
             go_get: None,
+            always_use_requested_range: false,
         }
     }
 
@@ -172,6 +193,20 @@ impl CacheEntry {
         self.go_get.take().is_some()
     }
 
+    /// Access the body of the cached item, if available.
+    pub async fn body(&self, from: Option<u64>, to: Option<u64>) -> Result<Body, crate::Error> {
+        let found = self
+            .found
+            .as_ref()
+            .ok_or(crate::Error::CacheError(Error::Missing))?;
+        found
+            .get_body()
+            .with_range(from, to)
+            .with_always_use_requested_range(self.always_use_requested_range)
+            .build()
+            .await
+    }
+
     /// Insert the provided body into the cache.
     ///
     /// Returns a CacheEntry where the new item is Found.
@@ -186,14 +221,15 @@ impl CacheEntry {
             key: self.key.clone(),
             found: Some(found),
             go_get: None,
+            always_use_requested_range: self.always_use_requested_range,
         })
     }
 
     /// Freshen the existing cache item according to the new write options,
     /// without changing the body.
-    pub fn update(&mut self, options: WriteOptions) -> Result<(), crate::Error> {
+    pub async fn update(&mut self, options: WriteOptions) -> Result<(), crate::Error> {
         let go_get = self.go_get.take().ok_or(Error::NotRevalidatable)?;
-        match go_get.update(options) {
+        match go_get.update(options).await {
             Ok(()) => Ok(()),
             Err((go_get, err)) => {
                 // On failure, preserve the obligation.
@@ -228,9 +264,8 @@ impl From<Arc<CacheData>> for Found {
 }
 
 impl Found {
-    /// Access the body of the cached object.
-    pub fn body(&self) -> Result<Body, crate::Error> {
-        self.data.as_ref().get_body()
+    fn get_body(&self) -> GetBodyBuilder {
+        self.data.as_ref().body()
     }
 
     /// Access the metadata of the cached object.
@@ -283,6 +318,7 @@ impl Cache {
             key: key.clone(),
             found,
             go_get: None,
+            always_use_requested_range: false,
         }
     }
 
@@ -301,11 +337,9 @@ impl Cache {
             .await;
         CacheEntry {
             key: key.clone(),
-            found: found.map(|data| Found {
-                data,
-                last_body_handle: None,
-            }),
+            found: found.map(|v| v.into()),
             go_get: obligation,
+            always_use_requested_range: false,
         }
     }
 
@@ -578,7 +612,7 @@ mod tests {
 
                 let nonempty = cache.lookup(&key, &HeaderMap::default()).await;
                 let found = nonempty.found().expect("should have found inserted key");
-                let got = found.body().unwrap().read_into_vec().await.unwrap();
+                let got = found.get_body().build().await.unwrap().read_into_vec().await.unwrap();
                 assert_eq!(got, value);
             });
         }
@@ -720,6 +754,7 @@ mod tests {
             stale_while_revalidate: Duration::from_secs(10),
             ..WriteOptions::default()
         })
+        .await
         .unwrap();
 
         // After this, should get the new response:
@@ -746,7 +781,7 @@ mod tests {
             stale_while_revalidate: Duration::from_secs(10),
             ..WriteOptions::default()
         };
-        txn1.update(opts.clone()).unwrap_err();
+        txn1.update(opts.clone()).await.unwrap_err();
 
         // But we should still be able to insert.
         txn1.insert(opts.clone(), Body::empty()).unwrap();