nikita
/
lavina
forked from lavina/lavina
1
0
Fork 0
Code Pull Requests Activity

use scoped for telemetry and all socket connections

This commit is contained in:
Nikita Vilunov 2023-09-06 17:55:24 +02:00
parent 1652333184
commit c114de7dfa
6 changed files with 91 additions and 111 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/core/player.rs
View File

@ -227,8 +227,7 @@ pub enum Updates {
}
/// Handle to a player registry — a shared data structure containing information about players.
#[derive(Clone)]
pub struct PlayerRegistry(Arc<RwLock<PlayerRegistryInner>>);
pub struct PlayerRegistry(RwLock<PlayerRegistryInner>);
impl PlayerRegistry {
pub fn empty(
room_registry: RoomRegistry,
@ -242,10 +241,10 @@ impl PlayerRegistry {
players: HashMap::new(),
metric_active_players,
};
Ok(PlayerRegistry(Arc::new(RwLock::new(inner))))
Ok(PlayerRegistry(RwLock::new(inner)))
}
pub async fn get_or_create_player(&mut self, id: PlayerId) -> PlayerHandle {
pub async fn get_or_create_player(&self, id: PlayerId) -> PlayerHandle {
let mut inner = self.0.write().unwrap();
if let Some((handle, _)) = inner.players.get(&id) {
handle.clone()
@ -257,7 +256,7 @@ impl PlayerRegistry {
}
}
pub async fn connect_to_player(&mut self, id: PlayerId) -> PlayerConnection {
pub async fn connect_to_player(&self, id: PlayerId) -> PlayerConnection {
let player_handle = self.get_or_create_player(id).await;
player_handle.subscribe().await
}

2
src/core/repo/mod.rs
View File

@ -33,7 +33,7 @@ impl Storage {
Ok(Storage { conn })
}
pub async fn retrieve_user_by_name(&mut self, name: &str) -> Result<Option<StoredUser>> {
pub async fn retrieve_user_by_name(&self, name: &str) -> Result<Option<StoredUser>> {
let mut executor = self.conn.lock().await;
let res = sqlx::query_as(
"select u.id, u.name, c.password

5
src/main.rs
View File

@ -55,11 +55,10 @@ async fn main() -> Result<()> {
let storage = Storage::open(storage_config).await?;
let rooms = RoomRegistry::new(&mut metrics, storage.clone())?;
let mut players = PlayerRegistry::empty(rooms.clone(), &mut metrics)?;
let telemetry_terminator =
util::telemetry::launch(telemetry_config, metrics.clone(), rooms.clone()).await?;
// unsafe: outer future is never dropped, scope is joined on `scope.collect`
let mut scope = unsafe { Scope::create() };
let telemetry_terminator = util::telemetry::launch(telemetry_config, &metrics, &rooms, &mut scope).await?;
let irc = projections::irc::launch(&irc_config, &players, &rooms, &metrics, &storage, &mut scope).await?;
let xmpp = projections::xmpp::launch(xmpp_config, &players, &rooms, &metrics, &mut scope).await?;
tracing::info!("Started");
@ -69,10 +68,10 @@ async fn main() -> Result<()> {
tracing::info!("Begin shutdown");
let _ = xmpp.send(());
let _ = irc.send(());
let _ = telemetry_terminator.send(());
let _ = scope.collect().await;
drop(scope);
telemetry_terminator.terminate().await?;
players.shutdown_all().await?;
drop(players);
drop(rooms);

47
src/projections/irc/mod.rs
View File

@ -2,7 +2,6 @@ use std::collections::HashMap;
use std::net::SocketAddr;
use futures_util::FutureExt;
use futures_util::future::join_all;
use prometheus::{IntCounter, IntGauge, Registry as MetricsRegistry};
use serde::Deserialize;
use tokio::io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
@ -17,7 +16,6 @@ use crate::prelude::*;
use crate::protos::irc::client::{client_message, ClientMessage};
use crate::protos::irc::server::{AwayStatus, ServerMessage, ServerMessageBody};
use crate::protos::irc::{Chan, Recipient};
use crate::util::Terminator;
#[cfg(test)]
mod test;
@ -44,10 +42,10 @@ async fn handle_socket(
config: ServerConfig,
mut stream: TcpStream,
socket_addr: &SocketAddr,
players: PlayerRegistry,
rooms: RoomRegistry,
players: &PlayerRegistry,
rooms: &RoomRegistry,
termination: Deferred<()>, // TODO use it to stop the connection gracefully
mut storage: Storage,
storage: &Storage,
) -> Result<()> {
let (reader, writer) = stream.split();
let mut reader: BufReader<ReadHalf> = BufReader::new(reader);
@ -67,7 +65,7 @@ async fn handle_socket(
writer.flush().await?;
let registered_user: Result<RegisteredUser> =
handle_registration(&mut reader, &mut writer, &mut storage).await;
handle_registration(&mut reader, &mut writer, &storage).await;
match registered_user {
Ok(user) => {
@ -84,7 +82,7 @@ async fn handle_socket(
async fn handle_registration<'a>(
reader: &mut BufReader<ReadHalf<'a>>,
writer: &mut BufWriter<WriteHalf<'a>>,
storage: &mut Storage,
storage: &Storage,
) -> Result<RegisteredUser> {
let mut buffer = vec![];
@ -177,8 +175,8 @@ async fn handle_registration<'a>(
async fn handle_registered_socket<'a>(
config: ServerConfig,
mut players: PlayerRegistry,
rooms: RoomRegistry,
players: &PlayerRegistry,
rooms: &RoomRegistry,
reader: &mut BufReader<ReadHalf<'a>>,
writer: &mut BufWriter<WriteHalf<'a>>,
user: RegisteredUser,
@ -716,6 +714,7 @@ pub async fn launch<'a>(
// TODO probably should separate logic for accepting new connection and storing them
// into two tasks so that they don't block each other
let mut actors = HashMap::new();
let mut scope = unsafe { Scope::create() };
loop {
select! {
biased;
@ -738,23 +737,22 @@ pub async fn launch<'a>(
continue;
}
let terminator = Terminator::spawn(|termination| {
let players = players.clone();
let rooms = rooms.clone();
let current_connections_clone = current_connections.clone();
let stopped_tx = stopped_tx.clone();
let storage = storage.clone();
async move {
match handle_socket(config, stream, &socket_addr, players, rooms, termination, storage).await {
let (a, rx) = oneshot();
let future = async move {
match handle_socket(config, stream, &socket_addr, players, rooms, rx, storage).await {
Ok(_) => log::info!("Connection terminated"),
Err(err) => log::warn!("Connection failed: {err}"),
}
current_connections_clone.dec();
stopped_tx.send(socket_addr).await?;
Ok(())
}
});
actors.insert(socket_addr, terminator);
};
scope.spawn(future.map(|_: Result<()>| ()));
actors.insert(socket_addr, a);
},
Err(err) => log::warn!("Failed to accept new connection: {err}"),
}
@ -763,15 +761,14 @@ pub async fn launch<'a>(
}
log::info!("Stopping IRC projection");
join_all(actors.into_iter().map(|(socket_addr, terminator)| async move {
log::debug!("Stopping IRC connection at {socket_addr}");
match terminator.terminate().await {
Ok(_) => log::debug!("Stopped IRC connection at {socket_addr}"),
Err(err) => {
log::warn!("IRC connection to {socket_addr} finished with error: {err}")
for (socket_addr, terminator) in actors {
match terminator.send(()) {
Ok(_) => log::debug!("Stopping IRC connection at {socket_addr}"),
Err(_) => log::debug!("IRC connection at {socket_addr} already stopped")
}
}
})).await;
let _ = scope.collect().await;
drop(scope);
log::info!("Stopped IRC projection");
Ok(())
};

41
src/projections/xmpp/mod.rs
View File

@ -8,7 +8,6 @@ use std::path::PathBuf;
use std::sync::Arc;
use futures_util::FutureExt;
use futures_util::future::join_all;
use prometheus::Registry as MetricsRegistry;
use quick_xml::events::{BytesDecl, Event};
use quick_xml::{NsReader, Writer};
@ -31,7 +30,6 @@ use crate::protos::xmpp::roster::RosterQuery;
use crate::protos::xmpp::session::Session;
use crate::protos::xmpp::stream::*;
use crate::util::xml::{Continuation, FromXml, Parser, ToXml};
use crate::util::Terminator;
use self::proto::{ClientPacket, IqClientBody};
@ -82,6 +80,7 @@ pub async fn launch<'a>(
let future = async move {
let (stopped_tx, mut stopped_rx) = channel(32);
let mut actors = HashMap::new();
let mut scope = unsafe { Scope::create() };
loop {
select! {
biased;
@ -99,22 +98,20 @@ pub async fn launch<'a>(
// TODO kill the older connection and restart it
continue;
}
let players = players.clone();
let rooms = rooms.clone();
let terminator = Terminator::spawn(|termination| {
let stopped_tx = stopped_tx.clone();
let loaded_config = loaded_config.clone();
async move {
match handle_socket(loaded_config, stream, &socket_addr, players, rooms, termination).await {
let (a, rx) = oneshot();
let future = async move {
match handle_socket(loaded_config, stream, &socket_addr, players, rooms, rx).await {
Ok(_) => log::info!("Connection terminated"),
Err(err) => log::warn!("Connection failed: {err}"),
}
stopped_tx.send(socket_addr).await?;
Ok(())
}
});
};
scope.spawn(future.map(|_: Result<()>| ()));
actors.insert(socket_addr, terminator);
actors.insert(socket_addr, a);
},
Err(err) => log::warn!("Failed to accept new connection: {err}"),
}
@ -122,22 +119,14 @@ pub async fn launch<'a>(
}
}
log::info!("Stopping XMPP projection");
join_all(
actors
.into_iter()
.map(|(socket_addr, terminator)| async move {
log::debug!("Stopping XMPP connection at {socket_addr}");
match terminator.terminate().await {
Ok(_) => log::debug!("Stopped XMPP connection at {socket_addr}"),
Err(err) => {
log::warn!(
"XMPP connection to {socket_addr} finished with error: {err}"
)
for (socket_addr, terminator) in actors {
match terminator.send(()) {
Ok(_) => log::debug!("Stopping XMPP connection at {socket_addr}"),
Err(_) => log::debug!("XMPP connection at {socket_addr} already stopped")
}
}
}),
)
.await;
let _ = scope.collect().await;
drop(scope);
log::info!("Stopped XMPP projection");
Ok(())
};
@ -150,8 +139,8 @@ async fn handle_socket(
config: Arc<LoadedConfig>,
mut stream: TcpStream,
socket_addr: &SocketAddr,
mut players: PlayerRegistry,
rooms: RoomRegistry,
players: &PlayerRegistry,
rooms: &RoomRegistry,
termination: Deferred<()>, // TODO use it to stop the connection gracefully
) -> Result<()> {
log::debug!("Received an XMPP connection from {socket_addr}");

48
src/util/telemetry.rs
View File

@ -15,7 +15,6 @@ use crate::core::room::RoomRegistry;
use crate::prelude::*;
use crate::util::http::*;
use crate::util::Terminator;
type BoxBody = http_body_util::combinators::BoxBody<Bytes, Infallible>;
type HttpResult<T> = std::result::Result<T, Infallible>;
@ -25,37 +24,28 @@ pub struct ServerConfig {
pub listen_on: SocketAddr,
}
pub async fn launch(
pub async fn launch<'a>(
config: ServerConfig,
metrics: MetricsRegistry,
rooms: RoomRegistry,
) -> Result<Terminator> {
metrics: &'a MetricsRegistry,
rooms: &'a RoomRegistry,
scope: &mut Scope<'a>,
) -> Result<Promise<()>> {
log::info!("Starting the telemetry service");
let listener = TcpListener::bind(config.listen_on).await?;
log::debug!("Listener started");
let terminator = Terminator::spawn(|rx| main_loop(listener, metrics, rooms, rx.map(|_| ())));
Ok(terminator)
}
async fn main_loop(
listener: TcpListener,
metrics: MetricsRegistry,
rooms: RoomRegistry,
termination: impl Future<Output = ()>,
) -> Result<()> {
pin!(termination);
let (signal, mut rx) = oneshot();
let future = async move {
let mut scope = unsafe { Scope::create() };
loop {
select! {
biased;
_ = &mut termination => break,
_ = &mut rx => break,
result = listener.accept() => {
let (stream, _) = result?;
let metrics = metrics.clone();
let rooms = rooms.clone();
tokio::task::spawn(async move {
let registry = metrics.clone();
let rooms = rooms.clone();
let server = http1::Builder::new().serve_connection(stream, service_fn(move |r| route(registry.clone(), rooms.clone(), r)));
scope.spawn(async move {
let server = http1::Builder::new().serve_connection(stream, service_fn(move |r| route(metrics, rooms, r)));
if let Err(err) = server.await {
tracing::error!("Error serving connection: {:?}", err);
}
@ -63,13 +53,19 @@ async fn main_loop(
},
}
}
let _ = scope.collect().await;
drop(scope);
log::info!("Terminating the telemetry service");
Ok(())
};
scope.spawn(future.map(|_: Result<()>| ()));
Ok(signal)
}
async fn route(
registry: MetricsRegistry,
rooms: RoomRegistry,
registry: &MetricsRegistry,
rooms: &RoomRegistry,
request: Request<hyper::body::Incoming>,
) -> std::result::Result<Response<BoxBody>, Infallible> {
match (request.method(), request.uri().path()) {
@ -79,7 +75,7 @@ async fn route(
}
}
fn endpoint_metrics(registry: MetricsRegistry) -> HttpResult<Response<Full<Bytes>>> {
fn endpoint_metrics(registry: &MetricsRegistry) -> HttpResult<Response<Full<Bytes>>> {
let mf = registry.gather();
let mut buffer = vec![];
TextEncoder
@ -88,7 +84,7 @@ fn endpoint_metrics(registry: MetricsRegistry) -> HttpResult<Response<Full<Bytes
Ok(Response::new(Full::new(Bytes::from(buffer))))
}
async fn endpoint_rooms(rooms: RoomRegistry) -> HttpResult<Response<Full<Bytes>>> {
async fn endpoint_rooms(rooms: &RoomRegistry) -> HttpResult<Response<Full<Bytes>>> {
let room_list = rooms.get_all_rooms().await;
let mut buffer = vec![];
serde_json::to_writer(&mut buffer, &room_list).expect("unexpected fail when writing to vec");