96 lines
2.9 KiB
Rust
96 lines
2.9 KiB
Rust
//! 仅定位成功收发的同步成本，不作为完整资源合同的 Rust 参照。
|
|
|
|
use futures_util::{Sink, Stream};
|
|
use std::{
|
|
future::{Future, poll_fn},
|
|
pin::Pin,
|
|
sync::Mutex,
|
|
task::{Context, Poll, ready},
|
|
};
|
|
use tokio::io::{AsyncRead, AsyncWrite};
|
|
use tokio_tungstenite::{
|
|
WebSocketStream,
|
|
tungstenite::{Error, Message},
|
|
};
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
pub trait Transport: AsyncRead + AsyncWrite + Unpin + Send {}
|
|
impl<Value: AsyncRead + AsyncWrite + Unpin + Send> Transport for Value {}
|
|
pub type Boxed = Box<dyn Transport>;
|
|
type Wire = WebSocketStream<Boxed>;
|
|
|
|
pub struct Shared<const CANCEL: bool> {
|
|
wire: Mutex<Wire>,
|
|
read: tokio::sync::Mutex<()>,
|
|
write: tokio::sync::Mutex<()>,
|
|
closed: CancellationToken,
|
|
closing: CancellationToken,
|
|
}
|
|
|
|
impl<const CANCEL: bool> Shared<CANCEL> {
|
|
pub fn new(wire: Wire) -> Self {
|
|
Self {
|
|
wire: Mutex::new(wire),
|
|
read: tokio::sync::Mutex::new(()),
|
|
write: tokio::sync::Mutex::new(()),
|
|
closed: CancellationToken::new(),
|
|
closing: CancellationToken::new(),
|
|
}
|
|
}
|
|
|
|
fn poll<Value>(
|
|
&self,
|
|
context: &mut Context<'_>,
|
|
call: impl FnOnce(Pin<&mut Wire>, &mut Context<'_>) -> Poll<Value>,
|
|
) -> Poll<Value> {
|
|
let mut wire = self.wire.lock().unwrap();
|
|
call(Pin::new(&mut *wire), context)
|
|
}
|
|
|
|
async fn active<Value>(
|
|
&self,
|
|
work: impl Future<Output = Result<Value, Error>>,
|
|
) -> Result<Value, Error> {
|
|
if CANCEL {
|
|
tokio::select! {
|
|
biased;
|
|
() = self.closed.cancelled() => Err(Error::ConnectionClosed),
|
|
() = self.closing.cancelled() => Err(Error::ConnectionClosed),
|
|
result = work => result,
|
|
}
|
|
} else {
|
|
work.await
|
|
}
|
|
}
|
|
|
|
pub async fn send(&self, frame: Message) -> Result<(), Error> {
|
|
self.active(async {
|
|
let _direction = self.write.lock().await;
|
|
let mut frame = Some(frame);
|
|
poll_fn(|context| {
|
|
self.poll(context, |mut wire, context| {
|
|
ready!(wire.as_mut().poll_ready(context))?;
|
|
Poll::Ready(wire.start_send(frame.take().unwrap()))
|
|
})
|
|
})
|
|
.await?;
|
|
poll_fn(|context| self.poll(context, |wire, context| wire.poll_flush(context))).await
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub async fn receive(&self) -> Result<Option<Message>, Error> {
|
|
self.active(async {
|
|
let _direction = self.read.lock().await;
|
|
poll_fn(|context| self.poll(context, |wire, context| wire.poll_next(context)))
|
|
.await
|
|
.transpose()
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub async fn close(&mut self) {
|
|
self.wire.get_mut().unwrap().close(None).await.unwrap();
|
|
}
|
|
}
|