diff --git a/fact/src/output/grpc.rs b/fact/src/output/grpc.rs index 969a15fd..af547ec1 100644 --- a/fact/src/output/grpc.rs +++ b/fact/src/output/grpc.rs @@ -9,7 +9,7 @@ use native_tls::{Certificate, Identity}; use openssl::{ec::EcKey, pkey::PKey}; use tokio::{ fs, - sync::{broadcast, watch}, + sync::{mpsc, oneshot, watch}, task::JoinHandle, time::sleep, }; @@ -21,8 +21,8 @@ use tonic::transport::Channel; use crate::{ config::{BackoffConfig, GrpcConfig}, - event::Event, metrics::EventCounter, + output::EventReceiver, }; struct Backoff { @@ -104,7 +104,7 @@ impl From<&BackoffConfig> for Backoff { } pub struct Client { - rx: broadcast::Receiver>, + subscriber: mpsc::Sender>, running: watch::Receiver, config: watch::Receiver, metrics: EventCounter, @@ -112,13 +112,13 @@ pub struct Client { impl Client { pub fn new( - rx: broadcast::Receiver>, + subscriber: mpsc::Sender>, running: watch::Receiver, metrics: EventCounter, config: watch::Receiver, ) -> Self { Client { - rx, + subscriber, running, config, metrics, @@ -230,19 +230,21 @@ impl Client { let mut client = FileActivityServiceClient::new(channel); let metrics = self.metrics.clone(); - let rx = - BroadcastStream::new(self.rx.resubscribe()).filter_map(move |event| match event { - Ok(event) => { - metrics.added(); - let event = Arc::unwrap_or_clone(event); - Some(event.into()) - } - Err(BroadcastStreamRecvError::Lagged(n)) => { - warn!("gRPC stream lagged, dropped {n} events"); - metrics.dropped_n(n); - None - } - }); + let (tx, rx) = oneshot::channel(); + self.subscriber.send(tx).await?; + let rx = rx.await?; + let rx = BroadcastStream::new(rx).filter_map(move |event| match event { + Ok(event) => { + metrics.added(); + let event = Arc::unwrap_or_clone(event); + Some(event.into()) + } + Err(BroadcastStreamRecvError::Lagged(n)) => { + warn!("gRPC stream lagged, dropped {n} events"); + metrics.dropped_n(n); + None + } + }); tokio::select! { res = client.communicate(rx) => { diff --git a/fact/src/output/mod.rs b/fact/src/output/mod.rs index 756c8eec..0090631a 100644 --- a/fact/src/output/mod.rs +++ b/fact/src/output/mod.rs @@ -12,6 +12,8 @@ use crate::{config::GrpcConfig, event::Event, metrics::OutputMetrics}; mod grpc; mod stdout; +type EventReceiver = broadcast::Receiver>; + /// Starts all the output tasks. /// /// Each task is responsible for managing its lifetime, handling @@ -24,10 +26,11 @@ pub fn start( stdout_enabled: bool, ) -> JoinHandle> { let (broad_tx, broad_rx) = broadcast::channel(100); + let (subs_req, mut subs_rx) = mpsc::channel(10); let mut run = running.clone(); let grpc_client = grpc::Client::new( - broad_rx.resubscribe(), + subs_req.clone(), running.clone(), metrics.grpc.clone(), config.clone(), @@ -36,40 +39,42 @@ pub fn start( // JSON client will only start if explicitly enabled or no other // output is active at startup if !grpc_client.is_enabled() || stdout_enabled { - stdout::Client::new( - broad_rx.resubscribe(), - running.clone(), - metrics.stdout.clone(), - ) - .start(); + stdout::Client::new(broad_rx, running.clone(), metrics.stdout.clone()).start(); } let mut grpc_handle = grpc_client.start(); tokio::spawn(async move { debug!("Starting output component..."); - loop { + let res = loop { tokio::select! { event = rx.recv() => { let Some(event) = event else { - break; + break Ok(()); }; if let Err(e) = broad_tx.send(Arc::new(event)) { warn!("Failed to forward output event: {e}"); } } + req = subs_rx.recv() => { + let Some(req) = req else { break Ok(()); }; + let rx = broad_tx.subscribe(); + if let Err(e) = req.send(rx) { + break Err(anyhow::anyhow!("Failed to subscribe worker: {e:?}")); + } + } res = grpc_handle.borrow_mut() => { match res { - Ok(Ok(_)) => break, + Ok(Ok(_)) => break Ok(()), Ok(Err(e)) => bail!("gRPC worker errored out: {e:?}"), Err(e) => bail!("gRPC task errored out: {e:?}"), } } - _ = run.changed() => if !*run.borrow() { break; } + _ = run.changed() => if !*run.borrow() { break Ok(()); } } - } + }; debug!("Stopping output component..."); - Ok(()) + res }) } diff --git a/fact/src/output/stdout.rs b/fact/src/output/stdout.rs index debc57fd..020c5534 100644 --- a/fact/src/output/stdout.rs +++ b/fact/src/output/stdout.rs @@ -1,25 +1,16 @@ -use std::sync::Arc; - use log::{info, warn}; -use tokio::sync::{ - broadcast::{self, error::RecvError}, - watch, -}; +use tokio::sync::{broadcast::error::RecvError, watch}; -use crate::{event::Event, metrics::EventCounter}; +use crate::{metrics::EventCounter, output::EventReceiver}; pub struct Client { - rx: broadcast::Receiver>, + rx: EventReceiver, running: watch::Receiver, metrics: EventCounter, } impl Client { - pub fn new( - rx: broadcast::Receiver>, - running: watch::Receiver, - metrics: EventCounter, - ) -> Self { + pub fn new(rx: EventReceiver, running: watch::Receiver, metrics: EventCounter) -> Self { Client { rx, running,