vecna
/
lox
mirror of https://gitlab.torproject.org/vecna/lox.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Adds graceful shutdown for all tasks, rdsys stream inconsistent shutdown

This commit is contained in:
onyinyang 2023-03-17 10:41:32 -04:00
parent 9616b95b5d
commit 9e78a90541
No known key found for this signature in database
GPG Key ID: 156A6435430C2036
1 changed files with 182 additions and 128 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

126
crates/lox-distributor/src/main.rs
View File

@ -18,7 +18,7 @@ use lox::{
}; };
use lox::{BridgeAuth, BridgeDb, OPENINV_LENGTH}; use lox::{BridgeAuth, BridgeDb, OPENINV_LENGTH};
use rand::RngCore; use rand::RngCore;
use rdsys_backend::{proto::ResourceDiff, start_stream, ResourceStream}; use rdsys_backend::{proto::ResourceDiff, start_stream};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::{ use std::{
convert::Infallible, convert::Infallible,
@ -29,13 +29,12 @@ use std::{
sync::{Arc, Mutex}, sync::{Arc, Mutex},
time::Duration, time::Duration,
}; };
use serde_json; use serde_json;
use serde_with::serde_as; use serde_with::serde_as;
use tokio::{ use tokio::{
spawn, spawn,
sync::{mpsc, oneshot}, sync::{broadcast, mpsc, oneshot},
time::sleep, task::yield_now, time::sleep, signal,
}; };
#[serde_as] #[serde_as]
@ -373,9 +372,21 @@ async fn shutdown_signal() {
tokio::signal::ctrl_c() tokio::signal::ctrl_c()
.await .await
.expect("failed to listen for ctrl+c signal"); .expect("failed to listen for ctrl+c signal");
println!("Shut down Lox Server");
} }
async fn rdsys_sender(rstream: ResourceStream, tx: Sender<ResourceDiff>) { async fn rdsys_stream(rtype: ResourceInfo, tx: Sender<ResourceDiff>, mut kill: broadcast::Receiver<()>) {
tokio:: select! {
start_rdsys_stream = rdsys_sender(rtype, tx) => start_rdsys_stream ,
_ = kill.recv() => {println!("Shut down rdsys stream"); return},
}
}
async fn rdsys_sender(rtype: ResourceInfo, tx: Sender<ResourceDiff>) {
let rstream = start_stream(rtype.endpoint, rtype.name, rtype.token, rtype.types)
.await
.expect("rdsys stream initialization failed. Start rdsys or check config.json");
sleep(Duration::from_millis(1)).await;
for diff in rstream { for diff in rstream {
println!("Received diff: {:?}", diff); //send this through a channel println!("Received diff: {:?}", diff); //send this through a channel
tx.send(diff).await.unwrap(); tx.send(diff).await.unwrap();
@ -383,6 +394,13 @@ async fn rdsys_sender(rstream: ResourceStream, tx: Sender<ResourceDiff>) {
} }
} }
async fn rdsys_bridge_parser(rdsys_tx: mpsc::Sender<Command>, rx: Receiver<ResourceDiff>, mut kill: broadcast::Receiver<()>) {
tokio:: select! {
start_bridge_parser = parse_bridges(rdsys_tx, rx) => start_bridge_parser ,
_ = kill.recv() => {println!("Shut down bridge_parser"); return},
}
}
async fn parse_bridges(rdsys_tx: mpsc::Sender<Command>, rx: Receiver<ResourceDiff>) { async fn parse_bridges(rdsys_tx: mpsc::Sender<Command>, rx: Receiver<ResourceDiff>) {
loop { loop {
let resourcediff = rx.recv().await.unwrap(); let resourcediff = rx.recv().await.unwrap();
@ -394,31 +412,16 @@ async fn parse_bridges(rdsys_tx: mpsc::Sender<Command>, rx: Receiver<ResourceDif
} }
} }
#[derive(Debug)] async fn create_context_manager(context_rx: mpsc::Receiver<Command>, mut kill: broadcast::Receiver<()>) {
enum Command { tokio:: select! {
Rdsys { create_context = context_manager(context_rx) => create_context,
resourcediff: ResourceDiff, _ = kill.recv() => {println!("Shut down context_manager"); return},
}, }
Request {
req: Request<Body>,
sender: oneshot::Sender<Result<Response<Body>, Infallible>>,
},
} }
// Run with cargo run -- config.json async fn context_manager(mut context_rx: mpsc::Receiver<Command>) {
#[tokio::main]
async fn main() {
let args: Vec<String> = env::args().collect();
let file = File::open(&args[1]).expect("Should have been able to read config.json file");
let reader = BufReader::new(file);
// Read the JSON contents of the file as a ResourceInfo
let rtype: ResourceInfo = serde_json::from_reader(reader).unwrap();
let (rdsys_tx, mut context_rx) = mpsc::channel(32);
let request_tx = rdsys_tx.clone();
let context_manager = spawn(async move {
// pass in distribution of open invite vs. hot spare buckets? // pass in distribution of open invite vs. hot spare buckets?
let bridgedb = BridgeDb::new(); let bridgedb = BridgeDb::new();
let lox_auth = BridgeAuth::new(bridgedb.pubkey); let lox_auth = BridgeAuth::new(bridgedb.pubkey);
@ -530,17 +533,61 @@ async fn main() {
} }
} }
} }
}
#[derive(Debug)]
enum Command {
Rdsys {
resourcediff: ResourceDiff,
},
Request {
req: Request<Body>,
sender: oneshot::Sender<Result<Response<Body>, Infallible>>,
},
}
#[tokio::main]
async fn main() {
let args: Vec<String> = env::args().collect();
let file = File::open(&args[1]).expect("Should have been able to read config.json file");
let reader = BufReader::new(file);
// Read the JSON contents of the file as a ResourceInfo
let rtype: ResourceInfo = serde_json::from_reader(reader).unwrap();
let (rdsys_tx, context_rx) = mpsc::channel(32);
let request_tx = rdsys_tx.clone();
// create the shutdown broadcast channel
let (shutdown_tx, mut shutdown_rx) = broadcast::channel(16);
let kill_stream = shutdown_tx.subscribe();
let kill_parser = shutdown_tx.subscribe();
let kill_context= shutdown_tx.subscribe();
// Listen for ctrl_c, send signal to broadcast shutdown to all threads by dropping shutdown_tx
let shutdown_handler = spawn(async move {
tokio::select! {
_ = signal::ctrl_c() => {
drop(shutdown_tx);
println!("Kill Sent, all threads should shutdown.");
_ = shutdown_rx.recv().await;
println!("Receiving shutdown signals. . .");
}
}
}); });
//Sender is resource stream and receiver is bridgedb function (add_openinv_bridges)
let (tx, rx) = async_channel::unbounded();
// to populate the bridge db
let rstream = start_stream(rtype.endpoint, rtype.name, rtype.token, rtype.types)
.await
.unwrap();
let rdsys_stream_handler = spawn(async { rdsys_sender(rstream, tx).await });
let rdsys_resource_receiver = spawn(async { parse_bridges(rdsys_tx, rx).await });
let context_manager = spawn(async move { create_context_manager(context_rx, kill_context).await });
let (tx, rx) = async_channel::unbounded();
let rdsys_stream_handler = spawn(async { rdsys_stream(rtype, tx, kill_stream).await });
let rdsys_resource_receiver = spawn(async { rdsys_bridge_parser(rdsys_tx, rx, kill_parser).await });
let make_service = make_service_fn(move |_conn: &AddrStream| { let make_service = make_service_fn(move |_conn: &AddrStream| {
let request_tx = request_tx.clone(); let request_tx = request_tx.clone();
@ -566,6 +613,13 @@ async fn main() {
if let Err(e) = graceful.await { if let Err(e) = graceful.await {
eprintln!("server error: {}", e); eprintln!("server error: {}", e);
} }
future::join_all([rdsys_stream_handler, rdsys_resource_receiver, context_manager]).await; future::join_all([
rdsys_stream_handler,
rdsys_resource_receiver,
context_manager,
shutdown_handler,
])
.await;
} }