From c98bda1b9c82a2d01dabded1e059f519bb2a3467 Mon Sep 17 00:00:00 2001 From: Mauro Ezequiel Moltrasio Date: Fri, 3 Jul 2026 16:32:45 +0200 Subject: [PATCH] feat(output): change resubscription method The new way for output components to resubscribe the broadcast channel distributing events is to send a oneshot sender back to the main task which will resubscribe from the broadcast sender and forward the receiver back. While this approach might seem overly complicated, it ensures there will be no lingering receivers laying around, which in turn allows for Arc::unwrap_or_clone to properly unwrap all messages when only one output is in use. --- fact/src/output/grpc.rs | 38 ++++++++++++++++++++------------------ fact/src/output/mod.rs | 31 ++++++++++++++++++------------- fact/src/output/stdout.rs | 17 ++++------------- 3 files changed, 42 insertions(+), 44 deletions(-) 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,