Rework fields, but make the borrow checker bigmad

elliottt · elliottt · commit 1fc34950a341 · 2023-09-18T11:35:13.000-07:00
diff --git a/crates/wasi-http/src/body.rs b/crates/wasi-http/src/body.rs
@@ -1,4 +1,7 @@
-use crate::bindings::http::types::{self, Headers, HeadersRef, Method, Scheme};
+use crate::{
+    bindings::http::types,
+    types::FieldMap,
+};
 use bytes::Bytes;
 use std::{pin, task};
 use tokio::sync::{mpsc, oneshot};
@@ -200,7 +203,7 @@ impl HostIncomingBody {
 
 pub struct HostFutureTrailers {
     pub worker: AbortOnDropJoinHandle<()>,
-    pub received: Option<Result<HeadersRef, types::Error>>,
+    pub received: Option<Result<FieldMap, types::Error>>,
     pub receiver: oneshot::Receiver<Result<hyper::HeaderMap, hyper::Error>>,
 }
 
@@ -221,7 +224,9 @@ impl HostFutureTrailers {
                 }
 
                 match Pin::new(&mut self.0.receiver).poll(cx) {
-                    Poll::Ready(Ok(Ok(headers))) => self.0.received = Some(Ok(headers)),
+                    Poll::Ready(Ok(Ok(headers))) => {
+                        self.0.received = Some(Ok(FieldMap::from(headers)))
+                    }
 
                     Poll::Ready(Ok(Err(e))) => {
                         self.0.received = Some(Err(types::Error::ProtocolError(format!(
diff --git a/crates/wasi-http/src/types.rs b/crates/wasi-http/src/types.rs
@@ -2,13 +2,13 @@
 //! implementation of the wasi-http API.
 
 use crate::{
-    bindings::http::types::{Headers, IncomingBody, Method, Scheme, FutureTrailers},
-    body::{HostIncomingBody, HostFutureTrailers},
+    bindings::http::types::{FutureTrailers, Headers, IncomingBody, Method, Scheme},
+    body::{HostFutureTrailers, HostIncomingBody},
 };
-use std::collections::HashMap;
 use std::pin::Pin;
 use std::task;
-use wasmtime_wasi::preview2::{pipe::AsyncReadStream, AbortOnDropJoinHandle, Table, TableError};
+use std::{any::Any, collections::HashMap};
+use wasmtime_wasi::preview2::{pipe::AsyncReadStream, AbortOnDropJoinHandle, Table, TableError, OccupiedEntry};
 
 const MAX_BUF_SIZE: usize = 65_536;
 
@@ -20,14 +20,12 @@ pub trait WasiHttpView: Send {
     fn table(&mut self) -> &mut Table;
 }
 
-pub type FieldsMap = HashMap<String, Vec<Vec<u8>>>;
-
 pub struct HostOutgoingRequest {
     pub method: Method,
     pub scheme: Option<Scheme>,
     pub path_with_query: String,
     pub authority: String,
-    pub headers: HostFields,
+    pub headers: FieldMap,
     pub body: Option<AsyncReadStream>,
 }
 
@@ -46,26 +44,15 @@ pub struct HostOutgoingRequest {
 
 pub struct HostIncomingResponse {
     pub status: u16,
-    pub headers: HeadersRef,
+    pub headers: FieldMap,
     pub body: Option<hyper::body::Incoming>,
     pub worker: AbortOnDropJoinHandle<anyhow::Result<()>>,
 }
 
-pub enum HeadersRef {
-    Value(hyper::HeaderMap),
-    Resource(Headers),
-}
-
-#[derive(Clone, Debug)]
-pub struct HostFields(pub HashMap<String, Vec<Vec<u8>>>);
-
-impl HostFields {
-    pub fn new() -> Self {
-        Self(FieldsMap::new())
-    }
-}
+#[derive(Clone)]
+pub struct FieldMap(pub HashMap<String, Vec<Vec<u8>>>);
 
-impl From<hyper::HeaderMap> for HostFields {
+impl From<hyper::HeaderMap> for FieldMap {
     fn from(headers: hyper::HeaderMap) -> Self {
         use std::collections::hash_map::Entry;
 
@@ -85,6 +72,21 @@ impl From<hyper::HeaderMap> for HostFields {
     }
 }
 
+pub enum HostFields {
+    Ref {
+        parent: u32,
+
+        // NOTE: there's not failure in the result here because we assume that HostFields will
+        // always be registered as a child of the entry with the `parent` id. This ensures that the
+        // entry will always exist while this `HostFields::Ref` entry exists in the table, thus we
+        // don't need to account for failure when fetching the fields ref from the parent.
+        get_fields: for<'a> fn(elem: &'a mut dyn Any) -> &'a mut FieldMap,
+    },
+    Owned {
+        fields: FieldMap,
+    },
+}
+
 pub struct IncomingResponseInternal {
     pub resp: hyper::Response<hyper::body::Incoming>,
     pub worker: AbortOnDropJoinHandle<anyhow::Result<()>>,
@@ -153,8 +155,7 @@ pub trait TableHttpExt {
     fn delete_incoming_response(&mut self, id: u32) -> Result<HostIncomingResponse, TableError>;
 
     fn push_fields(&mut self, fields: HostFields) -> Result<u32, TableError>;
-    fn get_fields(&self, id: u32) -> Result<&HostFields, TableError>;
-    fn get_fields_mut(&mut self, id: u32) -> Result<&mut HostFields, TableError>;
+    fn get_fields(&mut self, id: u32) -> Result<&mut FieldMap, TableError>;
     fn delete_fields(&mut self, id: u32) -> Result<HostFields, TableError>;
 
     fn push_future_incoming_response(
@@ -178,9 +179,18 @@ pub trait TableHttpExt {
     fn get_incoming_body(&mut self, id: IncomingBody) -> Result<&mut HostIncomingBody, TableError>;
     fn delete_incoming_body(&mut self, id: IncomingBody) -> Result<HostIncomingBody, TableError>;
 
-    fn push_future_trailers(&mut self, trailers: HostFutureTrailers) -> Result<FutureTrailers, TableError>;
-    fn get_future_trailers(&mut self, id: FutureTrailers) -> Result<&mut HostFutureTrailers, TableError>;
-    fn delete_future_trailers(&mut self, id: FutureTrailers) -> Result<HostFutureTrailers, TableError>;
+    fn push_future_trailers(
+        &mut self,
+        trailers: HostFutureTrailers,
+    ) -> Result<FutureTrailers, TableError>;
+    fn get_future_trailers(
+        &mut self,
+        id: FutureTrailers,
+    ) -> Result<&mut HostFutureTrailers, TableError>;
+    fn delete_future_trailers(
+        &mut self,
+        id: FutureTrailers,
+    ) -> Result<HostFutureTrailers, TableError>;
 }
 
 #[async_trait::async_trait]
@@ -223,13 +233,18 @@ impl TableHttpExt for Table {
     }
 
     fn push_fields(&mut self, fields: HostFields) -> Result<u32, TableError> {
-        self.push(Box::new(fields))
-    }
-    fn get_fields(&self, id: u32) -> Result<&HostFields, TableError> {
-        self.get::<HostFields>(id)
+        match fields {
+            HostFields::Ref { parent, .. } => self.push_child(Box::new(fields), parent),
+            HostFields::Owned { .. } => self.push(Box::new(fields)),
+        }
     }
-    fn get_fields_mut(&mut self, id: u32) -> Result<&mut HostFields, TableError> {
-        self.get_mut::<HostFields>(id)
+    fn get_fields<'a>(&'a mut self, id: u32) -> Result<&'a mut FieldMap, TableError> {
+        let (parent, get_fields) = match self.get_mut::<HostFields>(id)? {
+            HostFields::Ref { parent, get_fields } => (*parent, *get_fields),
+            HostFields::Owned { fields } => return Ok(fields),
+        };
+        let mut entry: OccupiedEntry<'a> = self.entry(parent).unwrap();
+        Ok(get_fields(entry.get_mut()))
     }
     fn delete_fields(&mut self, id: u32) -> Result<HostFields, TableError> {
         let fields = self.delete::<HostFields>(id)?;
@@ -273,15 +288,24 @@ impl TableHttpExt for Table {
         self.delete(id)
     }
 
-    fn push_future_trailers(&mut self, trailers: HostFutureTrailers) -> Result<FutureTrailers, TableError> {
+    fn push_future_trailers(
+        &mut self,
+        trailers: HostFutureTrailers,
+    ) -> Result<FutureTrailers, TableError> {
         self.push(Box::new(trailers))
     }
 
-    fn get_future_trailers(&mut self, id: FutureTrailers) -> Result<&mut HostFutureTrailers, TableError> {
+    fn get_future_trailers(
+        &mut self,
+        id: FutureTrailers,
+    ) -> Result<&mut HostFutureTrailers, TableError> {
         self.get_mut(id)
     }
 
-    fn delete_future_trailers(&mut self, id: FutureTrailers) -> Result<HostFutureTrailers, TableError> {
+    fn delete_future_trailers(
+        &mut self,
+        id: FutureTrailers,
+    ) -> Result<HostFutureTrailers, TableError> {
         self.delete(id)
     }
 }
diff --git a/crates/wasi-http/src/types_impl.rs b/crates/wasi-http/src/types_impl.rs
@@ -6,12 +6,13 @@ use crate::bindings::http::types::{
     Scheme, StatusCode, Trailers,
 };
 use crate::body::HostFutureTrailers;
+use crate::types::FieldMap;
 use crate::WasiHttpView;
 use crate::{
     body::HostIncomingBody,
     types::{
-        HeadersRef, HostFields, HostFutureIncomingResponse, HostIncomingResponse,
-        HostOutgoingRequest, TableHttpExt,
+        HostFields, HostFutureIncomingResponse, HostIncomingResponse, HostOutgoingRequest,
+        TableHttpExt,
     },
 };
 use anyhow::{anyhow, bail, Context};
@@ -31,14 +32,24 @@ impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
         Ok(())
     }
     async fn new_fields(&mut self, entries: Vec<(String, Vec<u8>)>) -> wasmtime::Result<Fields> {
-        let mut map = HostFields::new();
+        use std::collections::{hash_map::Entry, HashMap};
+
+        let mut map: HashMap<String, Vec<Vec<u8>>> = HashMap::new();
+
         for (key, value) in entries {
-            map.0.insert(key, vec![value.clone()]);
+            match map.entry(key) {
+                Entry::Occupied(mut entry) => entry.get_mut().push(value),
+                Entry::Vacant(entry) => {
+                    entry.insert(vec![value]);
+                }
+            }
         }
 
         let id = self
             .table()
-            .push_fields(map)
+            .push_fields(HostFields::Owned {
+                fields: FieldMap(map),
+            })
             .context("[new_fields] pushing fields")?;
         Ok(id)
     }
@@ -59,12 +70,12 @@ impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
         name: String,
         value: Vec<Vec<u8>>,
     ) -> wasmtime::Result<()> {
-        let m = self.table().get_fields_mut(fields)?;
+        let m = self.table().get_fields(fields)?;
         m.0.insert(name, value.clone());
         Ok(())
     }
     async fn fields_delete(&mut self, fields: Fields, name: String) -> wasmtime::Result<()> {
-        let m = self.table().get_fields_mut(fields)?;
+        let m = self.table().get_fields(fields)?;
         m.0.remove(&name);
         Ok(())
     }
@@ -76,7 +87,7 @@ impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
     ) -> wasmtime::Result<()> {
         let m = self
             .table()
-            .get_fields_mut(fields)
+            .get_fields(fields)
             .context("[fields_append] getting mutable fields")?;
         match m.0.get_mut(&name) {
             Some(v) => v.push(value),
@@ -99,12 +110,14 @@ impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
         Ok(result)
     }
     async fn fields_clone(&mut self, fields: Fields) -> wasmtime::Result<Fields> {
-        let table = self.table();
-        let m = table
+        let fields = self
+            .table()
             .get_fields(fields)
-            .context("[fields_clone] getting fields")?;
-        let id = table
-            .push_fields(m.clone())
+            .context("[fields_clone] getting fields")?
+            .clone();
+        let id = self
+            .table()
+            .push_fields(HostFields::Owned { fields })
             .context("[fields_clone] pushing fields")?;
         Ok(id)
     }
@@ -159,8 +172,8 @@ impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
         authority: Option<String>,
         headers: Headers,
     ) -> wasmtime::Result<OutgoingRequest> {
-        // We're taking ownership of the header data, so remove it from the table.
-        let headers = self.table().delete_fields(headers)?;
+        let headers = self.table().get_fields(headers)?.clone();
+
         let req = HostOutgoingRequest {
             path_with_query: path_with_query.unwrap_or("".to_string()),
             authority: authority.unwrap_or("".to_string()),
@@ -237,18 +250,19 @@ impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
         &mut self,
         response: IncomingResponse,
     ) -> wasmtime::Result<Headers> {
-        let r = self
+        let _ = self
             .table()
             .get_incoming_response_mut(response)
             .context("[incoming_response_headers] getting response")?;
 
-        let hdrs = match r.headers {
-            HeadersRef::Value(ref mut hdrs) => std::mem::take(hdrs),
-            HeadersRef::Resource(id) => return Ok(id),
-        };
+        fn get_fields(elem: &mut dyn Any) -> &mut FieldMap {
+            &mut elem.downcast_mut::<HostIncomingResponse>().unwrap().headers
+        }
 
-        let id = self.table().push_fields(HostFields::from(hdrs))?;
-        self.table().get_incoming_response_mut(response)?.headers = HeadersRef::Resource(id);
+        let id = self.table().push_fields(HostFields::Ref {
+            parent: response,
+            get_fields,
+        })?;
         Ok(id)
     }
     async fn incoming_response_consume(
@@ -311,20 +325,18 @@ impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
             return Ok(Some(Err(e.clone())));
         }
 
-        let hdrs = match res.unwrap() {
-            HeadersRef::Resource(id) => return Ok(Some(Ok(id))),
-            HeadersRef::Value(ref mut hdrs) => std::mem::take(hdrs),
-        };
-
         drop(res);
         drop(trailers);
 
-        let hdrs = self.table().push_fields(HostFields::from(hdrs))?;
+        fn get_fields(elem: &mut dyn Any) -> &mut FieldMap {
+            let trailers = elem.downcast_mut::<HostFutureTrailers>().unwrap();
+            trailers.received.as_mut().unwrap().as_mut().unwrap()
+        }
 
-        self.table()
-            .get_future_trailers(id)?
-            .received
-            .replace(Ok(HeadersRef::Resource(hdrs)));
+        let hdrs = self.table().push_fields(HostFields::Ref {
+            parent: id,
+            get_fields,
+        })?;
 
         Ok(Some(Ok(hdrs)))
     }
@@ -373,7 +385,7 @@ impl<T: WasiHttpView> crate::bindings::http::types::Host for T {
 
         let resp = self.table().push_incoming_response(HostIncomingResponse {
             status: parts.status.as_u16(),
-            headers: HeadersRef::Value(parts.headers),
+            headers: FieldMap::from(parts.headers),
             body: Some(body),
             worker: resp.worker,
         })?;
diff --git a/crates/wasi-http/wit/deps/http/types.wit b/crates/wasi-http/wit/deps/http/types.wit
@@ -90,7 +90,7 @@ interface types {
     path-with-query: option<string>,
     scheme: option<scheme>,
     authority: option<string>,
-    headers: /* own */ headers
+    headers: /* borrow */ headers
   ) -> outgoing-request
 
   // Will return the outgoing-body child at most once. If called more than
@@ -169,7 +169,7 @@ interface types {
   drop-outgoing-response: func(response: /* own */ outgoing-response)
   new-outgoing-response: func(
     status-code: status-code,
-    headers: /* own */ headers
+    headers: /* borrow */ headers
   ) -> outgoing-response
 
   /// Will give the child outgoing-response at most once. subsequent calls will
diff --git a/crates/wasi/src/preview2/table.rs b/crates/wasi/src/preview2/table.rs
@@ -203,7 +203,7 @@ impl Table {
     /// remove or replace the entry based on its contents. The methods available are a subset of
     /// [`std::collections::hash_map::OccupiedEntry`] - it does not give access to the key, it
     /// restricts replacing the entry to items of the same type, and it does not allow for deletion.
-    pub fn entry(&mut self, index: u32) -> Result<OccupiedEntry, TableError> {
+    pub fn entry<'a>(&'a mut self, index: u32) -> Result<OccupiedEntry<'a>, TableError> {
         if self.map.contains_key(&index) {
             Ok(OccupiedEntry { table: self, index })
         } else {
