fix: use redis aio
This commit is contained in:
trisua
parent
0df2cd7b56
commit
ecde5d3d46
6 changed files with 240 additions and 30 deletions
|
|
@ -15,7 +15,7 @@ serde = { version = "1.0.219", features = ["derive"] }
|
|||
tera = "1.20.0"
|
||||
tracing = "0.1.41"
|
||||
tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
|
||||
tower-http = { version = "0.6.2", features = ["trace", "fs"] }
|
||||
tower-http = { version = "0.6.2", features = ["trace", "fs", "catch-panic"] }
|
||||
axum = { version = "0.8.3", features = ["macros", "ws"] }
|
||||
tokio = { version = "1.44.2", features = ["macros", "rt-multi-thread"] }
|
||||
axum-extra = { version = "0.10.1", features = ["cookie", "multipart"] }
|
||||
|
|
@ -35,4 +35,5 @@ cf-turnstile = "0.2.0"
|
|||
contrasted = "0.1.2"
|
||||
futures-util = "0.3.31"
|
||||
|
||||
redis = { version = "0.30.0", optional = true }
|
||||
redis = { version = "0.30.0", features = ["aio", "tokio-comp"], optional = true }
|
||||
tower_governor = "0.7.0"
|
||||
|
|
|
|||
|
|
@ -11,10 +11,14 @@ pub use tetratto_core::*;
|
|||
use axum::{Extension, Router};
|
||||
use reqwest::Client;
|
||||
use tera::{Tera, Value};
|
||||
use tower_http::trace::{self, TraceLayer};
|
||||
use tower_http::{
|
||||
trace::{self, TraceLayer},
|
||||
catch_panic::CatchPanicLayer,
|
||||
};
|
||||
use tower_governor::{GovernorLayer, governor::GovernorConfigBuilder};
|
||||
use tracing::{Level, info};
|
||||
|
||||
use std::{collections::HashMap, env::var, process::exit, sync::Arc};
|
||||
use std::{collections::HashMap, env::var, net::SocketAddr, process::exit, sync::Arc, time::Duration};
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
pub(crate) type State = Arc<RwLock<(DataManager, Tera, Client)>>;
|
||||
|
|
@ -47,6 +51,27 @@ async fn main() {
|
|||
init_dirs(&config).await;
|
||||
let html_path = write_assets(&config).await;
|
||||
|
||||
// governor
|
||||
// (1000/125) * 32 = 256 requests/second
|
||||
let governor_config = Arc::new(
|
||||
GovernorConfigBuilder::default()
|
||||
.per_millisecond(125)
|
||||
.burst_size(32)
|
||||
.finish()
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
let governor_limiter = governor_config.limiter().clone();
|
||||
let governor_interval = Duration::from_secs(60);
|
||||
|
||||
std::thread::spawn(move || {
|
||||
loop {
|
||||
std::thread::sleep(governor_interval);
|
||||
tracing::info!("rate limiting storage size: {}", governor_limiter.len());
|
||||
governor_limiter.retain_recent();
|
||||
}
|
||||
});
|
||||
|
||||
// ...
|
||||
let database = DataManager::new(config.clone()).await.unwrap();
|
||||
database.init().await.unwrap();
|
||||
|
|
@ -78,7 +103,11 @@ async fn main() {
|
|||
TraceLayer::new_for_http()
|
||||
.make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
|
||||
.on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
|
||||
);
|
||||
)
|
||||
.layer(CatchPanicLayer::new())
|
||||
.layer(GovernorLayer {
|
||||
config: governor_config,
|
||||
});
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", config.port))
|
||||
.await
|
||||
|
|
@ -86,5 +115,10 @@ async fn main() {
|
|||
|
||||
info!("🐇 tetratto.");
|
||||
info!("listening on http://0.0.0.0:{}", config.port);
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
axum::serve(
|
||||
listener,
|
||||
app.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -443,15 +443,15 @@ pub async fn subscription_handler(
|
|||
|
||||
let data = data.clone();
|
||||
Ok(ws.on_upgrade(|socket| async move {
|
||||
tokio::spawn(async move {
|
||||
handle_socket(socket, data, user_id, id).await;
|
||||
});
|
||||
handle_socket(socket, data, user_id, id).await;
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(feature = "redis")]
|
||||
pub async fn handle_socket(socket: WebSocket, db: DataManager, user_id: String, stream_id: String) {
|
||||
let (mut sink, mut stream) = socket.split();
|
||||
let socket_id = tetratto_shared::hash::salt();
|
||||
db.2.incr("atto.active_connections:users".to_string()).await;
|
||||
|
||||
// get channel permissions
|
||||
let channel = format!("{user_id}/{stream_id}");
|
||||
|
|
@ -463,21 +463,25 @@ pub async fn handle_socket(socket: WebSocket, db: DataManager, user_id: String,
|
|||
continue;
|
||||
}
|
||||
|
||||
drop(stream);
|
||||
break;
|
||||
}
|
||||
});
|
||||
|
||||
let heartbeat_uri = format!("{channel}/{socket_id}");
|
||||
|
||||
let dbc = db.clone();
|
||||
let channel_c = channel.clone();
|
||||
let heartbeat_c = heartbeat_uri.clone();
|
||||
let mut redis_task = tokio::spawn(async move {
|
||||
// forward messages from redis to the socket
|
||||
let mut con = dbc.2.get_con().await;
|
||||
let mut pubsub = con.as_pubsub();
|
||||
pubsub.subscribe(channel_c).unwrap();
|
||||
let mut pubsub = dbc.2.client.get_async_pubsub().await.unwrap();
|
||||
|
||||
pubsub.subscribe(channel_c).await.unwrap();
|
||||
pubsub.subscribe(heartbeat_c).await.unwrap();
|
||||
|
||||
// listen for pubsub messages
|
||||
while let Ok(msg) = pubsub.get_message() {
|
||||
let mut pubsub = pubsub.into_on_message();
|
||||
while let Some(msg) = pubsub.next().await {
|
||||
// payload is a stringified SocketMessage
|
||||
let smsg = msg.get_payload::<String>().unwrap();
|
||||
let packet: SocketMessage = serde_json::from_str(&smsg).unwrap();
|
||||
|
|
@ -485,18 +489,15 @@ pub async fn handle_socket(socket: WebSocket, db: DataManager, user_id: String,
|
|||
if packet.method == SocketMethod::Forward(PacketType::Ping) {
|
||||
// forward with custom message
|
||||
if sink.send(WsMessage::Text("Ping".into())).await.is_err() {
|
||||
drop(sink);
|
||||
break;
|
||||
}
|
||||
} else if packet.method == SocketMethod::Message {
|
||||
if sink.send(WsMessage::Text(smsg.into())).await.is_err() {
|
||||
drop(sink);
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
// forward to client
|
||||
if sink.send(WsMessage::Text(smsg.into())).await.is_err() {
|
||||
drop(sink);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
|
@ -504,14 +505,14 @@ pub async fn handle_socket(socket: WebSocket, db: DataManager, user_id: String,
|
|||
});
|
||||
|
||||
let db2c = db.2.clone();
|
||||
let channel_c = channel.clone();
|
||||
let heartbeat_c = heartbeat_uri.clone();
|
||||
let heartbeat_task = tokio::spawn(async move {
|
||||
let mut con = db2c.get_con().await;
|
||||
let mut heartbeat = tokio::time::interval(Duration::from_secs(30));
|
||||
let mut heartbeat = tokio::time::interval(Duration::from_secs(10));
|
||||
|
||||
loop {
|
||||
con.publish::<&str, String, ()>(
|
||||
&channel_c,
|
||||
&heartbeat_c,
|
||||
serde_json::to_string(&SocketMessage {
|
||||
method: SocketMethod::Forward(PacketType::Ping),
|
||||
data: "Ping".to_string(),
|
||||
|
|
@ -524,7 +525,6 @@ pub async fn handle_socket(socket: WebSocket, db: DataManager, user_id: String,
|
|||
}
|
||||
});
|
||||
|
||||
db.2.incr("atto.active_connections:users".to_string()).await;
|
||||
tokio::select! {
|
||||
_ = (&mut recv_task) => redis_task.abort(),
|
||||
_ = (&mut redis_task) => recv_task.abort()
|
||||
|
|
|
|||
|
|
@ -187,14 +187,14 @@ pub async fn handle_socket(socket: WebSocket, db: DataManager, community_id: Str
|
|||
let dbc = db.clone();
|
||||
let mut redis_task = tokio::spawn(async move {
|
||||
// forward messages from redis to the socket
|
||||
let mut con = dbc.2.get_con().await;
|
||||
let mut pubsub = con.as_pubsub();
|
||||
let mut pubsub = dbc.2.client.get_async_pubsub().await.unwrap();
|
||||
|
||||
pubsub.subscribe(user.id).unwrap();
|
||||
pubsub.subscribe(community_id.clone()).unwrap();
|
||||
pubsub.subscribe(user.id).await.unwrap();
|
||||
pubsub.subscribe(community_id.clone()).await.unwrap();
|
||||
|
||||
// listen for pubsub messages
|
||||
while let Ok(msg) = pubsub.get_message() {
|
||||
let mut pubsub = pubsub.into_on_message();
|
||||
while let Some(msg) = pubsub.next().await {
|
||||
// payload is a stringified SocketMessage
|
||||
let smsg = msg.get_payload::<String>().unwrap();
|
||||
let packet: SocketMessage = serde_json::from_str(&smsg).unwrap();
|
||||
|
|
@ -241,7 +241,7 @@ pub async fn handle_socket(socket: WebSocket, db: DataManager, community_id: Str
|
|||
let db2c = db.2.clone();
|
||||
let heartbeat_task = tokio::spawn(async move {
|
||||
let mut con = db2c.get_con().await;
|
||||
let mut heartbeat = tokio::time::interval(Duration::from_secs(30));
|
||||
let mut heartbeat = tokio::time::interval(Duration::from_secs(10));
|
||||
|
||||
loop {
|
||||
con.publish::<usize, String, ()>(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue