Skip to content

Make streams owned by request/response that they are tied to. #6228

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pchickey merged 6 commits into from
Apr 21, 2023
Merged

Make streams owned by request/response that they are tied to. #6228

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c7b3f25
Make streams owned by request/response that they are tied to.
brendandburns Apr 18, 2023
a72d8e0
Address comments, fix tests.
brendandburns Apr 19, 2023
14b25d2
Address comment.
brendandburns Apr 20, 2023
22a2dd5
Update crates/wasi-http/src/streams_impl.rs
brendandburns Apr 20, 2023
bf2c3ba
Switch to BytesMut
brendandburns Apr 20, 2023
b855a4b
Merge branch 'ptr' of https://github.com/brendandburns/wasmtime into ptr
brendandburns Apr 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions crates/wasi-http/src/http_impl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use crate::r#struct::ActiveResponse;
pub use crate::r#struct::WasiHttp;
use crate::r#struct::{Stream, WasiHttp};
use crate::types::{RequestOptions, Scheme};
#[cfg(not(any(target_arch = "riscv64", target_arch = "s390x")))]
use anyhow::anyhow;
Expand Down Expand Up @@ -183,8 +183,10 @@ impl WasiHttp {
let body = Full::<Bytes>::new(
self.streams
.get(&request.body)
.unwrap_or(&Bytes::new())
.clone(),
.unwrap_or(&Stream::default())
.data
.clone()
.freeze(),
);
let t = timeout(first_bytes_timeout, sender.send_request(call.body(body)?)).await?;
let mut res = t?;
Expand Down Expand Up @@ -222,7 +224,7 @@ impl WasiHttp {
}
response.body = self.streams_id_base;
self.streams_id_base = self.streams_id_base + 1;
self.streams.insert(response.body, buf.freeze());
self.streams.insert(response.body, buf.freeze().into());
self.responses.insert(response_id, response);
Ok(response_id)
}
Expand Down
42 changes: 26 additions & 16 deletions crates/wasi-http/src/streams_impl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ use crate::poll::Pollable;
use crate::streams::{InputStream, OutputStream, StreamError};
use crate::WasiHttp;
use anyhow::{anyhow, bail};
use bytes::BufMut;
use std::vec::Vec;

impl crate::streams::Host for WasiHttp {
Expand All @@ -11,10 +10,14 @@ impl crate::streams::Host for WasiHttp {
stream: InputStream,
len: u64,
) -> wasmtime::Result<Result<(Vec<u8>, bool), StreamError>> {
let s = self
let st = self
.streams
.get_mut(&stream)
.ok_or_else(|| anyhow!("stream not found: {stream}"))?;
if st.closed {
bail!("stream is dropped!");
}
let s = &mut st.data;
if len == 0 {
Ok(Ok((bytes::Bytes::new().to_vec(), s.len() > 0)))
} else if s.len() > len.try_into()? {
Expand All @@ -31,10 +34,14 @@ impl crate::streams::Host for WasiHttp {
stream: InputStream,
len: u64,
) -> wasmtime::Result<Result<(u64, bool), StreamError>> {
let s = self
let st = self
.streams
.get_mut(&stream)
.ok_or_else(|| anyhow!("stream not found: {stream}"))?;
if st.closed {
bail!("stream is dropped!");
}
let s = &mut st.data;
if len == 0 {
Ok(Ok((0, s.len() > 0)))
} else if s.len() > len.try_into()? {
Expand All @@ -52,7 +59,11 @@ impl crate::streams::Host for WasiHttp {
}

fn drop_input_stream(&mut self, stream: InputStream) -> wasmtime::Result<()> {
self.streams.remove(&stream);
let st = self
.streams
.get_mut(&stream)
.ok_or_else(|| anyhow!("stream not found: {stream}"))?;
st.closed = true;
Ok(())
}

Expand All @@ -61,18 +72,13 @@ impl crate::streams::Host for WasiHttp {
this: OutputStream,
buf: Vec<u8>,
) -> wasmtime::Result<Result<u64, StreamError>> {
match self.streams.get(&this) {
Some(data) => {
let mut new = bytes::BytesMut::with_capacity(data.len() + buf.len());
new.put(data.clone());
new.put(bytes::Bytes::from(buf.clone()));
self.streams.insert(this, new.freeze());
}
None => {
self.streams.insert(this, bytes::Bytes::from(buf.clone()));
}
let len = buf.len();
let st = self.streams.entry(this).or_default();
if st.closed {
bail!("cannot write to closed stream");
}
Ok(Ok(buf.len().try_into()?))
st.data.extend_from_slice(buf.as_slice());
Ok(Ok(len.try_into()?))
}

fn write_zeroes(
Expand Down Expand Up @@ -111,7 +117,11 @@ impl crate::streams::Host for WasiHttp {
}

fn drop_output_stream(&mut self, stream: OutputStream) -> wasmtime::Result<()> {
self.streams.remove(&stream);
let st = self
.streams
.get_mut(&stream)
.ok_or_else(|| anyhow!("stream not found: {stream}"))?;
Copy link
Member

@rvolosatovs rvolosatovs Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Doesn't with_context work here?

Copy link
Contributor Author

@brendandburns brendandburns Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure I'm following your suggestion, can you clarify?

Copy link
Member

@rvolosatovs rvolosatovs Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, I'm referring to this method on the anyhow::Context trait https://docs.rs/anyhow/latest/anyhow/trait.Context.html#tymethod.with_context

If it's in scope, you should be able to call context or with_context (usually used with format! and similar) on most Result and all Option types directly, thus eliminating the need for ok_or_else and explicit anyhow! invocation

Copy link
Contributor

@pchickey pchickey Apr 20, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@rvolosatovs I made those ok_or_else(|| anyhow!(...)) suggestions in a prior PR and wasn't aware that Context could make that more succinct. Thank you.

@brendandburns The suggested change here is let st = self.streams.get_mut(&stream).with_context(|| format!("stream not found: {stream}"))?;, or maybe even ....get_mut(&stream).context("stream not found")? if leaving out the stream number is ok - and it probably is, because (at least when using the component linker, I'm not sure if it works in module mode) bindgen is inserting tracing statements which someone debugging can recover the values of the arguments from.

st.closed = true;
Ok(())
}
}
27 changes: 25 additions & 2 deletions crates/wasi-http/src/struct.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,13 @@
use crate::types::{Method, Scheme};
use bytes::Bytes;
use bytes::{BufMut, Bytes, BytesMut};
use std::collections::HashMap;

#[derive(Clone, Default)]
pub struct Stream {
Copy link
Member

@rvolosatovs rvolosatovs Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking at the usages of this struct, it seems that a From<Bytes> implementation could really clear this up

Copy link
Contributor Author

@brendandburns brendandburns Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added.

pub closed: bool,
pub data: BytesMut,
}

#[derive(Clone)]
pub struct WasiHttp {
pub request_id_base: u32,
Expand All @@ -11,7 +17,7 @@ pub struct WasiHttp {
pub requests: HashMap<u32, ActiveRequest>,
pub responses: HashMap<u32, ActiveResponse>,
pub fields: HashMap<u32, HashMap<String, Vec<String>>>,
pub streams: HashMap<u32, Bytes>,
pub streams: HashMap<u32, Stream>,
}

#[derive(Clone)]
Expand Down Expand Up @@ -66,6 +72,23 @@ impl ActiveResponse {
}
}

impl Stream {
pub fn new() -> Self {
Self::default()
}
}

impl From<Bytes> for Stream {
fn from(bytes: Bytes) -> Self {
let mut buf = BytesMut::with_capacity(bytes.len());
buf.put(bytes);
Self {
Copy link
Member

@rvolosatovs rvolosatovs Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we had a Default trait impl, this could have just called Self::default()

Copy link
Contributor Author

@brendandburns brendandburns Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done.

closed: false,
data: buf,
}
}
}

impl WasiHttp {
pub fn new() -> Self {
Self {
Expand Down
15 changes: 11 additions & 4 deletions crates/wasi-http/src/types_impl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
use crate::poll::Pollable;
use crate::r#struct::ActiveRequest;
use crate::r#struct::{ActiveRequest, Stream};
use crate::types::{
Error, Fields, FutureIncomingResponse, Headers, IncomingRequest, IncomingResponse,
IncomingStream, Method, OutgoingRequest, OutgoingResponse, OutgoingStream, ResponseOutparam,
Scheme, StatusCode, Trailers,
};
use crate::WasiHttp;
use anyhow::{anyhow, bail};
use std::collections::HashMap;
use std::collections::{hash_map::Entry, HashMap};

impl crate::types::Host for WasiHttp {
fn drop_fields(&mut self, fields: Fields) -> wasmtime::Result<()> {
Expand Down Expand Up @@ -123,7 +123,10 @@ impl crate::types::Host for WasiHttp {
bail!("unimplemented: drop_incoming_request")
}
fn drop_outgoing_request(&mut self, request: OutgoingRequest) -> wasmtime::Result<()> {
self.requests.remove(&request);
if let Entry::Occupied(e) = self.requests.entry(request) {
let r = e.remove();
self.streams.remove(&r.body);
}
Ok(())
}
fn incoming_request_method(&mut self, _request: IncomingRequest) -> wasmtime::Result<Method> {
Expand Down Expand Up @@ -192,6 +195,7 @@ impl crate::types::Host for WasiHttp {
if req.body == 0 {
req.body = self.streams_id_base;
self.streams_id_base = self.streams_id_base + 1;
self.streams.insert(req.body, Stream::default());
}
Ok(Ok(req.body))
}
Expand All @@ -206,7 +210,10 @@ impl crate::types::Host for WasiHttp {
bail!("unimplemented: set_response_outparam")
}
fn drop_incoming_response(&mut self, response: IncomingResponse) -> wasmtime::Result<()> {
self.responses.remove(&response);
if let Entry::Occupied(e) = self.responses.entry(response) {
let r = e.remove();
self.streams.remove(&r.body);
}
Ok(())
}
fn drop_outgoing_response(&mut self, _response: OutgoingResponse) -> wasmtime::Result<()> {
Expand Down