diff --git a/rc_chat_room/src/main.rs b/rc_chat_room/src/main.rs index 321d31d..fc6c280 100644 --- a/rc_chat_room/src/main.rs +++ b/rc_chat_room/src/main.rs @@ -34,7 +34,7 @@ async fn main() -> std::io::Result<()> { let chat_system = state::chat::ChatImpl::new(>::chat_system_config(&cubes)).expect("Bad chat config data"); - let server = std::sync::Arc::new(polariton_server::Server::new(operations::handler(chat_system, &cubes), polariton_server::events::EventsHandler::new())); + let server = std::sync::Arc::new(polariton_server::Server::new(operations::handler(chat_system.clone(), &cubes), polariton_server::events::EventsHandler::new())); let ip_addr: std::net::IpAddr = args.ip.parse().expect("Invalid IP address"); @@ -46,11 +46,11 @@ async fn main() -> std::io::Result<()> { if args.once { log::warn!("Handling first connection and then exiting"); let (socket, address) = listener.accept().await?; - process_socket(socket, address, server.clone(), users.clone()).await; + process_socket(socket, address, server.clone(), users.clone(), chat_system).await; } else { loop { let (socket, address) = listener.accept().await?; - tokio::spawn(process_socket(socket, address, server.clone(), users.clone())); + tokio::spawn(process_socket(socket, address, server.clone(), users.clone(), chat_system.clone())); } } server.join(); @@ -58,7 +58,7 @@ async fn main() -> std::io::Result<()> { Ok(()) } -async fn process_socket(mut socket: net::TcpStream, address: std::net::SocketAddr, server: std::sync::Arc>, users: std::sync::Arc) { +async fn process_socket(mut socket: net::TcpStream, address: std::net::SocketAddr, server: std::sync::Arc>, users: std::sync::Arc, chat: state::chat::ChatImpl) { log::debug!("Accepting connection from address {}", address); let enc = match do_connect_handshake(&mut socket).await { Some(x) => x, @@ -67,13 +67,13 @@ async fn process_socket(mut socket: net::TcpStream, address: std::net::SocketAdd return; } }; - ONLINE_USERS.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + ONLINE_USERS.store(chat.system().await.user_count() as u64 + 1, std::sync::atomic::Ordering::SeqCst); let (chann_tx, chann_rx) = tokio::sync::mpsc::unbounded_channel(); let user_state = std::sync::Arc::new(oj_rc_core::UserState::<()>::new(users, chann_tx.clone())); let (socket_r, socket_w) = socket.into_split(); server.handle_async_with_channel_join(socket_r, socket_w, user_state.clone(), polariton::packet::SerdesContext::from_boxed(Default::default(), enc), chann_tx, chann_rx).await; log::debug!("Goodbye connection from address {}", address); - ONLINE_USERS.fetch_sub(1, std::sync::atomic::Ordering::SeqCst); + ONLINE_USERS.store(chat.system().await.user_count() as u64 - 1, std::sync::atomic::Ordering::SeqCst); if let Ok(user_info) = user_state.user() { update_status(user_info.as_ref().as_ref()).await; } diff --git a/rc_services_room/src/main.rs b/rc_services_room/src/main.rs index b5acac1..8ac0ead 100644 --- a/rc_services_room/src/main.rs +++ b/rc_services_room/src/main.rs @@ -15,6 +15,8 @@ pub type UserTy = std::sync::Arc>; pub static START_TIMESTAMP_S: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0); pub static ONLINE_USERS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); +pub static LOGINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); +pub static USER_HANDLES: std::sync::Mutex>> = std::sync::Mutex::new(Vec::new()); pub struct InitConfig { pub cubes: oj_rc_core::persist::config::ConfigImpl, @@ -66,7 +68,8 @@ async fn main() -> std::io::Result<()> { } async fn process_socket(mut socket: net::TcpStream, address: std::net::SocketAddr, server: std::sync::Arc>, init_ctx: std::sync::Arc) { - log::debug!("Accepting connection from address {}", address); + let login_num = std::sync::Arc::new(LOGINS.fetch_add(1, std::sync::atomic::Ordering::SeqCst)); + log::debug!("Accepting connection from address {} (login #{})", address, login_num); let enc = match do_connect_handshake(&mut socket).await { Some(x) => x, None => { @@ -74,14 +77,32 @@ async fn process_socket(mut socket: net::TcpStream, address: std::net::SocketAdd return; } }; - ONLINE_USERS.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + { + if let Ok(mut handles) = USER_HANDLES.lock() { + handles.push(std::sync::Arc::downgrade(&login_num)); + ONLINE_USERS.store(handles.len() as u64, std::sync::atomic::Ordering::SeqCst); + } else { + // this should never happen + log::warn!("USER_HANDLES lock is poisoned, cannot track online users anymore (please restart)"); + } + } let (socket_r, socket_w) = socket.into_split(); let (chann_tx, chann_rx) = tokio::sync::mpsc::unbounded_channel(); let user_state = std::sync::Arc::new(oj_rc_core::UserState::<()>::new(init_ctx.users.clone(), chann_tx.clone())); let ctx = polariton::packet::SerdesContext::from_boxed(Default::default(), enc); server.handle_async_with_channel_join(socket_r, socket_w, user_state.clone(), ctx, chann_tx, chann_rx).await; - log::debug!("Goodbye connection from address {}", address); - ONLINE_USERS.fetch_sub(1, std::sync::atomic::Ordering::SeqCst); + log::debug!("Goodbye connection from address {} (login #{})", address, login_num); + drop(login_num); // explicit for good measure + { + if let Ok(mut handles) = USER_HANDLES.lock() { + handles.retain(|x| x.strong_count() != 0); + ONLINE_USERS.store(handles.len() as u64, std::sync::atomic::Ordering::SeqCst); + } else { + // this should never happen + ONLINE_USERS.fetch_sub(1, std::sync::atomic::Ordering::SeqCst); + log::warn!("USER_HANDLES lock is poisoned, cannot track online users anymore (please restart)"); + } + } if let Ok(user_info) = user_state.user() { update_status(user_info.as_ref().as_ref()).await; }