rust: intercept wrapper does not fail when reporting fails

Author: rizsotto

rust/intercept/src/bin/wrapper.rs

@@ -19,25 +19,25 @@
 
 extern crate core;
 
-use std::path::PathBuf;
-use anyhow::Result;
+use anyhow::{Context, Result};
 use intercept::reporter::{Reporter, TcpReporter};
+use intercept::KEY_DESTINATION;
+use std::path::{Path, PathBuf};
 
 fn main() -> Result<()> {
+    env_logger::init();
     // Find out what is the executable name the execution was started with
     let executable = std::env::args().next().unwrap();
+    log::info!("Executable as called: {:?}", executable);
     // Read the PATH variable and find the next executable with the same name
-    let real_executable = std::env::var("PATH")?
-        .split(':')
-        .map(|dir| std::path::Path::new(dir).join(&executable))
-        .filter(|path| path.exists())
-        .nth(1)
-        .ok_or_else(|| anyhow::anyhow!("Cannot find the real executable"))?;
-    // TODO: ^ This is a very naive way to find the real executable.
-    //        Make sure we don't call ourselves.
+    let real_executable = next_in_path(&executable)?;
+    log::info!("Executable to call: {:?}", real_executable);
 
     // Report the execution with the real executable
-    report_execution(&real_executable);
+    match into_execution(&real_executable).and_then(report) {
+        Ok(_) => log::info!("Execution reported"),
+        Err(e) => log::error!("Execution reporting failed: {}", e),
+    }
 
     // Execute the real executable with the same arguments
     let status = std::process::Command::new(real_executable)
@@ -47,29 +47,44 @@ fn main() -> Result<()> {
     std::process::exit(status.code().unwrap_or(1));
 }
 
-// TODO: Current error handling is very basic, it just panics on any error.
-//      More sophisticated error handling can be: logging the error and return.
-fn report_execution(path_buf: &PathBuf) {
-    // Get the reporter address from the environment
-    let reporter_address = std::env::var(INTERCEPT_REPORTER_ADDRESS)
-        .expect(format!("${} is not set", INTERCEPT_REPORTER_ADDRESS).as_str());
-    // Create a new reporter
-    let reporter = TcpReporter::new(reporter_address)
-        .expect("Cannot create reporter");
+// TODO: This is a very naive way to find the real executable.
+//        Make sure we don't call ourselves.
+fn next_in_path(executable: &String) -> Result<PathBuf> {
+    let real_executable = std::env::var("PATH")?
+        .split(':')
+        .map(|dir| Path::new(dir).join(&executable))
+        .filter(|path| path.exists())
+        .nth(1)
+        .ok_or_else(|| anyhow::anyhow!("Cannot find the real executable"))?;
+    Ok(real_executable)
+}
 
-    // Report the execution
-    let execution = intercept::Event {
+fn report(execution: intercept::Execution) -> Result<()> {
+    let event = intercept::Event {
         pid: intercept::ProcessId(std::process::id() as u32),
-        execution: intercept::Execution {
-            executable: path_buf.clone(),
-            arguments: std::env::args().collect(),
-            working_dir: std::env::current_dir().expect("Cannot get current directory"),
-            environment: std::env::vars().collect(),
-        },
+        execution,
     };
-    reporter.report(execution)
-        .expect("Cannot report execution");
+
+    // Get the reporter address from the environment
+    std::env::var(KEY_DESTINATION)
+        .with_context(|| format!("${} is missing from the environment", KEY_DESTINATION))
+        // Create a new reporter
+        .and_then(|reporter_address| TcpReporter::new(reporter_address))
+        .with_context(|| "Cannot create TCP execution reporter")
+        // Report the execution
+        .and_then(|reporter| reporter.report(event))
+        .with_context(|| "Sending execution failed")?;
+
+    Ok(())
 }
 
-// declare a const string for the INTERCEPT_REPORTER_ADDRESS environment name
-const INTERCEPT_REPORTER_ADDRESS: &str = "INTERCEPT_REPORTER_ADDRESS";
+fn into_execution(path_buf: &Path) -> Result<intercept::Execution> {
+    std::env::current_dir()
+        .with_context(|| "Cannot get current directory")
+        .map(|working_dir| intercept::Execution {
+            executable: path_buf.to_path_buf(),
+            arguments: std::env::args().collect(),
+            working_dir,
+            environment: std::env::vars().collect(),
+        })
+}

rust/intercept/src/collector.rs

@@ -20,18 +20,32 @@
 use std::net::{SocketAddr, TcpListener, TcpStream};
 
 use crossbeam::channel::Sender;
-use serde::{Deserialize, Serialize};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
 use super::Envelope;
 
-#[derive(Serialize, Deserialize, Debug, PartialEq)]
-pub struct SessionLocator(pub String);
-
+/// Represents the local sink of supervised process events.
+///
+/// The collector is responsible for collecting the events from the reporters.
+///
+/// To share the collector between threads, we use the `Arc` type to wrap the
+/// collector. This way we can clone the collector and send it to other threads.
 pub trait EventCollector {
-    fn address(&self) -> SessionLocator;
+    /// Returns the address of the collector.
+    ///
+    /// The address is in the format of `ip:port`.
+    fn address(&self) -> String;
+
+    /// Collects the events from the reporters.
+    ///
+    /// The events are sent to the given destination channel.
+    ///
+    /// The function returns when the collector is stopped. The collector is stopped
+    /// when the `stop` method invoked (from another thread).
     fn collect(&self, destination: Sender<Envelope>) -> Result<(), anyhow::Error>;
+
+    /// Stops the collector.
     fn stop(&self) -> Result<(), anyhow::Error>;
 }
 
@@ -42,6 +56,10 @@ pub struct EventCollectorOnTcp {
 }
 
 impl EventCollectorOnTcp {
+    /// Creates a new TCP event collector.
+    ///
+    /// The collector listens on a random port on the loopback interface.
+    /// The address of the collector can be obtained by the `address` method.
     pub fn new() -> Result<Self, anyhow::Error> {
         let shutdown = Arc::new(AtomicBool::new(false));
         let listener = TcpListener::bind("127.0.0.1:0")?;
@@ -69,12 +87,18 @@ impl EventCollectorOnTcp {
 }
 
 impl EventCollector for EventCollectorOnTcp {
-    fn address(&self) -> SessionLocator {
-        SessionLocator(self.address.to_string())
+    fn address(&self) -> String {
+        self.address.to_string()
     }
 
+    /// Single-threaded implementation of the collector.
+    ///
+    /// The collector listens on the TCP port and accepts incoming connections.
+    /// When a connection is accepted, the collector reads the events from the
+    /// connection and sends them to the destination channel.
     fn collect(&self, destination: Sender<Envelope>) -> Result<(), anyhow::Error> {
         for stream in self.listener.incoming() {
+            // This has to be the first thing to do, in order to implement the stop method!
             if self.shutdown.load(Ordering::Relaxed) {
                 break;
             }
@@ -97,6 +121,11 @@ impl EventCollector for EventCollectorOnTcp {
         Ok(())
     }
 
+    /// Stops the collector by flipping the shutdown flag and connecting to the collector.
+    ///
+    /// The collector is stopped when the `collect` method sees the shutdown flag.
+    /// To signal the collector to stop, we connect to the collector to unblock the
+    /// `accept` call to check the shutdown flag.
     fn stop(&self) -> Result<(), anyhow::Error> {
         self.shutdown.store(true, Ordering::Relaxed);
         let _ = TcpStream::connect(self.address)?;

rust/intercept/src/lib.rs

@@ -21,6 +21,7 @@ use std::io::{Read, Write};
 use std::path::PathBuf;
 
 use chrono::Utc;
+use rand::random;
 use serde::{Deserialize, Serialize};
 
 pub mod collector;
@@ -34,6 +35,13 @@ pub mod reporter;
 #[derive(Serialize, Deserialize, Debug, PartialEq, Clone)]
 pub struct ReporterId(pub u64);
 
+impl ReporterId {
+    pub fn new() -> Self {
+        let id = random::<u64>();
+        ReporterId(id)
+    }
+}
+
 /// Process id is a OS identifier for a process.
 #[derive(Serialize, Deserialize, Debug, PartialEq, Clone)]
 pub struct ProcessId(pub u32);
@@ -113,3 +121,6 @@ impl Envelope {
         Ok(length)
     }
 }
+
+/// Declare the environment variable name for the reporter address.
+pub const KEY_DESTINATION: &str = "INTERCEPT_REPORTER_ADDRESS";

rust/intercept/src/reporter.rs

@@ -19,22 +19,11 @@
 
 use std::net::TcpStream;
 
-use rand::random;
-
 use super::{Envelope, Event, ReporterId};
 
-impl ReporterId {
-    pub fn new() -> Self {
-        let id = random::<u64>();
-        ReporterId(id)
-    }
-}
-
-// Represents the remote sink of supervised process events.
-//
-// Events from a process execution can be sent from many actors (mostly
-// supervisor processes). The events are collected in a common place
-// in order to reconstruct of final report of a build process.
+/// Represents the remote sink of supervised process events.
+///
+/// This allows the reporters to send events to a remote collector.
 pub trait Reporter {
     fn report(&self, event: Event) -> Result<(), anyhow::Error>;
 }
@@ -45,6 +34,10 @@ pub struct TcpReporter {
 }
 
 impl TcpReporter {
+    /// Creates a new TCP reporter instance.
+    ///
+    /// It does not open the TCP connection yet. Stores the destination
+    /// address and creates a unique reporter id.
     pub fn new(destination: String) -> Result<Self, anyhow::Error> {
         let reporter_id = ReporterId::new();
         let result = TcpReporter {
@@ -56,6 +49,10 @@ impl TcpReporter {
 }
 
 impl Reporter for TcpReporter {
+    /// Sends an event to the remote collector.
+    ///
+    /// The event is wrapped in an envelope and sent to the remote collector.
+    /// The TCP connection is opened and closed for each event.
     fn report(&self, event: Event) -> Result<(), anyhow::Error> {
         let envelope = Envelope::new(&self.reporter_id, event);
         let mut socket = TcpStream::connect(self.destination.clone())?;

rust/intercept/tests/test.rs

@@ -23,7 +23,7 @@ mod test {
     #[test]
     fn tcp_reporter_and_collectors_work() {
         let collector = EventCollectorOnTcp::new().unwrap();
-        let reporter = TcpReporter::new(collector.address().0).unwrap();
+        let reporter = TcpReporter::new(collector.address()).unwrap();
 
         // Create wrapper to share the collector across threads.
         let thread_collector = Arc::new(collector);