From aa0194f840c83942a9070fc342b5bd413ee29712 Mon Sep 17 00:00:00 2001
From: Nicolas Avrutin <rasputin@google.com>
Date: Sat, 15 Jul 2023 21:51:46 -0400
Subject: [PATCH] subscriber: add with_binary_name and with_process_id options.

These options are useful when multiple processes are logging to the same
destination (stderr, in the common case). This can happen when a process
launches a child process or when a user simply launches multiple
processes in the same terminal.

Fixes #2447.

Implement these two options.
---
 tracing-subscriber/src/fmt/fmt_subscriber.rs |  30 ++++++
 tracing-subscriber/src/fmt/format/mod.rs     | 107 +++++++++++++++++++
 tracing-subscriber/src/fmt/mod.rs            |  30 ++++++
 3 files changed, 167 insertions(+)

diff --git a/tracing-subscriber/src/fmt/fmt_subscriber.rs b/tracing-subscriber/src/fmt/fmt_subscriber.rs
index 366660fb82..c47301413e 100644
--- a/tracing-subscriber/src/fmt/fmt_subscriber.rs
+++ b/tracing-subscriber/src/fmt/fmt_subscriber.rs
@@ -490,6 +490,36 @@ where
         }
     }
 
+    /// Sets whether or not the binary name of the process is displayed when
+    /// formatting events. If binary_name is None, argv[0] will be used,
+    /// otherwise the spcified string will be used.
+    ///
+    /// [argv\[0\]]: std::env::args
+    pub fn with_binary_name(
+        self,
+        display_binary_name: bool,
+        binary_name: Option<&str>,
+    ) -> Subscriber<C, N, format::Format<L, T>, W> {
+        Subscriber {
+            fmt_event: self.fmt_event.with_binary_name(display_binary_name, binary_name),
+            ..self
+        }
+    }
+
+    /// Sets whether or not the [process ID] of the current thread is displayed
+    /// when formatting events.
+    ///
+    /// [process ID]: std::process::id
+    pub fn with_process_id(
+        self,
+        display_process_id: bool,
+    ) -> Subscriber<C, N, format::Format<L, T>, W> {
+        Subscriber {
+            fmt_event: self.fmt_event.with_process_id(display_process_id),
+            ..self
+        }
+    }
+
     /// Sets whether or not the [thread ID] of the current thread is displayed
     /// when formatting events.
     ///
diff --git a/tracing-subscriber/src/fmt/format/mod.rs b/tracing-subscriber/src/fmt/format/mod.rs
index c9760138ab..68de0fe884 100644
--- a/tracing-subscriber/src/fmt/format/mod.rs
+++ b/tracing-subscriber/src/fmt/format/mod.rs
@@ -410,6 +410,9 @@ pub struct Format<F = Full, T = SystemTime> {
     pub(crate) display_timestamp: bool,
     pub(crate) display_target: bool,
     pub(crate) display_level: bool,
+    pub(crate) display_binary_name: bool,
+    pub(crate) binary_name: Option<String>,
+    pub(crate) display_process_id: bool,
     pub(crate) display_thread_id: bool,
     pub(crate) display_thread_name: bool,
     pub(crate) display_filename: bool,
@@ -588,6 +591,9 @@ impl Default for Format<Full, SystemTime> {
             display_timestamp: true,
             display_target: true,
             display_level: true,
+            display_binary_name: false,
+            binary_name: None,
+            display_process_id: false,
             display_thread_id: false,
             display_thread_name: false,
             display_filename: false,
@@ -608,6 +614,9 @@ impl<F, T> Format<F, T> {
             display_target: false,
             display_timestamp: self.display_timestamp,
             display_level: self.display_level,
+            display_binary_name: self.display_binary_name,
+            binary_name: self.binary_name,
+            display_process_id: self.display_process_id,
             display_thread_id: self.display_thread_id,
             display_thread_name: self.display_thread_name,
             display_filename: self.display_filename,
@@ -647,6 +656,9 @@ impl<F, T> Format<F, T> {
             display_target: self.display_target,
             display_timestamp: self.display_timestamp,
             display_level: self.display_level,
+            display_binary_name: self.display_binary_name,
+            binary_name: self.binary_name,
+            display_process_id: self.display_process_id,
             display_thread_id: self.display_thread_id,
             display_thread_name: self.display_thread_name,
             display_filename: true,
@@ -679,6 +691,9 @@ impl<F, T> Format<F, T> {
             display_target: self.display_target,
             display_timestamp: self.display_timestamp,
             display_level: self.display_level,
+            display_binary_name: self.display_binary_name,
+            binary_name: self.binary_name,
+            display_process_id: self.display_process_id,
             display_thread_id: self.display_thread_id,
             display_thread_name: self.display_thread_name,
             display_filename: self.display_filename,
@@ -708,6 +723,9 @@ impl<F, T> Format<F, T> {
             display_target: self.display_target,
             display_timestamp: self.display_timestamp,
             display_level: self.display_level,
+            display_binary_name: self.display_binary_name,
+            binary_name: self.binary_name,
+            display_process_id: self.display_process_id,
             display_thread_id: self.display_thread_id,
             display_thread_name: self.display_thread_name,
             display_filename: self.display_filename,
@@ -724,6 +742,9 @@ impl<F, T> Format<F, T> {
             display_timestamp: false,
             display_target: self.display_target,
             display_level: self.display_level,
+            display_binary_name: self.display_binary_name,
+            binary_name: self.binary_name,
+            display_process_id: self.display_process_id,
             display_thread_id: self.display_thread_id,
             display_thread_name: self.display_thread_name,
             display_filename: self.display_filename,
@@ -755,6 +776,33 @@ impl<F, T> Format<F, T> {
         }
     }
 
+    /// Sets whether or not the binary name of the process is displayed when
+    /// formatting events. If binary_name is None, argv[0] will be used,
+    /// otherwise the spcified string will be used.
+    ///
+    /// [argv\[0\]]: std::env::args
+    pub fn with_binary_name(self,
+                            display_binary_name: bool,
+                            binary_name: Option<&str>,
+    ) -> Format<F, T> {
+        Format {
+            display_binary_name,
+            binary_name: binary_name.map(str::to_string),
+            ..self
+        }
+    }
+
+    /// Sets whether or not the [process ID] of the current thread is displayed
+    /// when formatting events.
+    ///
+    /// [process ID]: std::process::id
+    pub fn with_process_id(self, display_process_id: bool) -> Format<F, T> {
+        Format {
+            display_process_id,
+            ..self
+        }
+    }
+
     /// Sets whether or not the [thread ID] of the current thread is displayed
     /// when formatting events.
     ///
@@ -940,6 +988,20 @@ where
         self.format_timestamp(&mut writer)?;
         self.format_level(*meta.level(), &mut writer)?;
 
+        if self.display_binary_name {
+            if let Some(binary_name) = &self.binary_name {
+                write!(writer, "{} ", binary_name)?;
+            } else if let Some(argv0) = std::env::args().next() {
+                write!(writer, "{} ", argv0)?;
+            } else {
+                write!(writer, "<missing_argv[0] ")?;
+            }
+        }
+
+        if self.display_process_id {
+            write!(writer, "PID({}) ", std::process::id())?;
+        }
+
         if self.display_thread_name {
             let current_thread = std::thread::current();
             match current_thread.name() {
@@ -1662,6 +1724,8 @@ pub(super) mod test {
             .without_time()
             .with_level(false)
             .with_target(false)
+            .with_binary_name(false, None)
+            .with_process_id(false)
             .with_thread_ids(false)
             .with_thread_names(false);
         #[cfg(feature = "ansi")]
@@ -1777,6 +1841,49 @@ pub(super) mod test {
         assert_info_hello(subscriber, make_writer, expected);
     }
 
+    #[test]
+    fn with_automatic_binary_name() {
+        let make_writer = MockMakeWriter::default();
+        let subscriber = crate::fmt::Collector::builder()
+            .with_writer(make_writer.clone())
+            .with_binary_name(true, None)
+            .with_ansi(false)
+            .with_timer(MockTime);
+        let expected =
+            format!("fake time  INFO {} tracing_subscriber::fmt::format::test: hello\n",
+                    std::env::args().next().unwrap());
+
+        assert_info_hello(subscriber, make_writer, &expected);
+    }
+
+    #[test]
+    fn with_manual_binary_name() {
+        let make_writer = MockMakeWriter::default();
+        let subscriber = crate::fmt::Collector::builder()
+            .with_writer(make_writer.clone())
+            .with_binary_name(true, Some("a_nice_rust_binary"))
+            .with_ansi(false)
+            .with_timer(MockTime);
+        let expected =
+            "fake time  INFO a_nice_rust_binary tracing_subscriber::fmt::format::test: hello\n";
+
+        assert_info_hello(subscriber, make_writer, expected);
+    }
+
+    #[test]
+    fn with_process_id() {
+        let make_writer = MockMakeWriter::default();
+        let subscriber = crate::fmt::Collector::builder()
+            .with_writer(make_writer.clone())
+            .with_process_id(true)
+            .with_ansi(false)
+            .with_timer(MockTime);
+        let expected =
+            "fake time  INFO PID(NUMERIC) tracing_subscriber::fmt::format::test: hello\n";
+
+        assert_info_hello_ignore_numeric(subscriber, make_writer, expected);
+    }
+
     #[test]
     fn with_thread_ids() {
         let make_writer = MockMakeWriter::default();
diff --git a/tracing-subscriber/src/fmt/mod.rs b/tracing-subscriber/src/fmt/mod.rs
index 38acf4d092..72a9e9817d 100644
--- a/tracing-subscriber/src/fmt/mod.rs
+++ b/tracing-subscriber/src/fmt/mod.rs
@@ -703,6 +703,36 @@ where
         }
     }
 
+    /// Sets whether or not the binary name of the process is displayed when
+    /// formatting events. If binary_name is None, argv[0] will be used,
+    /// otherwise the spcified string will be used.
+    ///
+    /// [argv\[0\]]: std::env::args
+    pub fn with_binary_name(
+        self,
+        display_binary_name: bool,
+        binary_name: Option<&str>,
+    ) -> CollectorBuilder<N, format::Format<L, T>, F, W> {
+        CollectorBuilder {
+            inner: self.inner.with_binary_name(display_binary_name, binary_name),
+            ..self
+        }
+    }
+
+    /// Sets whether or not the [process ID] of the current thread is displayed
+    /// when formatting events.
+    ///
+    /// [process ID]: std::process::id
+    pub fn with_process_id(
+        self,
+        display_process_id: bool,
+    ) -> CollectorBuilder<N, format::Format<L, T>, F, W> {
+        CollectorBuilder {
+            inner: self.inner.with_process_id(display_process_id),
+            ..self
+        }
+    }
+
     /// Sets whether or not the [name] of the current thread is displayed
     /// when formatting events.
     ///