74 lines
1.7 KiB
Rust
74 lines
1.7 KiB
Rust
// src/main.rs
|
|
#[macro_use]
|
|
extern crate rocket;
|
|
|
|
use rocket::response::stream::{Event, EventStream};
|
|
use rocket::serde::json::Json;
|
|
use rocket::{Build, Rocket};
|
|
use serde::Deserialize;
|
|
use std::sync::Arc;
|
|
use tokio::sync::broadcast;
|
|
|
|
/// ---------- shared broadcaster ----------
|
|
struct ChatBroadcaster {
|
|
sender: broadcast::Sender<String>,
|
|
}
|
|
|
|
impl ChatBroadcaster {
|
|
fn new(buffer_size: usize) -> Self {
|
|
let (sender, _rx) = broadcast::channel::<String>(buffer_size);
|
|
Self { sender }
|
|
}
|
|
|
|
async fn publish(&self, msg: String) {
|
|
let _ = self.sender.send(msg);
|
|
}
|
|
|
|
fn subscribe(&self) -> broadcast::Receiver<String> {
|
|
self.sender.subscribe()
|
|
}
|
|
}
|
|
|
|
/// ---------- Rocket routes ----------
|
|
#[derive(Debug, Deserialize)]
|
|
struct ChatMsg {
|
|
text: String,
|
|
}
|
|
|
|
#[post("/chat", format = "json", data = "<msg>")]
|
|
async fn post_message(
|
|
msg: Json<ChatMsg>,
|
|
chat: &rocket::State<Arc<ChatBroadcaster>>,
|
|
) -> &'static str {
|
|
let text = msg.text.clone();
|
|
chat.publish(text).await;
|
|
"Message sent"
|
|
}
|
|
|
|
#[get("/events")]
|
|
async fn event_stream(chat: &rocket::State<Arc<ChatBroadcaster>>) -> EventStream![] {
|
|
let mut rx = chat.subscribe();
|
|
|
|
EventStream! {
|
|
loop {
|
|
match rx.recv().await {
|
|
Ok(msg) => yield Event::data(msg),
|
|
Err(broadcast::error::RecvError::Lagged(_)) => {
|
|
yield Event::comment("lagged");
|
|
}
|
|
Err(broadcast::error::RecvError::Closed) => break,
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// ---------- launch ----------
|
|
#[launch]
|
|
fn rocket() -> Rocket<Build> {
|
|
let chat = Arc::new(ChatBroadcaster::new(32));
|
|
|
|
rocket::build()
|
|
.manage(chat)
|
|
.mount("/", routes![post_message, event_stream])
|
|
}
|