expose count of connected end servers to testing code

main
(on company time) 2022-12-19 14:58:12 -06:00
parent 0992d76fdc
commit e3ff600b51
3 changed files with 43 additions and 2 deletions

View File

@ -28,10 +28,16 @@ pub struct Opt {
pub struct App { pub struct App {
endpoint: quinn::Endpoint, endpoint: quinn::Endpoint,
listen_addr: SocketAddr, listen_addr: SocketAddr,
pub (crate) metrics: Arc <RwLock <Metrics>>,
server_cert: Vec <u8>, server_cert: Vec <u8>,
tcp_listen_port: Option <u16>, tcp_listen_port: Option <u16>,
} }
#[derive (Default)]
pub (crate) struct Metrics {
pub (crate) connected_end_servers: usize,
}
impl App { impl App {
pub fn new (opt: Opt) -> anyhow::Result <Self> { pub fn new (opt: Opt) -> anyhow::Result <Self> {
let listen_addr = opt.listen_addr.unwrap_or_else (|| String::from ("0.0.0.0:30380")).parse ()?; let listen_addr = opt.listen_addr.unwrap_or_else (|| String::from ("0.0.0.0:30380")).parse ()?;
@ -42,6 +48,7 @@ impl App {
Ok (Self { Ok (Self {
endpoint, endpoint,
listen_addr, listen_addr,
metrics: Default::default (),
server_cert, server_cert,
tcp_listen_port: opt.tcp_listen_port, tcp_listen_port: opt.tcp_listen_port,
}) })
@ -59,11 +66,13 @@ impl App {
let Self { let Self {
endpoint, endpoint,
listen_addr, listen_addr,
metrics,
server_cert, server_cert,
tcp_listen_port, tcp_listen_port,
} = self; } = self;
let relay_state = RelayState::default (); let mut relay_state = RelayState::default ();
relay_state.metrics = metrics;
if let Err (e) = relay_state.reload_config ().await { if let Err (e) = relay_state.reload_config ().await {
error! ("{:?}", e); error! ("{:?}", e);
} }
@ -110,6 +119,7 @@ impl App {
// Each new peer QUIC connection gets its own task // Each new peer QUIC connection gets its own task
tokio::spawn (async move { tokio::spawn (async move {
let active = relay_state.stats.quic.connect (); let active = relay_state.stats.quic.connect ();
debug! ("QUIC connections: {}", active); debug! ("QUIC connections: {}", active);
match handle_quic_connection (Arc::clone (&relay_state), conn).await { match handle_quic_connection (Arc::clone (&relay_state), conn).await {
@ -228,6 +238,7 @@ struct RelayState {
config: arc_swap::ArcSwap <Config>, config: arc_swap::ArcSwap <Config>,
p4_server_proxies: Mutex <HashMap <PeerId, P4State>>, p4_server_proxies: Mutex <HashMap <PeerId, P4State>>,
direc_cookies: Mutex <HashMap <Vec <u8>, DirecState>>, direc_cookies: Mutex <HashMap <Vec <u8>, DirecState>>,
metrics: Arc <RwLock <Metrics>>,
stats: Stats, stats: Stats,
http_client: reqwest::Client, http_client: reqwest::Client,
} }
@ -431,7 +442,17 @@ async fn handle_quic_connection (
// TODO: Check authorization for P4 peers // TODO: Check authorization for P4 peers
protocol::p3_authorize_p4_peer (&mut send).await?; protocol::p3_authorize_p4_peer (&mut send).await?;
let metrics = Arc::clone (&relay_state.metrics);
{
let mut m = metrics.write ().await;
m.connected_end_servers += 1;
}
handle_p4_connection (relay_state, conn, peer).await?; handle_p4_connection (relay_state, conn, peer).await?;
{
let mut m = metrics.write ().await;
m.connected_end_servers -= 1;
}
}, },
} }

View File

@ -37,6 +37,7 @@ pub use tokio::{
}, },
sync::{ sync::{
Mutex, Mutex,
RwLock,
mpsc, mpsc,
}, },
task::JoinHandle, task::JoinHandle,

View File

@ -1,3 +1,5 @@
use crate::prelude::*;
#[test] #[test]
fn end_to_end () -> anyhow::Result <()> { fn end_to_end () -> anyhow::Result <()> {
let rt = tokio::runtime::Runtime::new ()?; let rt = tokio::runtime::Runtime::new ()?;
@ -14,12 +16,20 @@ async fn end_to_end_async () -> anyhow::Result <()> {
tcp_listen_port: None, tcp_listen_port: None,
}; };
let relay_app = relay::App::new (relay_opt)?; let relay_app = relay::App::new (relay_opt)?;
let relay_quic_port = relay_app.listen_addr ().port (); let relay_quic_port = relay_app.listen_addr ().port ();
let relay_cert = Vec::from (relay_app.server_cert ()); let relay_cert = Vec::from (relay_app.server_cert ());
let relay_metrics = Arc::clone (&relay_app.metrics);
let task_relay = tokio::spawn (async move { let task_relay = tokio::spawn (async move {
relay_app.run ().await relay_app.run ().await
}); });
{
let m = relay_metrics.read ().await;
assert_eq! (m.connected_end_servers, 0);
}
// Connect with wrong port, should fail // Connect with wrong port, should fail
let server_conf = server::Config { let server_conf = server::Config {
@ -48,6 +58,11 @@ async fn end_to_end_async () -> anyhow::Result <()> {
assert! (server_err.is_err ()); assert! (server_err.is_err ());
{
let m = relay_metrics.read ().await;
assert_eq! (m.connected_end_servers, 0);
}
// Connect properly // Connect properly
let server_conf = server::Config { let server_conf = server::Config {
@ -58,8 +73,12 @@ async fn end_to_end_async () -> anyhow::Result <()> {
use_udp_over_tcp: false, use_udp_over_tcp: false,
}; };
let (server, _) = server::P4EndServer::connect (server_conf).await?; let (server, server_shutdown_tx) = server::P4EndServer::connect (server_conf).await?;
{
let m = relay_metrics.read ().await;
assert_eq! (m.connected_end_servers, 1);
}
Ok (()) Ok (())
} }