Finish implementing cancel

Note that this implements a recently-changed behavior in the compute
platform. Previously, a `PendingTransaction::drop` would *block* on whichever
request got the `GoGet` obligation; more recently (and in Viceroy),
the `PendingTransaction::drop` does not block on any other transaction.

Author: cceckman-at-fastly

lib/src/cache.rs

@@ -8,6 +8,7 @@ use crate::{
     body::Body,
     wiggle_abi::types::{BodyHandle, CacheOverrideTag},
 };
+
 use http::{HeaderMap, HeaderValue};
 
 mod store;
@@ -112,6 +113,15 @@ pub struct CacheEntry {
 }
 
 impl CacheEntry {
+    /// Return a stub entry to hold in CacheBusy.
+    pub fn stub(&self) -> CacheEntry {
+        Self {
+            key: self.key.clone(),
+            found: None,
+            go_get: None,
+        }
+    }
+
     /// Returns the key used to generate this CacheEntry.
     pub fn key(&self) -> &CacheKey {
         &self.key

lib/src/component/cache.rs

@@ -301,10 +301,17 @@ impl api::Host for ComponentCtx {
         &mut self,
         handle: api::BusyHandle,
     ) -> Result<api::Handle, types::Error> {
-        let busy_handle: CacheBusyHandle = handle.into();
-        let handle: CacheHandle = busy_handle.into();
-        let _ = self.session.cache_entry_mut(handle).await?;
-        Ok(handle.into())
+        let handle = handle.into();
+        // Swap out for a distinct handle, so we don't hit a repeated `close`+`close_busy`:
+        let entry = self.session.cache_entry_mut(handle).await?;
+        let mut other_entry = entry.stub();
+        std::mem::swap(entry, &mut other_entry);
+        let task = PeekableTask::spawn(Box::pin(async move { Ok(other_entry) })).await;
+        let h: CacheHandle = self
+            .session
+            .insert_cache_op(PendingCacheTask::new(task))
+            .into();
+        Ok(h.into())
     }
 
     async fn transaction_insert(
@@ -377,18 +384,20 @@ impl api::Host for ComponentCtx {
         }
     }
 
-    async fn close_busy(&mut self, _handle: api::BusyHandle) -> Result<(), types::Error> {
-        Err(Error::Unsupported {
-            msg: "Cache API primitives not yet supported",
-        }
-        .into())
+    async fn close_busy(&mut self, handle: api::BusyHandle) -> Result<(), types::Error> {
+        // Don't wait for the transaction to complete; drop the future to cancel.
+        let _ = self.session.take_cache_entry(handle.into())?;
+        Ok(())
     }
 
-    async fn close(&mut self, _handle: api::Handle) -> Result<(), types::Error> {
-        Err(Error::Unsupported {
-            msg: "Cache API primitives not yet supported",
-        }
-        .into())
+    async fn close(&mut self, handle: api::Handle) -> Result<(), types::Error> {
+        let _ = self
+            .session
+            .take_cache_entry(handle.into())?
+            .task()
+            .recv()
+            .await?;
+        Ok(())
     }
 
     async fn get_state(&mut self, handle: api::Handle) -> Result<api::LookupState, types::Error> {

lib/src/wiggle_abi/cache.rs

@@ -225,10 +225,16 @@ impl FastlyCache for Session {
         options_mask: types::CacheLookupOptionsMask,
         options: wiggle::GuestPtr<types::CacheLookupOptions>,
     ) -> Result<types::CacheHandle, Error> {
-        let busy = self
-            .transaction_lookup_async(memory, cache_key, options_mask, options)
-            .await?;
-        self.cache_busy_handle_wait(memory, busy).await
+        let headers = load_lookup_options(self, memory, options_mask, options)?;
+        let key = load_cache_key(memory, cache_key)?;
+        let cache = Arc::clone(self.cache());
+
+        let entry = cache.transaction_lookup(&key, &headers).await;
+
+        let task = PeekableTask::spawn(Box::pin(async move { Ok(entry) })).await;
+        let task = PendingCacheTask::new(task);
+        let handle = self.insert_cache_op(task);
+        Ok(handle.into())
     }
 
     async fn transaction_lookup_async(
@@ -257,8 +263,12 @@ impl FastlyCache for Session {
         handle: types::CacheBusyHandle,
     ) -> Result<types::CacheHandle, Error> {
         let handle = handle.into();
-        let _ = self.cache_entry_mut(handle).await?;
-        Ok(handle)
+        // Swap out for a distinct handle, so we don't hit a repeated `close`+`close_busy`:
+        let entry = self.cache_entry_mut(handle).await?;
+        let mut other_entry = entry.stub();
+        std::mem::swap(entry, &mut other_entry);
+        let task = PeekableTask::spawn(Box::pin(async move { Ok(other_entry) })).await;
+        Ok(self.insert_cache_op(PendingCacheTask::new(task)).into())
     }
 
     async fn transaction_insert(
@@ -325,18 +335,24 @@ impl FastlyCache for Session {
         memory: &mut wiggle::GuestMemory<'_>,
         handle: types::CacheHandle,
     ) -> Result<(), Error> {
-        Err(Error::NotAvailable("Cache API primitives"))
+        let entry = self.cache_entry_mut(handle.into()).await?;
+        if let Some(_) = entry.take_go_get() {
+            Ok(())
+        } else {
+            Err(Error::CacheError(crate::cache::Error::CannotWrite).into())
+        }
     }
 
     async fn close_busy(
         &mut self,
         memory: &mut wiggle::GuestMemory<'_>,
         handle: types::CacheBusyHandle,
     ) -> Result<(), Error> {
-        Err(Error::NotAvailable("Cache API primitives"))
+        // Don't wait for the transaction to complete; drop the future to cancel.
+        let _ = self.take_cache_entry(handle.into())?;
+        Ok(())
     }
 
-    /// Wait for the lookup to be complete, then discard the results.
     async fn close(
         &mut self,
         memory: &mut wiggle::GuestMemory<'_>,

test-fixtures/src/bin/cache_transactional.rs

@@ -5,7 +5,8 @@ use uuid::Uuid;
 
 fn main() {
     test_racing_transactions();
-    test_implicit_cancel();
+    test_implicit_cancel_of_fetch();
+    test_implicit_cancel_of_pending();
     test_explicit_cancel();
 }
 
@@ -62,18 +63,33 @@ fn test_racing_transactions() {
     assert_eq!(&read, "hello");
 }
 
-fn test_implicit_cancel() {
+fn test_implicit_cancel_of_fetch() {
     let key = new_key();
 
     let busy1 = Transaction::lookup(key.clone()).execute_async().unwrap();
     let busy2 = Transaction::lookup(key.clone()).execute_async().unwrap();
-    let (tx, pending) = ready_and_pending(busy1, busy2);
+    let (t1, pending) = ready_and_pending(busy1, busy2);
 
     // Cancel via dropping:
-    assert!(tx.must_insert_or_update());
-    std::mem::drop(tx);
-    //let t2 = pending.wait().unwrap();
-    //assert!(!t2.found().is_some());
+    assert!(t1.must_insert_or_update());
+    std::mem::drop(t1);
+    let t2 = pending.wait().unwrap();
+    assert!(t2.found().is_none());
+    assert!(t2.must_insert_or_update());
+}
+
+fn test_implicit_cancel_of_pending() {
+    let key = new_key();
+
+    let busy1 = Transaction::lookup(key.clone()).execute_async().unwrap();
+    let busy2 = Transaction::lookup(key.clone()).execute_async().unwrap();
+    let (t1, pending) = ready_and_pending(busy1, busy2);
+
+    // Cancel the blocked request via dropping.
+    // Fun fact, this was a bug in compute platform that we fixed when writing the Viceroy
+    // implementation!
+    std::mem::drop(pending);
+    assert!(t1.must_insert_or_update());
 }
 
 fn test_explicit_cancel() {