121 lines
4.0 KiB
Rust
121 lines
4.0 KiB
Rust
// 独立协议对端，不调用 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");
|
|
});
|
|
}
|