Introduce channel between subscriber and client

Summary: Rather than making direct calls to HealthCheckClient, send the events from HealthCheckSubscriber over a channel. The channel is bound to 100 events. If the client is unable to keep up, we drop the channel and client and stop the health checks. There are further safety mechanism in further diffs.

Reviewed By: JakobDegen

Differential Revision: D73570880

fbshipit-source-id: 595b83b83fa340b8ab23fa3ddb9008568c22c61c

Author: Rajneesh Lakkundi

app/buck2_client_ctx/src/subscribers/health_check_subscriber.rs

@@ -7,22 +7,31 @@
  * of this source tree.
  */
 
-#![allow(dead_code)] // TODO(rajneeshl): Remove this when we start forwarding events to client
+#![allow(dead_code)]
+// TODO(rajneeshl): Remove this when we start forwarding events to client
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use buck2_core::soft_error;
 use buck2_data::buck_event::Data::*;
 use buck2_error::BuckErrorContext;
 use buck2_events::BuckEvent;
 use buck2_health_check::health_check_client::HealthCheckClient;
+use buck2_health_check::health_check_client::StreamingHealthCheckClient;
+use buck2_health_check::interface::HealthCheckContextEvent;
+use buck2_health_check::interface::HealthCheckEvent;
 use buck2_health_check::report::DisplayReport;
 use tokio::sync::mpsc::Sender;
+use tokio::sync::mpsc::error::TrySendError;
 
 use crate::subscribers::subscriber::EventSubscriber;
 
+const EVENT_CHANNEL_SIZE: usize = 100;
+
 /// This subscriber is responsible for forwarding events to the health check client
 pub struct HealthCheckSubscriber {
-    health_check_client: HealthCheckClient,
+    health_check_client: Option<Box<dyn HealthCheckClient>>,
+    event_sender: Option<Sender<HealthCheckEvent>>,
 }
 
 #[async_trait]
@@ -39,27 +48,36 @@ impl HealthCheckSubscriber {
     pub fn new(
         tags_sender: Sender<Vec<String>>,
         display_reports_sender: Sender<Vec<DisplayReport>>,
+    ) -> Box<Self> {
+        let (tx, rx) = tokio::sync::mpsc::channel(EVENT_CHANNEL_SIZE);
+        let client =
+            StreamingHealthCheckClient::new(Some(tags_sender), Some(display_reports_sender), rx);
+        Self::new_with_client(Some(Box::new(client)), tx)
+    }
+
+    fn new_with_client(
+        health_check_client: Option<Box<dyn HealthCheckClient>>,
+        events_tx: Sender<HealthCheckEvent>,
     ) -> Box<Self> {
         Box::new(Self {
-            health_check_client: HealthCheckClient::new(
-                Some(tags_sender),
-                Some(display_reports_sender),
-            ),
+            health_check_client,
+            event_sender: Some(events_tx),
         })
     }
 
     async fn handle_event(&mut self, event: &Arc<BuckEvent>) -> buck2_error::Result<()> {
-        // Capture all errors in the subscribers and don't propagate them any further.
-        // Health check failures should not affect the command in any way.
+        if self.event_sender.is_none() || self.health_check_client.is_none() {
+            return Ok(());
+        }
 
-        let _res = match event.data() {
+        let health_check_event = match event.data() {
             SpanStart(start) => match &start.data {
                 Some(buck2_data::span_start_event::Data::Command(command)) => {
-                    self.health_check_client
-                        .update_command_data(command.clone())
-                        .await
+                    Some(HealthCheckEvent::HealthCheckContextEvent(
+                        HealthCheckContextEvent::CommandStart(command.clone()),
+                    ))
                 }
-                _ => Ok(()),
+                _ => None,
             },
             SpanEnd(end) => {
                 use buck2_data::span_end_event::Data::*;
@@ -68,25 +86,30 @@ impl HealthCheckSubscriber {
                     .as_ref()
                     .buck_error_context("Missing `data` in SpanEnd")?
                 {
-                    FileWatcher(file_watcher) => {
-                        if let Some(merge_base) = file_watcher
-                            .stats
+                    FileWatcher(file_watcher) => file_watcher
+                        .stats
+                        .as_ref()
+                        .and_then(|stats| stats.branched_from_revision.as_ref())
+                        .map(|merge_base| {
+                            HealthCheckEvent::HealthCheckContextEvent(
+                                HealthCheckContextEvent::BranchedFromRevision(merge_base.clone()),
+                            )
+                        }),
+                    ActionExecution(action_execution_end) => {
+                        let has_excess_cache_miss = action_execution_end
+                            .invalidation_info
                             .as_ref()
-                            .and_then(|stats| stats.branched_from_revision.as_ref())
-                        {
-                            self.health_check_client
-                                .update_branched_from_revision(&merge_base)
-                                .await
+                            .is_some_and(|v| v.changed_file.is_none());
+
+                        if has_excess_cache_miss {
+                            Some(HealthCheckEvent::HealthCheckContextEvent(
+                                HealthCheckContextEvent::HasExcessCacheMisses(),
+                            ))
                         } else {
-                            Ok(())
+                            None
                         }
                     }
-                    ActionExecution(action_execution_end) => {
-                        self.health_check_client
-                            .update_excess_cache_misses(action_execution_end)
-                            .await
-                    }
-                    _ => Ok(()),
+                    _ => None,
                 }
             }
             Instant(instant) => {
@@ -96,23 +119,166 @@ impl HealthCheckSubscriber {
                     .as_ref()
                     .buck_error_context("Missing `data` in `Instant`")?
                 {
-                    SystemInfo(system_info) => {
-                        self.health_check_client
-                            .update_experiment_configurations(system_info.clone())
-                            .await
-                    }
+                    SystemInfo(system_info) => Some(HealthCheckEvent::HealthCheckContextEvent(
+                        HealthCheckContextEvent::ExperimentConfigurations(system_info.clone()),
+                    )),
                     TargetPatterns(target_patterns) => {
-                        self.health_check_client
-                            .update_parsed_target_patterns(target_patterns.clone())
-                            .await
+                        Some(HealthCheckEvent::HealthCheckContextEvent(
+                            HealthCheckContextEvent::ParsedTargetPatterns(target_patterns.clone()),
+                        ))
                     }
-                    Snapshot(snapshot) => self.health_check_client.run_checks(snapshot).await,
-                    _ => Ok(()),
+                    Snapshot(_) => Some(HealthCheckEvent::Snapshot()),
+                    _ => None,
                 }
             }
-            _ => Ok(()),
+            _ => None,
         };
 
+        if let (Some(health_check_event), Some(event_sender)) =
+            (health_check_event, &mut self.event_sender)
+        {
+            match event_sender.try_send(health_check_event) {
+                Ok(_) => {}
+                Err(TrySendError::Full(_)) => {
+                    self.close_client_connection_with_error_report(
+                        "Health check event channel full. Disabling health checks.",
+                    );
+                }
+                Err(TrySendError::Closed(_)) => {
+                    self.close_client_connection_with_error_report(
+                        "Health check event receiver closed. Disabling health checks.",
+                    );
+                }
+            }
+        }
+        Ok(())
+    }
+
+    pub fn close_client_connection_with_error_report(&mut self, error: &str) {
+        self.event_sender.take();
+        self.health_check_client.take();
+        let _ignored = soft_error!(
+            "health_check_subscriber_error",
+            buck2_error::buck2_error!(buck2_error::ErrorTag::HealthCheck, "{}", error)
+        );
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::SystemTime;
+
+    use buck2_health_check::interface::HealthCheckType;
+    use buck2_health_check::report::DisplayReport;
+    use buck2_health_check::report::HealthIssue;
+    use buck2_health_check::report::Severity;
+    use buck2_wrapper_common::invocation_id::TraceId;
+    use tokio::sync::mpsc::Receiver;
+    use tokio::sync::mpsc::{self};
+    use tokio::task::JoinHandle;
+
+    use super::*;
+
+    struct TestHealthCheckClient {
+        handle: JoinHandle<()>,
+    }
+
+    impl TestHealthCheckClient {
+        fn new(
+            tags_tx: Sender<Vec<String>>,
+            display_reports_tx: Sender<Vec<DisplayReport>>,
+            mut event_rx: Receiver<HealthCheckEvent>,
+        ) -> Self {
+            let handle = tokio::spawn(async move {
+                while let Some(event) = event_rx.recv().await {
+                    match event {
+                        HealthCheckEvent::Snapshot() => {
+                            // Send test tags
+                            let _unused = tags_tx
+                                .send(vec!["test_tag1".to_owned(), "test_tag2".to_owned()])
+                                .await;
+
+                            let _unused = display_reports_tx.send(test_reports()).await;
+                        }
+                        HealthCheckEvent::HealthCheckContextEvent(_) => {
+                            // Process context events if needed for tests
+                        }
+                    }
+                }
+            });
+            Self { handle }
+        }
+    }
+
+    impl HealthCheckClient for TestHealthCheckClient {}
+
+    fn test_reports() -> Vec<DisplayReport> {
+        vec![
+            DisplayReport {
+                health_check_type: HealthCheckType::StableRevision,
+                health_issue: Some(HealthIssue {
+                    severity: Severity::Warning,
+                    message: "Test report 1".to_owned(),
+                    remediation: None,
+                }),
+            },
+            DisplayReport {
+                health_check_type: HealthCheckType::LowDiskSpace,
+                health_issue: Some(HealthIssue {
+                    severity: Severity::Info,
+                    message: "Test report 2".to_owned(),
+                    remediation: None,
+                }),
+            },
+        ]
+    }
+
+    fn test_event(data: buck2_data::buck_event::Data) -> Arc<BuckEvent> {
+        Arc::new(BuckEvent::new(
+            SystemTime::now(),
+            TraceId::new(),
+            None,
+            None,
+            data,
+        ))
+    }
+
+    #[tokio::test]
+    async fn test_health_check_subscriber() -> buck2_error::Result<()> {
+        let (tags_tx, mut tags_rx) = mpsc::channel::<Vec<String>>(10);
+        let (reports_tx, mut reports_rx) = mpsc::channel::<Vec<DisplayReport>>(10);
+        let (events_tx, events_rx) = mpsc::channel::<HealthCheckEvent>(EVENT_CHANNEL_SIZE);
+
+        // Create client and subscriber
+        let client = TestHealthCheckClient::new(tags_tx, reports_tx, events_rx);
+        let mut subscriber =
+            HealthCheckSubscriber::new_with_client(Some(Box::new(client)), events_tx);
+
+        // Create a snapshot event
+        let event = test_event(
+            buck2_data::InstantEvent {
+                data: Some(Box::new(buck2_data::Snapshot::default()).into()),
+            }
+            .into(),
+        );
+
+        // Send the event
+        subscriber.handle_event(&event).await?;
+
+        // Verify tags were received
+        let tags =
+            tokio::time::timeout(std::time::Duration::from_millis(100), tags_rx.recv()).await?;
+        assert_eq!(tags.unwrap(), vec!["test_tag1", "test_tag2"]);
+
+        // Verify reports were received
+        let reports =
+            tokio::time::timeout(std::time::Duration::from_millis(100), reports_rx.recv()).await?;
+        assert!(reports.is_some());
+        let reports = reports.unwrap();
+        assert_eq!(reports.len(), 2);
+        assert!(reports[0].health_issue.is_some());
+        assert!(reports[1].health_issue.is_some());
+
         Ok(())
     }
 }