-
Notifications
You must be signed in to change notification settings - Fork 144
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.
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
feat(net): AsyncWrite and AsyncRead for WebSocket #379
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
|
@@ -27,6 +27,7 @@ gloo-utils = { version = "0.2", path = "../utils", default-features = false } | |||||
wasm-bindgen-futures = "0.4" | ||||||
futures-core = { version = "0.3", optional = true } | ||||||
futures-sink = { version = "0.3", optional = true } | ||||||
futures-io = { version = "0.3", optional = true } | ||||||
|
||||||
thiserror = "1.0" | ||||||
|
||||||
|
@@ -99,3 +100,5 @@ eventsource = [ | |||||
'web-sys/EventSource', | ||||||
'web-sys/MessageEvent', | ||||||
] | ||||||
# As of now, only implements `AsyncRead` and `AsyncWrite` on `WebSocket` | ||||||
io-util = ["futures-io"] | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
This ensures There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I agree with you. However, the rest of the Cargo.toml is currently not using |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -57,6 +57,11 @@ pub struct WebSocket { | |
Closure<dyn FnMut(web_sys::Event)>, | ||
Closure<dyn FnMut(web_sys::CloseEvent)>, | ||
), | ||
/// Leftover bytes when using `AsyncRead`. | ||
/// | ||
/// These bytes are drained and returned in subsequent calls to `poll_read`. | ||
#[cfg(feature = "io-util")] | ||
pub(super) read_pending_bytes: Option<Vec<u8>>, // Same size as `Vec<u8>` alone thanks to niche optimization | ||
Comment on lines
+60
to
+64
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. An alternative to |
||
} | ||
|
||
impl WebSocket { | ||
|
@@ -196,6 +201,8 @@ impl WebSocket { | |
error_callback, | ||
close_callback, | ||
), | ||
#[cfg(feature = "io-util")] | ||
read_pending_bytes: None, | ||
}) | ||
} | ||
|
||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,168 @@ | ||
use core::cmp; | ||
use core::pin::Pin; | ||
use core::task::{Context, Poll}; | ||
use std::io; | ||
|
||
use futures_core::{ready, Stream as _}; | ||
use futures_io::{AsyncRead, AsyncWrite}; | ||
use futures_sink::Sink; | ||
|
||
use crate::websocket::futures::WebSocket; | ||
use crate::websocket::{Message as WebSocketMessage, WebSocketError}; | ||
|
||
impl WebSocket { | ||
/// Returns whether there are pending bytes left after calling [`AsyncRead::poll_read`] on this WebSocket. | ||
/// | ||
/// When calling [`AsyncRead::poll_read`], [`Stream::poll_next`](futures_core::Stream::poll_next) is called | ||
/// under the hood, and when the received item is too big to fit into the provided buffer, leftover bytes are | ||
/// stored. These leftover bytes are returned by subsequent calls to [`AsyncRead::poll_read`]. | ||
#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] | ||
pub fn has_pending_bytes(&self) -> bool { | ||
self.read_pending_bytes.is_some() | ||
} | ||
} | ||
|
||
macro_rules! try_in_poll_io { | ||
($expr:expr) => {{ | ||
match $expr { | ||
Ok(o) => o, | ||
// WebSocket is closed, nothing more to read or write | ||
Err(WebSocketError::ConnectionClose(event)) if event.was_clean => { | ||
return Poll::Ready(Ok(0)); | ||
} | ||
Err(e) => return Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))), | ||
} | ||
}}; | ||
} | ||
|
||
#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] | ||
impl AsyncRead for WebSocket { | ||
fn poll_read( | ||
mut self: Pin<&mut Self>, | ||
cx: &mut Context<'_>, | ||
buf: &mut [u8], | ||
) -> Poll<io::Result<usize>> { | ||
let mut data = if let Some(data) = self.as_mut().get_mut().read_pending_bytes.take() { | ||
data | ||
} else { | ||
match ready!(self.as_mut().poll_next(cx)) { | ||
Some(item) => match try_in_poll_io!(item) { | ||
WebSocketMessage::Text(s) => s.into_bytes(), | ||
WebSocketMessage::Bytes(data) => data, | ||
}, | ||
None => return Poll::Ready(Ok(0)), | ||
} | ||
}; | ||
|
||
let bytes_to_copy = cmp::min(buf.len(), data.len()); | ||
buf[..bytes_to_copy].copy_from_slice(&data[..bytes_to_copy]); | ||
|
||
if data.len() > bytes_to_copy { | ||
data.drain(..bytes_to_copy); | ||
self.get_mut().read_pending_bytes = Some(data); | ||
} | ||
|
||
Poll::Ready(Ok(bytes_to_copy)) | ||
} | ||
} | ||
|
||
#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] | ||
impl AsyncWrite for WebSocket { | ||
fn poll_write( | ||
mut self: Pin<&mut Self>, | ||
cx: &mut Context<'_>, | ||
buf: &[u8], | ||
) -> Poll<io::Result<usize>> { | ||
// try flushing preemptively | ||
let _ = AsyncWrite::poll_flush(self.as_mut(), cx); | ||
|
||
// make sure sink is ready to send | ||
try_in_poll_io!(ready!(self.as_mut().poll_ready(cx))); | ||
|
||
// actually submit new item | ||
try_in_poll_io!(self.start_send(WebSocketMessage::Bytes(buf.to_vec()))); | ||
// ^ if no error occurred, message is accepted and queued when calling `start_send` | ||
// (i.e.: `to_vec` is called only once) | ||
|
||
Poll::Ready(Ok(buf.len())) | ||
} | ||
|
||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { | ||
let res = ready!(Sink::poll_flush(self, cx)); | ||
Poll::Ready(ws_result_to_io_result(res)) | ||
} | ||
|
||
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { | ||
let res = ready!(Sink::poll_close(self, cx)); | ||
Poll::Ready(ws_result_to_io_result(res)) | ||
} | ||
} | ||
|
||
fn ws_result_to_io_result(res: Result<(), WebSocketError>) -> io::Result<()> { | ||
match res { | ||
Ok(()) => Ok(()), | ||
Err(WebSocketError::ConnectionClose(_)) => Ok(()), | ||
Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)), | ||
} | ||
} | ||
|
||
#[cfg(test)] | ||
mod tests { | ||
use super::*; | ||
use futures::{AsyncReadExt, AsyncWriteExt, StreamExt}; | ||
use wasm_bindgen_futures::spawn_local; | ||
use wasm_bindgen_test::*; | ||
|
||
wasm_bindgen_test_configure!(run_in_browser); | ||
|
||
#[wasm_bindgen_test] | ||
async fn check_read_write() { | ||
let ws_echo_server_url = | ||
option_env!("WS_ECHO_SERVER_URL").expect("Did you set WS_ECHO_SERVER_URL?"); | ||
|
||
let mut ws = WebSocket::open(ws_echo_server_url).unwrap(); | ||
|
||
// ignore first message | ||
// the echo-server uses it to send it's info in the first message | ||
let _ = ws.next().await.unwrap(); | ||
|
||
let (mut reader, mut writer) = AsyncReadExt::split(ws); | ||
|
||
spawn_local(async move { | ||
writer.write_all(b"test 1").await.unwrap(); | ||
writer.write_all(b"test 2").await.unwrap(); | ||
}); | ||
|
||
spawn_local(async move { | ||
let mut buf = [0u8; 6]; | ||
reader.read_exact(&mut buf).await.unwrap(); | ||
assert_eq!(&buf, b"test 1"); | ||
reader.read_exact(&mut buf).await.unwrap(); | ||
assert_eq!(&buf, b"test 2"); | ||
}); | ||
} | ||
|
||
#[wasm_bindgen_test] | ||
async fn with_pending_bytes() { | ||
let ws_echo_server_url = | ||
option_env!("WS_ECHO_SERVER_URL").expect("Did you set WS_ECHO_SERVER_URL?"); | ||
|
||
let mut ws = WebSocket::open(ws_echo_server_url).unwrap(); | ||
|
||
// ignore first message | ||
// the echo-server uses it to send it's info in the first message | ||
let _ = ws.next().await.unwrap(); | ||
|
||
ws.write_all(b"1234567890").await.unwrap(); | ||
|
||
let mut buf = [0u8; 5]; | ||
|
||
ws.read_exact(&mut buf).await.unwrap(); | ||
assert_eq!(&buf, b"12345"); | ||
assert!(ws.has_pending_bytes()); | ||
|
||
ws.read_exact(&mut buf).await.unwrap(); | ||
assert_eq!(&buf, b"67890"); | ||
assert!(!ws.has_pending_bytes()); | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I figured it was best to provide this as an optional feature so that
futures-io
is not added in the dependency graph of all dependents