oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Introduce a concrete Conn type for peer's websocket connection

Browse Source 
This is mostly to set us up to test the rpc::Client's reconnect
logic.

There are multiple ways that the `rpc::Client` may establish
its websocket connection: (SSL in production, plain TCP during
local development, and using an in-memory connection for tests).
Now we can represent all of those connections using a common type.

Also, several long methods no longer need to be generic, which
is good for compile time.
This commit is contained in:
Max Brunsfeld
2021-09-08 18:14:47 -07:00
parent c3e29e0a2d
commit b6eac57f63
9 changed files with 196 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files
zrpc/src/conn.rs
+54
View File
@@ -0,0 +1,54 @@
use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
use futures::{SinkExt as _, StreamExt as _};
pub struct Conn {
pub(crate) tx:
Box<dyn 'static + Send + Unpin + futures::Sink<WebSocketMessage, Error = WebSocketError>>,
pub(crate) rx: Box<
dyn 'static
+ Send
+ Unpin
+ futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>,
>,
}
impl Conn {
pub fn new<S>(stream: S) -> Self
where
S: 'static
+ Send
+ Unpin
+ futures::Sink<WebSocketMessage, Error = WebSocketError>
+ futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>,
{
let (tx, rx) = stream.split();
Self {
tx: Box::new(tx),
rx: Box::new(rx),
}
}
pub async fn send(&mut self, message: WebSocketMessage) -> Result<(), WebSocketError> {
self.tx.send(message).await
}
#[cfg(any(test, feature = "test-support"))]
pub fn in_memory() -> (Self, Self) {
use futures::SinkExt as _;
use futures::StreamExt as _;
use std::io::{Error, ErrorKind};
let (a_tx, a_rx) = futures::channel::mpsc::unbounded::<WebSocketMessage>();
let (b_tx, b_rx) = futures::channel::mpsc::unbounded::<WebSocketMessage>();
(
Self {
tx: Box::new(a_tx.sink_map_err(|e| Error::new(ErrorKind::Other, e).into())),
rx: Box::new(b_rx.map(Ok)),
},
Self {
tx: Box::new(b_tx.sink_map_err(|e| Error::new(ErrorKind::Other, e).into())),
rx: Box::new(a_rx.map(Ok)),
},
)
}
}
zrpc/src/lib.rs
+2 -3
View File
@@ -1,7 +1,6 @@
pub mod auth;
mod conn;
mod peer;
pub mod proto;
#[cfg(any(test, feature = "test-support"))]
pub mod test;
pub use conn::Conn;
pub use peer::*;
zrpc/src/peer.rs
+18 -24
View File
@@ -1,8 +1,8 @@
use crate::proto::{self, AnyTypedEnvelope, EnvelopedMessage, MessageStream, RequestMessage};
use super::proto::{self, AnyTypedEnvelope, EnvelopedMessage, MessageStream, RequestMessage};
use super::Conn;
use anyhow::{anyhow, Context, Result};
use async_lock::{Mutex, RwLock};
use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
use futures::{FutureExt, StreamExt};
use futures::FutureExt as _;
use postage::{
mpsc,
prelude::{Sink as _, Stream as _},
@@ -98,21 +98,14 @@ impl Peer {
})
}
pub async fn add_connection<Conn>(
pub async fn add_connection(
self: &Arc<Self>,
conn: Conn,
) -> (
ConnectionId,
impl Future<Output = anyhow::Result<()>> + Send,
mpsc::Receiver<Box<dyn AnyTypedEnvelope>>,
)
where
Conn: futures::Sink<WebSocketMessage, Error = WebSocketError>
+ futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>
+ Send
+ Unpin,
{
let (tx, rx) = conn.split();
) {
let connection_id = ConnectionId(
self.next_connection_id
.fetch_add(1, atomic::Ordering::SeqCst),
@@ -124,8 +117,8 @@ impl Peer {
next_message_id: Default::default(),
response_channels: Default::default(),
};
let mut writer = MessageStream::new(tx);
let mut reader = MessageStream::new(rx);
let mut writer = MessageStream::new(conn.tx);
let mut reader = MessageStream::new(conn.rx);
let this = self.clone();
let response_channels = connection.response_channels.clone();
@@ -347,7 +340,9 @@ impl Peer {
#[cfg(test)]
mod tests {
use super::*;
use crate::{test, TypedEnvelope};
use crate::TypedEnvelope;
use async_tungstenite::tungstenite::Message as WebSocketMessage;
use futures::StreamExt as _;
#[test]
fn test_request_response() {
@@ -357,12 +352,12 @@ mod tests {
let client1 = Peer::new();
let client2 = Peer::new();
let (client1_to_server_conn, server_to_client_1_conn) = test::Channel::bidirectional();
let (client1_to_server_conn, server_to_client_1_conn) = Conn::in_memory();
let (client1_conn_id, io_task1, _) =
client1.add_connection(client1_to_server_conn).await;
let (_, io_task2, incoming1) = server.add_connection(server_to_client_1_conn).await;
let (client2_to_server_conn, server_to_client_2_conn) = test::Channel::bidirectional();
let (client2_to_server_conn, server_to_client_2_conn) = Conn::in_memory();
let (client2_conn_id, io_task3, _) =
client2.add_connection(client2_to_server_conn).await;
let (_, io_task4, incoming2) = server.add_connection(server_to_client_2_conn).await;
@@ -497,7 +492,7 @@ mod tests {
#[test]
fn test_disconnect() {
smol::block_on(async move {
let (client_conn, mut server_conn) = test::Channel::bidirectional();
let (client_conn, mut server_conn) = Conn::in_memory();
let client = Peer::new();
let (connection_id, io_handler, mut incoming) =
@@ -521,18 +516,17 @@ mod tests {
io_ended_rx.recv().await;
messages_ended_rx.recv().await;
assert!(
futures::SinkExt::send(&mut server_conn, WebSocketMessage::Binary(vec![]))
.await
.is_err()
);
assert!(server_conn
.send(WebSocketMessage::Binary(vec![]))
.await
.is_err());
});
}
#[test]
fn test_io_error() {
smol::block_on(async move {
let (client_conn, server_conn) = test::Channel::bidirectional();
let (client_conn, server_conn) = Conn::in_memory();
drop(server_conn);
let client = Peer::new();
zrpc/src/proto.rs
-27
View File
@@ -247,30 +247,3 @@ impl From<SystemTime> for Timestamp {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test;
#[test]
fn test_round_trip_message() {
smol::block_on(async {
let stream = test::Channel::new();
let message1 = Ping { id: 5 }.into_envelope(3, None, None);
let message2 = OpenBuffer {
worktree_id: 0,
path: "some/path".to_string(),
}
.into_envelope(5, None, None);
let mut message_stream = MessageStream::new(stream);
message_stream.write_message(&message1).await.unwrap();
message_stream.write_message(&message2).await.unwrap();
let decoded_message1 = message_stream.read_message().await.unwrap();
let decoded_message2 = message_stream.read_message().await.unwrap();
assert_eq!(decoded_message1, message1);
assert_eq!(decoded_message2, message2);
});
}
}
zrpc/src/test.rs
-64
View File
@@ -1,64 +0,0 @@
use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
use std::{
io,
pin::Pin,
task::{Context, Poll},
};
pub struct Channel {
tx: futures::channel::mpsc::UnboundedSender<WebSocketMessage>,
rx: futures::channel::mpsc::UnboundedReceiver<WebSocketMessage>,
}
impl Channel {
pub fn new() -> Self {
let (tx, rx) = futures::channel::mpsc::unbounded();
Self { tx, rx }
}
pub fn bidirectional() -> (Self, Self) {
let (a_tx, a_rx) = futures::channel::mpsc::unbounded();
let (b_tx, b_rx) = futures::channel::mpsc::unbounded();
let a = Self { tx: a_tx, rx: b_rx };
let b = Self { tx: b_tx, rx: a_rx };
(a, b)
}
}
impl futures::Sink<WebSocketMessage> for Channel {
type Error = WebSocketError;
fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.tx)
.poll_ready(cx)
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
}
fn start_send(mut self: Pin<&mut Self>, item: WebSocketMessage) -> Result<(), Self::Error> {
Pin::new(&mut self.tx)
.start_send(item)
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.tx)
.poll_flush(cx)
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
}
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.tx)
.poll_close(cx)
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
}
}
impl futures::Stream for Channel {
type Item = Result<WebSocketMessage, WebSocketError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.rx)
.poll_next(cx)
.map(|i| i.map(|i| Ok(i)))
}
}