feat(output): change resubscription method (#1011)

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.

Author: Molter73

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,21 +104,21 @@ impl From<&BackoffConfig> for Backoff {
 }
 
 pub struct Client {
-    rx: broadcast::Receiver<Arc<Event>>,
+    subscriber: mpsc::Sender<oneshot::Sender<EventReceiver>>,
     running: watch::Receiver<bool>,
     config: watch::Receiver<GrpcConfig>,
     metrics: EventCounter,
 }
 
 impl Client {
     pub fn new(
-        rx: broadcast::Receiver<Arc<Event>>,
+        subscriber: mpsc::Sender<oneshot::Sender<EventReceiver>>,
         running: watch::Receiver<bool>,
         metrics: EventCounter,
         config: watch::Receiver<GrpcConfig>,
     ) -> 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) => {

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<Arc<Event>>;
+
 /// 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<anyhow::Result<()>> {
     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
     })
 }

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<Arc<Event>>,
+    rx: EventReceiver,
     running: watch::Receiver<bool>,
     metrics: EventCounter,
 }
 
 impl Client {
-    pub fn new(
-        rx: broadcast::Receiver<Arc<Event>>,
-        running: watch::Receiver<bool>,
-        metrics: EventCounter,
-    ) -> Self {
+    pub fn new(rx: EventReceiver, running: watch::Receiver<bool>, metrics: EventCounter) -> Self {
         Client {
             rx,
             running,