Initial Dever 0.0.1

test/rust/performance/websocket/layers.rs

@@ -0,0 +1,95 @@
+//! 仅定位成功收发的同步成本，不作为完整资源合同的 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();
+    }
+}