toml
[dependencies]
libc = "0.2.98"
tokio = { version="1.8.1", features=["sync", "time", "rt-multi-thread","macros"] }
tokio-tungstenite = "0.15.0"
tungstenite = "0.14.0"
futures-util = "0.3.15"
env_logger = "0.8.4"
log = "0.4.14"
chrono = "0.4.19"
html
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8"/>
<title>Chat Client(WebSocket)</title>
<!-- Modified from: https://github.com/actix/examples/blob/a66c05448eace8b1ea53c7495b27604e7e91281c/websockets/chat-broker/static/index.html -->
<style>
:root {
font-size: 18px;
}
input[type='text'] {
font-size: inherit;
}
#log {
width: 30em;
height: 20em;
overflow: auto;
margin: 0.5em 0;
border: 1px solid black;
}
#status {
padding: 0 0.2em;
}
#url, #text {
width: 17em;
padding: 0.5em;
}
.msg {
margin: 0;
padding: 0.25em 0.5em;
}
.msg--status {
background-color: #ffffc9;
}
.msg--message {
background-color: #d2f4ff;
}
.msg--error {
background-color: pink;
}
</style>
</head>
<body>
<h2>Chat Client</h2>
<form id="urlform">
<label for="url">输入ws地址，并点击确定</label>
<br>
<input type="text" id="url" value="ws://127.0.0.1:9007"/>
<input type="submit" value="确定"/>
</form>
<div>
<span>Status:</span>
<span id="status">断开连接</span>
<button id="connect">连接</button>
</div>
<div id="log"></div>
<form id="chatform">
<label for="text"></label>
<input type="text" id="text"/>
<input type="submit" id="send"/>
</form>
<script>
const status = document.querySelector('#status')
const connectButton = document.querySelector('#connect')
connectButton.disabled = true;
const log = document.querySelector('#log')
const textInput = document.querySelector('#text')
/** @type {WebSocket | null} */
let socket = null
let websocketUrl = "";
function logStatus(msg, type = 'status') {
log.innerHTML += `<p class="msg msg--${type}">${msg}</p>`
log.scrollTop += 1000
}
function connect() {
disconnect()
logStatus('Connecting...')
socket = new WebSocket(websocketUrl)
socket.onopen = () => {
logStatus('Connected')
updateConnectionStatus()
}
socket.onmessage = (ev) => {
logStatus('Received: ' + ev.data, 'message')
}
socket.onclose = () => {
logStatus('Disconnected')
socket = null
updateConnectionStatus()
}
}
function disconnect() {
if (socket) {
logStatus('Disconnecting...')
socket.close()
socket = null
updateConnectionStatus()
}
}
function updateConnectionStatus() {
if (socket) {
status.style.backgroundColor = 'transparent'
status.style.color = 'green'
status.textContent = `connected`
connectButton.innerHTML = 'Disconnect'
textInput.focus()
} else {
status.style.backgroundColor = 'red'
status.style.color = 'white'
status.textContent = 'disconnected'
connectButton.textContent = 'Connect'
}
}
document.querySelector('#urlform').addEventListener('submit', (event)=>{
event.preventDefault()
connectButton.disabled = false
websocketUrl = document.querySelector('#url').value
})
connectButton.addEventListener('click', () => {
if (socket) {
disconnect()
} else {
connect()
}
updateConnectionStatus()
})
document.querySelector('#chatform').addEventListener('submit', (ev) => {
ev.preventDefault()
const text = textInput.value
logStatus('Sending: ' + text)
socket.send(text)
textInput.value = ''
textInput.focus()
})
updateConnectionStatus()
</script>
</body>
</html>
main.rs
#![warn(clippy::nursery, clippy::pedantic)]
use std::collections::HashMap;
use futures_util::{stream::SplitSink, SinkExt, StreamExt};
use std::sync::Arc;
use std::time::Instant;
use tokio_tungstenite::{accept_async, WebSocketStream};
use tungstenite::Message as WsMessage;
use log::info;
use tokio::sync::RwLock;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
const WS_SERVER_ADDR: &str = "127.0.0.1:9007";
//定义日志级别
env_logger::Builder::new().filter_level(log::LevelFilter::Info).init();
let listener = tokio::net::TcpListener::bind(WS_SERVER_ADDR).await?;
info!("Listening on: {}", WS_SERVER_ADDR);
//消息通道
let (new_msg_notify_sender, _) = tokio::sync::broadcast::channel::<ChatMessage>(5000);
//全局app
let app_state = Arc::new(AppState {
new_msg_notify_sender,
//历史消息放内存，只是为了测试用
chat_history: tokio::sync::RwLock::new(Vec::new()),
});
//接受连接
while let Ok((stream, _addr)) = listener.accept().await {
tokio::spawn(accept_connection(stream, app_state.clone()));
}
Ok(())
}
//等待连接，并加入到全局用户在线列表
async fn accept_connection(stream: tokio::net::TcpStream, app_state: Arc<AppState>) -> tungstenite::Result<()> {
//发送，接收
let (ws_sender, mut ws_receiver)
= accept_async(stream).await.unwrap().split();
let mut new_msg_notify_receiver = app_state.new_msg_notify_sender.subscribe();
//客户端结构体
let mut connection_state = ConnectionState {
connection_id: gen_random_user_id(),
ws_sender,
last_heartbeat_timestamp: Instant::now(),
};
// 当有新的连接连接上，发送历史消息 -->可选
for old_chat_msg in &*app_state.chat_history.read().await {
connection_state
.ws_sender
.send(WsMessage::Text(old_chat_msg.to_string()))
.await?;
}
//定时器
let mut interval = tokio::time::interval(std::time::Duration::from_secs(1));
loop {
tokio::select! {
msg = ws_receiver.next() => {
match msg {
Some(msg) => {
let msg = msg?;
match msg {
WsMessage::Text(msg) => connection_state.handle_client_msg(msg, &app_state).await?,
WsMessage::Ping(_) => {
connection_state.ws_sender.send(WsMessage::Pong(Vec::new())).await?;
},
WsMessage::Pong(_) => {
connection_state.last_heartbeat_timestamp = Instant::now();
},
WsMessage::Close(_) => println!("用户退出"),
WsMessage::Binary(_) => unreachable!()
}
}
None => break,
}
}
// 接收新消息
Ok(msg) = new_msg_notify_receiver.recv() => {
connection_state.ws_sender.send(WsMessage::Text(msg.to_string())).await?;
}
// 定时器发送心跳
_ = interval.tick() => {
if connection_state.last_heartbeat_timestamp.elapsed().as_secs() > ConnectionState::HEARTBEAT_TIMEOUT {
log::info!("server close dead connection");
break;
}
connection_state.ws_sender.send(WsMessage::Ping(Vec::new())).await?;
}
}
}
Ok(())
}
#[derive(Clone, Debug)]
struct ChatMessage {
created_at: chrono::NaiveDateTime,
user_id: i32,
message: String,
}
impl ToString for ChatMessage {
fn to_string(&self) -> String {
format!(
"[user_id:{} {}]: {}",
self.user_id,
self.created_at.format("%H:%M:%S"),
self.message
)
}
}
//全局
struct AppState {
new_msg_notify_sender: tokio::sync::broadcast::Sender<ChatMessage>,
chat_history: tokio::sync::RwLock<Vec<ChatMessage>>,
//全局在线用户，后续可以手动踢出，异常断线剔除
//online_users: RwLock<HashMap<i32,ConnectionState>>
}
//客户端结构体
struct ConnectionState {
//用户连接id
connection_id: i32,
// 客户端最后一次心跳活跃时间
last_heartbeat_timestamp: Instant,
// 客户端的连接四元组
ws_sender: SplitSink<WebSocketStream<tokio::net::TcpStream>, WsMessage>,
}
impl ConnectionState {
const HEARTBEAT_TIMEOUT: u64 = 15;
async fn handle_client_msg(&mut self, msg: String, app_state: &Arc<AppState>) -> tungstenite::Result<()> {
//封装消息结构体
let chat_message = ChatMessage {
created_at: chrono::Local::now().naive_utc(),
user_id: self.connection_id.clone(),
message: msg.clone(),
};
app_state.new_msg_notify_sender
.send(chat_message.clone())
.unwrap();
app_state.chat_history.write().await.push(chat_message);
Ok(())
}
}
//生成一个客户端id
#[allow(clippy::cast_sign_loss)]
#[allow(clippy::cast_possible_truncation)]
fn gen_random_user_id() -> i32 {
static SRAND_INIT: std::sync::Once = std::sync::Once::new();
SRAND_INIT.call_once(|| unsafe {
libc::srand(libc::time(std::ptr::null_mut()) as u32);
});
unsafe { libc::rand() }
}
**粗体** _斜体_ [链接](http://example.com) `代码` - 列表 > 引用。你还可以使用@来通知其他用户。