diff --git a/crates/wasi-http/src/http_impl.rs b/crates/wasi-http/src/http_impl.rs index 171b5897192b..4a327f6d4eba 100644 --- a/crates/wasi-http/src/http_impl.rs +++ b/crates/wasi-http/src/http_impl.rs @@ -183,7 +183,7 @@ impl WasiHttp { let body = Full::::new( self.streams .get(&request.body) - .unwrap_or(&Stream::new()) + .unwrap_or(&Stream::default()) .data .clone(), ); diff --git a/crates/wasi-http/src/streams_impl.rs b/crates/wasi-http/src/streams_impl.rs index 35531882dc8b..fbfef635c424 100644 --- a/crates/wasi-http/src/streams_impl.rs +++ b/crates/wasi-http/src/streams_impl.rs @@ -1,5 +1,4 @@ use crate::poll::Pollable; -use crate::r#struct::Stream; use crate::streams::{InputStream, OutputStream, StreamError}; use crate::WasiHttp; use anyhow::{anyhow, bail}; @@ -74,34 +73,23 @@ impl crate::streams::Host for WasiHttp { this: OutputStream, buf: Vec, ) -> wasmtime::Result> { + let len = buf.len(); match self.streams.get(&this) { Some(st) => { if st.closed { bail!("stream is dropped!"); } - let data = &st.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, - Stream { - closed: false, - data: new.freeze(), - }, - ); + let new_len = st.data.len() + len; + let mut new = bytes::BytesMut::with_capacity(new_len); + new.put(st.data.clone()); + new.put(bytes::Bytes::from(buf)); + self.streams.insert(this, new.freeze().into()); } None => { - self.streams.insert( - this, - Stream { - closed: false, - data: bytes::Bytes::from(buf.clone()), - }, - ); + self.streams.insert(this, bytes::Bytes::from(buf).into()); } } - Ok(Ok(buf.len().try_into()?)) + Ok(Ok(len.try_into()?)) } fn write_zeroes( diff --git a/crates/wasi-http/src/struct.rs b/crates/wasi-http/src/struct.rs index d259fce299cf..9e38af307549 100644 --- a/crates/wasi-http/src/struct.rs +++ b/crates/wasi-http/src/struct.rs @@ -2,7 +2,7 @@ use crate::types::{Method, Scheme}; use bytes::Bytes; use std::collections::HashMap; -#[derive(Clone)] +#[derive(Clone, Default)] pub struct Stream { pub closed: bool, pub data: Bytes, @@ -74,9 +74,15 @@ impl ActiveResponse { impl Stream { pub fn new() -> Self { + Self::default() + } +} + +impl From for Stream { + fn from(bytes: Bytes) -> Self { Self { closed: false, - data: Bytes::new(), + data: bytes, } } } diff --git a/crates/wasi-http/src/types_impl.rs b/crates/wasi-http/src/types_impl.rs index 3915a1ba3c7f..32cad6b0299a 100644 --- a/crates/wasi-http/src/types_impl.rs +++ b/crates/wasi-http/src/types_impl.rs @@ -1,5 +1,5 @@ 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, @@ -7,7 +7,7 @@ use crate::types::{ }; use crate::WasiHttp; use anyhow::{anyhow, bail}; -use std::collections::HashMap; +use std::collections::{HashMap, hash_map::Entry}; impl crate::types::Host for WasiHttp { fn drop_fields(&mut self, fields: Fields) -> wasmtime::Result<()> { @@ -123,12 +123,9 @@ impl crate::types::Host for WasiHttp { bail!("unimplemented: drop_incoming_request") } fn drop_outgoing_request(&mut self, request: OutgoingRequest) -> wasmtime::Result<()> { - match self.requests.get(&request) { - Some(r) => { - self.streams.remove(&r.body); - self.requests.remove(&request); - } - None => { /* pass */ } + if let Entry::Occupied(e) = self.requests.entry(request) { + let r = e.remove(); + self.streams.remove(&r.body); } Ok(()) } @@ -198,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)) } @@ -212,12 +210,9 @@ impl crate::types::Host for WasiHttp { bail!("unimplemented: set_response_outparam") } fn drop_incoming_response(&mut self, response: IncomingResponse) -> wasmtime::Result<()> { - match self.responses.get(&response) { - Some(r) => { - self.streams.remove(&r.body); - self.responses.remove(&response); - } - None => { /* pass */ } + if let Entry::Occupied(e) = self.responses.entry(response) { + let r = e.remove(); + self.streams.remove(&r.body); } Ok(()) }