From 8db5a4c0b6d85660062aeeb235d418dbc57fcfca Mon Sep 17 00:00:00 2001
From: Lucas Kent <rubickent@gmail.com>
Date: Wed, 14 Dec 2022 09:35:47 +1100
Subject: [PATCH] Send panic info over json when json logging enabled (#954)

---
 Cargo.lock                                  |  1 +
 shotover-proxy/Cargo.toml                   |  1 +
 shotover-proxy/src/lib.rs                   |  1 +
 shotover-proxy/src/runner.rs                | 10 +++
 shotover-proxy/src/tracing_panic_handler.rs | 85 +++++++++++++++++++++
 5 files changed, 98 insertions(+)
 create mode 100644 shotover-proxy/src/tracing_panic_handler.rs

diff --git a/Cargo.lock b/Cargo.lock
index 976f6356c..6e0b45171 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2913,6 +2913,7 @@ dependencies = [
  "anyhow",
  "async-recursion",
  "async-trait",
+ "backtrace",
  "base64",
  "bigdecimal 0.3.0",
  "bincode",
diff --git a/shotover-proxy/Cargo.toml b/shotover-proxy/Cargo.toml
index 129d68b76..b646dcd8f 100644
--- a/shotover-proxy/Cargo.toml
+++ b/shotover-proxy/Cargo.toml
@@ -41,6 +41,7 @@ split-iter = "0.1.0"
 # Error handling
 thiserror = "1.0"
 anyhow = "1.0.31"
+backtrace = "0.3.66"
 
 # Parsers
 cql3-parser = "0.3.1"
diff --git a/shotover-proxy/src/lib.rs b/shotover-proxy/src/lib.rs
index 5fa16d016..c802133e8 100644
--- a/shotover-proxy/src/lib.rs
+++ b/shotover-proxy/src/lib.rs
@@ -33,4 +33,5 @@ mod server;
 pub mod sources;
 pub mod tcp;
 pub mod tls;
+pub mod tracing_panic_handler;
 pub mod transforms;
diff --git a/shotover-proxy/src/runner.rs b/shotover-proxy/src/runner.rs
index 889c140b1..c82a017dc 100644
--- a/shotover-proxy/src/runner.rs
+++ b/shotover-proxy/src/runner.rs
@@ -224,6 +224,16 @@ impl TracingState {
             }
         };
 
+        // When in json mode we need to process panics as events instead of printing directly to stdout.
+        // This is so that:
+        // * We dont include invalid json in stdout
+        // * panics can be received by whatever is processing the json events
+        //
+        // We dont do this for LogFormat::Human because the default panic messages are more readable for humans
+        if let LogFormat::Json = format {
+            crate::tracing_panic_handler::setup();
+        }
+
         Ok(TracingState { guard, handle })
     }
 }
diff --git a/shotover-proxy/src/tracing_panic_handler.rs b/shotover-proxy/src/tracing_panic_handler.rs
new file mode 100644
index 000000000..ef5052085
--- /dev/null
+++ b/shotover-proxy/src/tracing_panic_handler.rs
@@ -0,0 +1,85 @@
+use backtrace::{Backtrace, BacktraceFmt, BacktraceFrame, BytesOrWideString, PrintFmt};
+use std::fmt;
+
+pub fn setup() {
+    std::panic::set_hook(Box::new(|panic| {
+        let backtrace = BacktraceFormatter(Backtrace::new());
+        // If the panic has a source location, record it as structured fields.
+        if let Some(location) = panic.location() {
+            tracing::error!(
+                message = %panic,
+                panic.file = location.file(),
+                panic.line = location.line(),
+                panic.column = location.column(),
+                panic.backtrace = format!("{backtrace}"),
+            );
+        } else {
+            tracing::error!(
+                message = %panic,
+                panic.backtrace = format!("{backtrace}"),
+            );
+        }
+    }));
+}
+
+/// The std::backtrace::Backtrace formatting is really noisy because it includes all the pre-main and panic handling frames.
+/// Internal panics have logic to remove that but that is missing from std::backtrace::Backtrace.
+/// https://github.com/rust-lang/rust/issues/105413
+///
+/// As a workaround we use the backtrace crate and manually perform the required formatting
+struct BacktraceFormatter(Backtrace);
+
+// based on https://github.com/rust-lang/backtrace-rs/blob/5be2e8ba9cf6e391c5fa45219fc091b4075eb6be/src/capture.rs#L371
+impl fmt::Display for BacktraceFormatter {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // When printing paths we try to strip the cwd if it exists, otherwise
+        // we just print the path as-is. Note that we also only do this for the
+        // short format, because if it's full we presumably want to print
+        // everything.
+        let cwd = std::env::current_dir();
+        let mut print_path = move |fmt: &mut fmt::Formatter<'_>, path: BytesOrWideString<'_>| {
+            let path = path.into_path_buf();
+            if let Ok(cwd) = &cwd {
+                if let Ok(suffix) = path.strip_prefix(cwd) {
+                    return fmt::Display::fmt(&suffix.display(), fmt);
+                }
+            }
+            fmt::Display::fmt(&path.display(), fmt)
+        };
+
+        let mut f = BacktraceFmt::new(f, PrintFmt::Short, &mut print_path);
+        f.add_context()?;
+        let mut include_frames = false;
+        for frame in self.0.frames() {
+            if frame_is_named(
+                frame,
+                "std::sys_common::backtrace::__rust_begin_short_backtrace",
+            ) {
+                break;
+            }
+
+            if include_frames {
+                f.frame().backtrace_frame(frame)?;
+            }
+
+            if frame_is_named(
+                frame,
+                "std::sys_common::backtrace::__rust_end_short_backtrace",
+            ) {
+                include_frames = true;
+            }
+        }
+        f.finish()
+    }
+}
+
+fn frame_is_named(frame: &BacktraceFrame, name: &str) -> bool {
+    for symbol in frame.symbols() {
+        if let Some(full_name) = symbol.name() {
+            if format!("{}", full_name).starts_with(name) {
+                return true;
+            }
+        }
+    }
+    false
+}