500 lines
18 KiB
Rust
500 lines
18 KiB
Rust
use crate::proto::{self, EnvelopedMessage, MessageStream, RequestMessage};
|
|
use anyhow::{anyhow, Result};
|
|
use async_lock::{Mutex, RwLock};
|
|
use futures::{
|
|
future::{BoxFuture, Either},
|
|
AsyncRead, AsyncWrite, FutureExt,
|
|
};
|
|
use postage::{
|
|
barrier, mpsc, oneshot,
|
|
prelude::{Sink, Stream},
|
|
};
|
|
use std::{
|
|
any::TypeId,
|
|
collections::{HashMap, HashSet},
|
|
fmt,
|
|
future::Future,
|
|
pin::Pin,
|
|
sync::{
|
|
atomic::{self, AtomicU32},
|
|
Arc,
|
|
},
|
|
};
|
|
|
|
type BoxedWriter = Pin<Box<dyn AsyncWrite + 'static + Send>>;
|
|
type BoxedReader = Pin<Box<dyn AsyncRead + 'static + Send>>;
|
|
|
|
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
|
|
pub struct ConnectionId(u32);
|
|
|
|
struct Connection {
|
|
writer: Mutex<MessageStream<BoxedWriter>>,
|
|
reader: Mutex<MessageStream<BoxedReader>>,
|
|
response_channels: Mutex<HashMap<u32, oneshot::Sender<proto::Envelope>>>,
|
|
next_message_id: AtomicU32,
|
|
}
|
|
|
|
type MessageHandler = Box<
|
|
dyn Send + Sync + Fn(&mut Option<proto::Envelope>, ConnectionId) -> Option<BoxFuture<bool>>,
|
|
>;
|
|
|
|
pub struct TypedEnvelope<T> {
|
|
pub id: u32,
|
|
pub connection_id: ConnectionId,
|
|
pub payload: T,
|
|
}
|
|
|
|
pub struct Peer {
|
|
connections: RwLock<HashMap<ConnectionId, Arc<Connection>>>,
|
|
connection_close_barriers: RwLock<HashMap<ConnectionId, barrier::Sender>>,
|
|
message_handlers: RwLock<Vec<MessageHandler>>,
|
|
handler_types: Mutex<HashSet<TypeId>>,
|
|
next_connection_id: AtomicU32,
|
|
}
|
|
|
|
impl Peer {
|
|
pub fn new() -> Arc<Self> {
|
|
Arc::new(Self {
|
|
connections: Default::default(),
|
|
connection_close_barriers: Default::default(),
|
|
message_handlers: Default::default(),
|
|
handler_types: Default::default(),
|
|
next_connection_id: Default::default(),
|
|
})
|
|
}
|
|
|
|
pub async fn add_message_handler<T: EnvelopedMessage>(
|
|
&self,
|
|
) -> mpsc::Receiver<TypedEnvelope<T>> {
|
|
if !self.handler_types.lock().await.insert(TypeId::of::<T>()) {
|
|
panic!("duplicate handler type");
|
|
}
|
|
|
|
let (tx, rx) = mpsc::channel(256);
|
|
self.message_handlers
|
|
.write()
|
|
.await
|
|
.push(Box::new(move |envelope, connection_id| {
|
|
if envelope.as_ref().map_or(false, T::matches_envelope) {
|
|
let envelope = Option::take(envelope).unwrap();
|
|
let mut tx = tx.clone();
|
|
Some(
|
|
async move {
|
|
tx.send(TypedEnvelope {
|
|
id: envelope.id,
|
|
connection_id,
|
|
payload: T::from_envelope(envelope).unwrap(),
|
|
})
|
|
.await
|
|
.is_err()
|
|
}
|
|
.boxed(),
|
|
)
|
|
} else {
|
|
None
|
|
}
|
|
}));
|
|
rx
|
|
}
|
|
|
|
pub async fn add_connection<Conn>(self: &Arc<Self>, conn: Conn) -> ConnectionId
|
|
where
|
|
Conn: Clone + AsyncRead + AsyncWrite + Unpin + Send + 'static,
|
|
{
|
|
let connection_id = ConnectionId(
|
|
self.next_connection_id
|
|
.fetch_add(1, atomic::Ordering::SeqCst),
|
|
);
|
|
self.connections.write().await.insert(
|
|
connection_id,
|
|
Arc::new(Connection {
|
|
reader: Mutex::new(MessageStream::new(Box::pin(conn.clone()))),
|
|
writer: Mutex::new(MessageStream::new(Box::pin(conn.clone()))),
|
|
response_channels: Default::default(),
|
|
next_message_id: Default::default(),
|
|
}),
|
|
);
|
|
connection_id
|
|
}
|
|
|
|
pub async fn disconnect(&self, connection_id: ConnectionId) {
|
|
self.connections.write().await.remove(&connection_id);
|
|
self.connection_close_barriers
|
|
.write()
|
|
.await
|
|
.remove(&connection_id);
|
|
}
|
|
|
|
pub fn handle_messages(
|
|
self: &Arc<Self>,
|
|
connection_id: ConnectionId,
|
|
) -> impl Future<Output = Result<()>> + 'static {
|
|
let (close_tx, mut close_rx) = barrier::channel();
|
|
let this = self.clone();
|
|
async move {
|
|
this.connection_close_barriers
|
|
.write()
|
|
.await
|
|
.insert(connection_id, close_tx);
|
|
let connection = this.connection(connection_id).await?;
|
|
let closed = close_rx.recv();
|
|
futures::pin_mut!(closed);
|
|
|
|
loop {
|
|
let mut reader = connection.reader.lock().await;
|
|
let read_message = reader.read_message();
|
|
futures::pin_mut!(read_message);
|
|
|
|
match futures::future::select(read_message, &mut closed).await {
|
|
Either::Left((Ok(incoming), _)) => {
|
|
if let Some(responding_to) = incoming.responding_to {
|
|
let channel = connection
|
|
.response_channels
|
|
.lock()
|
|
.await
|
|
.remove(&responding_to);
|
|
if let Some(mut tx) = channel {
|
|
tx.send(incoming).await.ok();
|
|
} else {
|
|
log::warn!(
|
|
"received RPC response to unknown request {}",
|
|
responding_to
|
|
);
|
|
}
|
|
} else {
|
|
let mut envelope = Some(incoming);
|
|
let mut handler_index = None;
|
|
let mut handler_was_dropped = false;
|
|
for (i, handler) in
|
|
this.message_handlers.read().await.iter().enumerate()
|
|
{
|
|
if let Some(future) = handler(&mut envelope, connection_id) {
|
|
handler_was_dropped = future.await;
|
|
handler_index = Some(i);
|
|
break;
|
|
}
|
|
}
|
|
|
|
if let Some(handler_index) = handler_index {
|
|
if handler_was_dropped {
|
|
drop(this.message_handlers.write().await.remove(handler_index));
|
|
}
|
|
} else {
|
|
log::warn!("unhandled message: {:?}", envelope.unwrap().payload);
|
|
}
|
|
}
|
|
}
|
|
Either::Left((Err(error), _)) => {
|
|
log::warn!("received invalid RPC message: {}", error);
|
|
Err(error)?;
|
|
}
|
|
Either::Right(_) => return Ok(()),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn receive<M: EnvelopedMessage>(
|
|
self: &Arc<Self>,
|
|
connection_id: ConnectionId,
|
|
) -> Result<TypedEnvelope<M>> {
|
|
let connection = self.connection(connection_id).await?;
|
|
let envelope = connection.reader.lock().await.read_message().await?;
|
|
let id = envelope.id;
|
|
let payload =
|
|
M::from_envelope(envelope).ok_or_else(|| anyhow!("unexpected message type"))?;
|
|
Ok(TypedEnvelope {
|
|
id,
|
|
connection_id,
|
|
payload,
|
|
})
|
|
}
|
|
|
|
pub fn request<T: RequestMessage>(
|
|
self: &Arc<Self>,
|
|
connection_id: ConnectionId,
|
|
req: T,
|
|
) -> impl Future<Output = Result<T::Response>> {
|
|
let this = self.clone();
|
|
let (tx, mut rx) = oneshot::channel();
|
|
async move {
|
|
let connection = this.connection(connection_id).await?;
|
|
let message_id = connection
|
|
.next_message_id
|
|
.fetch_add(1, atomic::Ordering::SeqCst);
|
|
connection
|
|
.response_channels
|
|
.lock()
|
|
.await
|
|
.insert(message_id, tx);
|
|
connection
|
|
.writer
|
|
.lock()
|
|
.await
|
|
.write_message(&req.into_envelope(message_id, None))
|
|
.await?;
|
|
let response = rx
|
|
.recv()
|
|
.await
|
|
.expect("response channel was unexpectedly dropped");
|
|
T::Response::from_envelope(response)
|
|
.ok_or_else(|| anyhow!("received response of the wrong type"))
|
|
}
|
|
}
|
|
|
|
pub fn send<T: EnvelopedMessage>(
|
|
self: &Arc<Self>,
|
|
connection_id: ConnectionId,
|
|
message: T,
|
|
) -> impl Future<Output = Result<()>> {
|
|
let this = self.clone();
|
|
async move {
|
|
let connection = this.connection(connection_id).await?;
|
|
let message_id = connection
|
|
.next_message_id
|
|
.fetch_add(1, atomic::Ordering::SeqCst);
|
|
connection
|
|
.writer
|
|
.lock()
|
|
.await
|
|
.write_message(&message.into_envelope(message_id, None))
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
pub fn respond<T: RequestMessage>(
|
|
self: &Arc<Self>,
|
|
request: TypedEnvelope<T>,
|
|
response: T::Response,
|
|
) -> impl Future<Output = Result<()>> {
|
|
let this = self.clone();
|
|
async move {
|
|
let connection = this.connection(request.connection_id).await?;
|
|
let message_id = connection
|
|
.next_message_id
|
|
.fetch_add(1, atomic::Ordering::SeqCst);
|
|
connection
|
|
.writer
|
|
.lock()
|
|
.await
|
|
.write_message(&response.into_envelope(message_id, Some(request.id)))
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
async fn connection(&self, id: ConnectionId) -> Result<Arc<Connection>> {
|
|
Ok(self
|
|
.connections
|
|
.read()
|
|
.await
|
|
.get(&id)
|
|
.ok_or_else(|| anyhow!("unknown connection: {}", id.0))?
|
|
.clone())
|
|
}
|
|
}
|
|
|
|
impl fmt::Display for ConnectionId {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
self.0.fmt(f)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use smol::{
|
|
io::AsyncWriteExt,
|
|
net::unix::{UnixListener, UnixStream},
|
|
};
|
|
use std::io;
|
|
use tempdir::TempDir;
|
|
|
|
#[test]
|
|
fn test_request_response() {
|
|
smol::block_on(async move {
|
|
// create socket
|
|
let socket_dir_path = TempDir::new("test-request-response").unwrap();
|
|
let socket_path = socket_dir_path.path().join("test.sock");
|
|
let listener = UnixListener::bind(&socket_path).unwrap();
|
|
|
|
// create 2 clients connected to 1 server
|
|
let server = Peer::new();
|
|
let client1 = Peer::new();
|
|
let client2 = Peer::new();
|
|
let client1_conn_id = client1
|
|
.add_connection(UnixStream::connect(&socket_path).await.unwrap())
|
|
.await;
|
|
let client2_conn_id = client2
|
|
.add_connection(UnixStream::connect(&socket_path).await.unwrap())
|
|
.await;
|
|
let server_conn_id1 = server
|
|
.add_connection(listener.accept().await.unwrap().0)
|
|
.await;
|
|
let server_conn_id2 = server
|
|
.add_connection(listener.accept().await.unwrap().0)
|
|
.await;
|
|
smol::spawn(client1.handle_messages(client1_conn_id)).detach();
|
|
smol::spawn(client2.handle_messages(client2_conn_id)).detach();
|
|
smol::spawn(server.handle_messages(server_conn_id1)).detach();
|
|
smol::spawn(server.handle_messages(server_conn_id2)).detach();
|
|
|
|
// define the expected requests and responses
|
|
let request1 = proto::Auth {
|
|
user_id: 1,
|
|
access_token: "token-1".to_string(),
|
|
};
|
|
let response1 = proto::AuthResponse {
|
|
credentials_valid: true,
|
|
};
|
|
let request2 = proto::Auth {
|
|
user_id: 2,
|
|
access_token: "token-2".to_string(),
|
|
};
|
|
let response2 = proto::AuthResponse {
|
|
credentials_valid: false,
|
|
};
|
|
let request3 = proto::OpenBuffer {
|
|
worktree_id: 102,
|
|
path: "path/two".to_string(),
|
|
};
|
|
let response3 = proto::OpenBufferResponse {
|
|
buffer: Some(proto::Buffer {
|
|
id: 1001,
|
|
path: "path/two".to_string(),
|
|
content: "path/two content".to_string(),
|
|
history: vec![],
|
|
}),
|
|
};
|
|
let request4 = proto::OpenBuffer {
|
|
worktree_id: 101,
|
|
path: "path/one".to_string(),
|
|
};
|
|
let response4 = proto::OpenBufferResponse {
|
|
buffer: Some(proto::Buffer {
|
|
id: 1002,
|
|
path: "path/one".to_string(),
|
|
content: "path/one content".to_string(),
|
|
history: vec![],
|
|
}),
|
|
};
|
|
|
|
// on the server, respond to two requests for each client
|
|
let mut open_buffer_rx = server.add_message_handler::<proto::OpenBuffer>().await;
|
|
let mut auth_rx = server.add_message_handler::<proto::Auth>().await;
|
|
let (mut server_done_tx, mut server_done_rx) = oneshot::channel::<()>();
|
|
smol::spawn({
|
|
let request1 = request1.clone();
|
|
let request2 = request2.clone();
|
|
let request3 = request3.clone();
|
|
let request4 = request4.clone();
|
|
let response1 = response1.clone();
|
|
let response2 = response2.clone();
|
|
let response3 = response3.clone();
|
|
let response4 = response4.clone();
|
|
async move {
|
|
let msg = auth_rx.recv().await.unwrap();
|
|
assert_eq!(msg.payload, request1);
|
|
server.respond(msg, response1.clone()).await.unwrap();
|
|
|
|
let msg = auth_rx.recv().await.unwrap();
|
|
assert_eq!(msg.payload, request2.clone());
|
|
server.respond(msg, response2.clone()).await.unwrap();
|
|
|
|
let msg = open_buffer_rx.recv().await.unwrap();
|
|
assert_eq!(msg.payload, request3.clone());
|
|
server.respond(msg, response3.clone()).await.unwrap();
|
|
|
|
let msg = open_buffer_rx.recv().await.unwrap();
|
|
assert_eq!(msg.payload, request4.clone());
|
|
server.respond(msg, response4.clone()).await.unwrap();
|
|
|
|
server_done_tx.send(()).await.unwrap();
|
|
}
|
|
})
|
|
.detach();
|
|
|
|
assert_eq!(
|
|
client1.request(client1_conn_id, request1).await.unwrap(),
|
|
response1
|
|
);
|
|
assert_eq!(
|
|
client2.request(client2_conn_id, request2).await.unwrap(),
|
|
response2
|
|
);
|
|
assert_eq!(
|
|
client2.request(client2_conn_id, request3).await.unwrap(),
|
|
response3
|
|
);
|
|
assert_eq!(
|
|
client1.request(client1_conn_id, request4).await.unwrap(),
|
|
response4
|
|
);
|
|
|
|
client1.disconnect(client1_conn_id).await;
|
|
client2.disconnect(client1_conn_id).await;
|
|
|
|
server_done_rx.recv().await.unwrap();
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_disconnect() {
|
|
smol::block_on(async move {
|
|
let socket_dir_path = TempDir::new("drop-client").unwrap();
|
|
let socket_path = socket_dir_path.path().join(".sock");
|
|
let listener = UnixListener::bind(&socket_path).unwrap();
|
|
let client_conn = UnixStream::connect(&socket_path).await.unwrap();
|
|
let (mut server_conn, _) = listener.accept().await.unwrap();
|
|
|
|
let client = Peer::new();
|
|
let connection_id = client.add_connection(client_conn).await;
|
|
let (mut incoming_messages_ended_tx, mut incoming_messages_ended_rx) =
|
|
barrier::channel();
|
|
let handle_messages = client.handle_messages(connection_id);
|
|
smol::spawn(async move {
|
|
handle_messages.await.ok();
|
|
incoming_messages_ended_tx.send(()).await.unwrap();
|
|
})
|
|
.detach();
|
|
client.disconnect(connection_id).await;
|
|
|
|
incoming_messages_ended_rx.recv().await;
|
|
|
|
let err = server_conn.write(&[]).await.unwrap_err();
|
|
assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_io_error() {
|
|
smol::block_on(async move {
|
|
let socket_dir_path = TempDir::new("io-error").unwrap();
|
|
let socket_path = socket_dir_path.path().join(".sock");
|
|
let _listener = UnixListener::bind(&socket_path).unwrap();
|
|
let mut client_conn = UnixStream::connect(&socket_path).await.unwrap();
|
|
client_conn.close().await.unwrap();
|
|
|
|
let client = Peer::new();
|
|
let connection_id = client.add_connection(client_conn).await;
|
|
smol::spawn(client.handle_messages(connection_id)).detach();
|
|
|
|
let err = client
|
|
.request(
|
|
connection_id,
|
|
proto::Auth {
|
|
user_id: 42,
|
|
access_token: "token".to_string(),
|
|
},
|
|
)
|
|
.await
|
|
.unwrap_err();
|
|
assert_eq!(
|
|
err.downcast_ref::<io::Error>().unwrap().kind(),
|
|
io::ErrorKind::BrokenPipe
|
|
);
|
|
});
|
|
}
|
|
}
|