Add new lookup_wait_v2 to fix generation parsing bug

crates/adapter/src/fastly/core.rs

@@ -2511,7 +2511,7 @@ pub mod fastly_kv_store {
     #[repr(C)]
     pub struct InsertConfig {
         pub mode: InsertMode,
-        pub if_generation_match: u32,
+        pub if_generation_match: u64,
         pub metadata: *const u8,
         pub metadata_len: u32,
         pub time_to_live_sec: u32,
@@ -2602,9 +2602,10 @@ pub mod fastly_kv_store {
         pub struct InsertConfigOptions: u32 {
             const RESERVED = 1 << 0;
             const BACKGROUND_FETCH = 1 << 1;
-            const IF_GENERATION_MATCH = 1 << 2;
+            const RESERVED_2 = 1 << 2;
             const METADATA = 1 << 3;
             const TIME_TO_LIVE_SEC = 1 << 4;
+            const IF_GENERATION_MATCH = 1 << 5;
         }
         /// `ListConfigOptions` codings.
         #[derive(Default)]
@@ -2791,6 +2792,63 @@ pub mod fastly_kv_store {
             }
         );
 
+        let body = res.body();
+        let generation = 0;
+
+        unsafe {
+            *body_handle_out = body;
+            *generation_out = generation;
+        }
+
+        FastlyStatus::OK
+    }
+
+    #[export_name = "fastly_kv_store#lookup_wait_v2"]
+    pub fn lookup_wait_v2(
+        pending_handle: PendingObjectStoreLookupHandle,
+        body_handle_out: *mut BodyHandle,
+        metadata_out: *mut u8,
+        metadata_len: usize,
+        nwritten_out: *mut usize,
+        generation_out: *mut u64,
+        kv_error_out: *mut KvError,
+    ) -> FastlyStatus {
+        let res = match kv_store::lookup_wait(pending_handle) {
+            Ok((res, status)) => {
+                unsafe {
+                    *kv_error_out = status.into();
+                }
+
+                let Some(res) = res else {
+                    return FastlyStatus::OK;
+                };
+
+                res
+            }
+            Err(e) => {
+                unsafe {
+                    *kv_error_out = KvError::Uninitialized;
+                }
+
+                return e.into();
+            }
+        };
+
+        with_buffer!(
+            metadata_out,
+            metadata_len,
+            { res.metadata(u64::try_from(metadata_len).trapping_unwrap()) },
+            |res| {
+                let buf = handle_buffer_len!(res, nwritten_out);
+
+                unsafe {
+                    *nwritten_out = buf.as_ref().map(Vec::len).unwrap_or(0);
+                }
+
+                std::mem::forget(buf);
+            }
+        );
+
         let body = res.body();
         let generation = res.generation();
 

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

@@ -878,6 +878,8 @@
         (result $err (expected (error $fastly_status)))
     )
 
+    ;; deprecated, generation always returns 0
+    ;; via a failed u32::parse of a u64
     (@interface func (export "lookup_wait")
         (param $handle $kv_store_lookup_handle)
         (param $body_handle_out (@witx pointer $body_handle))
@@ -889,6 +891,17 @@
         (result $err (expected (error $fastly_status)))
     )
 
+    (@interface func (export "lookup_wait_v2")
+        (param $handle $kv_store_lookup_handle)
+        (param $body_handle_out (@witx pointer $body_handle))
+        (param $metadata_buf (@witx pointer (@witx char8)))
+        (param $metadata_buf_len (@witx usize))
+        (param $nwritten_out (@witx pointer (@witx usize)))
+        (param $generation_out (@witx pointer u64))
+        (param $kv_error_out (@witx pointer $kv_error))
+        (result $err (expected (error $fastly_status)))
+    )
+
     (@interface func (export "insert")
         (param $store $kv_store_handle)
         (param $key string)

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

@@ -422,9 +422,11 @@
     (flags (@witx repr u32)
        $reserved
        $background_fetch
-       $if_generation_match
+       ;; reserved_2 was previously if_generation_match (u32)
+       $reserved_2
        $metadata
        $time_to_live_sec
+       $if_generation_match
        ))
 
 (typename $kv_insert_mode
@@ -437,10 +439,11 @@
 (typename $kv_insert_config
   (record
     (field $mode $kv_insert_mode)
-    (field $if_generation_match u32)
+    (field $unused u32)
     (field $metadata (@witx pointer (@witx char8)))
     (field $metadata_len u32)
     (field $time_to_live_sec u32)
+    (field $if_generation_match u64)
     ))
 
 (typename $kv_list_config_options

lib/src/component/kv_store.rs

@@ -22,7 +22,7 @@ use {
 pub struct LookupResult {
     body: http_body::BodyHandle,
     metadata: Option<Vec<u8>>,
-    generation: u32,
+    generation: u64,
 }
 
 #[async_trait::async_trait]
@@ -54,7 +54,7 @@ impl kv_store::HostLookupResult for ComponentCtx {
     async fn generation(
         &mut self,
         rep: wasmtime::component::Resource<kv_store::LookupResult>,
-    ) -> wasmtime::Result<u32> {
+    ) -> wasmtime::Result<u64> {
         Ok(self.table().get(&rep)?.generation)
     }
 

lib/src/object_store.rs

@@ -14,7 +14,7 @@ pub struct ObjectValue {
     pub body: Vec<u8>,
     pub metadata: Vec<u8>,
     pub metadata_len: usize,
-    pub generation: u32,
+    pub generation: u64,
     pub expiration: Option<SystemTime>,
 }
 
@@ -90,7 +90,7 @@ impl ObjectStores {
         obj_key: ObjectKey,
         obj: Vec<u8>,
         mode: KvInsertMode,
-        generation: Option<u32>,
+        generation: Option<u64>,
         metadata: Option<Vec<u8>>,
         ttl: Option<std::time::Duration>,
     ) -> Result<(), KvStoreError> {
@@ -153,7 +153,7 @@ impl ObjectStores {
             generation: SystemTime::now()
                 .duration_since(SystemTime::UNIX_EPOCH)
                 .unwrap()
-                .as_nanos() as u32,
+                .as_nanos() as u64,
             expiration: exp,
         };
 

lib/src/session.rs

@@ -700,7 +700,7 @@ impl Session {
         obj_key: ObjectKey,
         obj: Vec<u8>,
         mode: Option<KvInsertMode>,
-        generation: Option<u32>,
+        generation: Option<u64>,
         metadata: Option<Vec<u8>>,
         ttl: Option<Duration>,
     ) -> Result<(), KvStoreError> {

lib/src/wiggle_abi.rs

@@ -78,7 +78,7 @@ wiggle::from_witx!({
     async: {
         fastly_async_io::{select},
         fastly_object_store::{delete_async, pending_delete_wait, insert, insert_async, pending_insert_wait, lookup_async, pending_lookup_wait, list},
-        fastly_kv_store::{lookup, lookup_wait, insert, insert_wait, delete, delete_wait, list, list_wait},
+        fastly_kv_store::{lookup, lookup_wait, lookup_wait_v2, insert, insert_wait, delete, delete_wait, list, list_wait},
         fastly_http_body::{append, read, write},
         fastly_http_cache::{lookup, transaction_lookup, insert, transaction_insert, transaction_insert_and_stream_back, transaction_update, transaction_update_and_return_fresh, transaction_record_not_cacheable, transaction_abandon, found, close, get_suggested_backend_request, get_suggested_cache_options, prepare_response_for_storage, get_found_response, get_state, get_length, get_max_age_ns, get_stale_while_revalidate_ns, get_age_ns, get_hits, get_sensitive_data, get_surrogate_keys, get_vary_rule},
         fastly_http_req::{

lib/src/wiggle_abi/kv_store_impl.rs

@@ -81,6 +81,57 @@ impl FastlyKvStore for Session {
             .recv()
             .await?;
 
+        match resp {
+            Ok(value) => {
+                let body_handle = self.insert_body(value.body.into());
+
+                memory.write(body_handle_out, body_handle)?;
+                match value.metadata_len {
+                    0 => memory.write(nwritten_out, 0)?,
+                    len => {
+                        let meta_len_u32 =
+                            u32::try_from(len).expect("metadata len is outside the bounds of u32");
+                        memory.write(nwritten_out, meta_len_u32)?;
+                        if meta_len_u32 > metadata_buf_len {
+                            return Err(Error::BufferLengthError {
+                                buf: "metadata",
+                                len: "specified length",
+                            });
+                        }
+                        memory.copy_from_slice(
+                            &value.metadata,
+                            metadata_buf.as_array(meta_len_u32),
+                        )?;
+                    }
+                }
+                memory.write(generation_out, 0)?;
+                memory.write(kv_error_out, KvError::Ok)?;
+                Ok(())
+            }
+            Err(e) => {
+                memory.write(kv_error_out, (&e).into())?;
+                Ok(())
+            }
+        }
+    }
+
+    async fn lookup_wait_v2(
+        &mut self,
+        memory: &mut GuestMemory<'_>,
+        pending_kv_lookup_handle: KvStoreLookupHandle,
+        body_handle_out: GuestPtr<BodyHandle>,
+        metadata_buf: GuestPtr<u8>,
+        metadata_buf_len: u32,
+        nwritten_out: GuestPtr<u32>,
+        generation_out: GuestPtr<u64>,
+        kv_error_out: GuestPtr<KvError>,
+    ) -> Result<(), Error> {
+        let resp = self
+            .take_pending_kv_lookup(pending_kv_lookup_handle.into())?
+            .task()
+            .recv()
+            .await?;
+
         match resp {
             Ok(value) => {
                 let body_handle = self.insert_body(value.body.into());

lib/wit/deps/fastly/compute.wit

@@ -746,7 +746,7 @@ interface kv-store {
   resource lookup-result {
     body: func() -> body-handle;
     metadata: func(max-len: u64) -> result<option<list<u8>>, error>;
-    generation: func() -> u32;
+    generation: func() -> u64;
   }
 
   lookup-wait: func(
@@ -770,7 +770,7 @@ interface kv-store {
 
   record insert-config {
     mode: insert-mode,
-    if-generation-match: u32,
+    if-generation-match: u64,
     metadata: list<u8>,
     time-to-live-sec: u32,
   }