Initial Dever 0.0.1

test/rust/performance/websocket/peer.rs

@@ -0,0 +1,120 @@
+// 独立协议对端，不调用 Dever HTTP/WebSocket；绑到另一个 CPU，避免双方争抢工作线程。
+use futures_util::{SinkExt, StreamExt};
+use std::{
+    io::{BufRead, BufReader, Write},
+    process::{Child, Command, Stdio},
+};
+use tokio::net::{TcpListener, TcpStream};
+use tokio_tungstenite::{
+    WebSocketStream,
+    tungstenite::{Message, protocol::WebSocketConfig},
+};
+
+pub fn config() -> WebSocketConfig {
+    WebSocketConfig::default()
+        .read_buffer_size(4096)
+        .write_buffer_size(0)
+        .max_write_buffer_size(1048576 + 4096)
+        .max_message_size(Some(1048576))
+        .max_frame_size(Some(1048576))
+}
+
+pub struct Peer {
+    child: Child,
+    pub port: u16,
+}
+
+impl Peer {
+    pub fn start(connections: usize) -> Self {
+        let cpu = std::env::var("DEVER_WS_PEER_CPU").expect("set a separate peer CPU");
+        let executable = std::env::var_os("DEVER_WS_PEER_BINARY")
+            .map(std::path::PathBuf::from)
+            .unwrap_or_else(|| std::env::current_exe().unwrap());
+        let mut child = Command::new("taskset")
+            .args(["-c", &cpu])
+            .arg(executable)
+            .arg("--peer")
+            .env("DEVER_WS_PEER_CONNECTIONS", connections.to_string())
+            .stdout(Stdio::piped())
+            .spawn()
+            .unwrap();
+        let mut ready = String::new();
+        BufReader::new(child.stdout.take().unwrap())
+            .read_line(&mut ready)
+            .unwrap();
+        let port = ready.trim().parse().expect("peer did not report its port");
+        Self { child, port }
+    }
+
+    pub fn finish(&mut self) {
+        assert!(self.child.wait().unwrap().success());
+    }
+}
+
+impl Drop for Peer {
+    fn drop(&mut self) {
+        // 断言失败也回收本次测试创建的子进程。
+        let _ = self.child.kill();
+        let _ = self.child.wait();
+    }
+}
+
+async fn serve(mut socket: WebSocketStream<TcpStream>) {
+    while let Some(command) = socket.next().await {
+        let command = command.unwrap();
+        if command.is_close() {
+            socket.flush().await.unwrap();
+            return;
+        }
+        let (mode, count) = command.to_text().unwrap().split_once(':').unwrap();
+        let count: usize = count.parse().unwrap();
+        assert!(mode == "batch" || mode == "roundtrip");
+        let expected = socket.next().await.unwrap().unwrap();
+        socket.send(Message::Text("ready".into())).await.unwrap();
+        for _ in 0..count {
+            let actual = socket.next().await.unwrap().unwrap();
+            assert_eq!(actual, expected);
+            if mode == "roundtrip" {
+                socket.send(actual).await.unwrap();
+            }
+        }
+        if mode == "batch" {
+            socket.send(expected).await.unwrap();
+        }
+    }
+}
+
+pub fn run() {
+    let connections: usize = std::env::var("DEVER_WS_PEER_CONNECTIONS")
+        .expect("set expected peer connections")
+        .parse()
+        .unwrap();
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+    runtime.block_on(async {
+        tokio::time::timeout(std::time::Duration::from_secs(120), async {
+            let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
+            println!("{}", listener.local_addr().unwrap().port());
+            std::io::stdout().flush().unwrap();
+            let mut sessions = Vec::new();
+            for _ in 0..connections {
+                let (stream, _) = listener.accept().await.unwrap();
+                stream.set_nodelay(true).unwrap();
+                sessions.push(tokio::spawn(async move {
+                    let socket =
+                        tokio_tungstenite::accept_async_with_config(stream, Some(config()))
+                            .await
+                            .unwrap();
+                    serve(socket).await;
+                }));
+            }
+            for session in sessions {
+                session.await.unwrap();
+            }
+        })
+        .await
+        .expect("protocol peer timed out");
+    });
+}