Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Rework fields, but make the borrow checker mad
Browse files Browse the repository at this point in the history
elliottt committed Sep 18, 2023
1 parent cd8ede9 commit fd0f015
Showing 5 changed files with 130 additions and 75 deletions.
11 changes: 8 additions & 3 deletions crates/wasi-http/src/body.rs
Original file line number Diff line number Diff line change
@@ -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!(
101 changes: 65 additions & 36 deletions crates/wasi-http/src/types.rs
Original file line number Diff line number Diff line change
@@ -2,13 +2,15 @@
//! 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, OccupiedEntry, Table, TableError,
};

const MAX_BUF_SIZE: usize = 65_536;

@@ -20,14 +22,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 +46,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 +74,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 + 'static)) -> &'a mut FieldMap,
},
Owned {
fields: FieldMap,
},
}

pub struct IncomingResponseInternal {
pub resp: hyper::Response<hyper::body::Incoming>,
pub worker: AbortOnDropJoinHandle<anyhow::Result<()>>,
@@ -153,8 +157,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 +181,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 +235,21 @@ 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(&mut self, id: u32) -> Result<&mut FieldMap, TableError> {
match self.get_mut::<HostFields>(id)? {
HostFields::Ref { parent, get_fields } => {
let parent = *parent;
let get_fields = *get_fields;
let entry = self.get_any_mut(parent)?;
Ok(get_fields(entry))
}
HostFields::Owned { fields } => Ok(fields),
}
}
fn delete_fields(&mut self, id: u32) -> Result<HostFields, TableError> {
let fields = self.delete::<HostFields>(id)?;
@@ -273,15 +293,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)
}
}
78 changes: 45 additions & 33 deletions crates/wasi-http/src/types_impl.rs
Original file line number Diff line number Diff line change
@@ -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,
})?;
4 changes: 2 additions & 2 deletions crates/wasi-http/wit/deps/http/types.wit
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit fd0f015

Please sign in to comment.