From ee383638bcaf3b2352d6d78559abeb2dbbbe941e Mon Sep 17 00:00:00 2001
From: Jon Ross-Perkins <jperkins@google.com>
Date: Wed, 25 Sep 2024 14:00:36 -0700
Subject: [PATCH] Flush pending diagnostics on crash. (#4337)

This risks diagnsotic formatting crashing, but I think we more
frequently see cases where it'd be interesting to know what diagnostics
were being delayed as part of the default sorting.
---
 toolchain/diagnostics/diagnostic_consumer.h |  2 ++
 toolchain/driver/BUILD                      |  1 +
 toolchain/driver/compile_subcommand.cpp     | 24 +++++++++++++++++++++
 3 files changed, 27 insertions(+)

diff --git a/toolchain/diagnostics/diagnostic_consumer.h b/toolchain/diagnostics/diagnostic_consumer.h
index 9ffb984939b9e..e43f2af12f229 100644
--- a/toolchain/diagnostics/diagnostic_consumer.h
+++ b/toolchain/diagnostics/diagnostic_consumer.h
@@ -45,6 +45,8 @@ class StreamDiagnosticConsumer : public DiagnosticConsumer {
   auto HandleDiagnostic(Diagnostic diagnostic) -> void override;
   auto Flush() -> void override { stream_->flush(); }
 
+  auto set_stream(llvm::raw_ostream* stream) -> void { stream_ = stream; }
+
  private:
   auto Print(const DiagnosticMessage& message, llvm::StringRef prefix) -> void;
 
diff --git a/toolchain/driver/BUILD b/toolchain/driver/BUILD
index 64a73d60ccc3a..e0d87e6f12802 100644
--- a/toolchain/driver/BUILD
+++ b/toolchain/driver/BUILD
@@ -103,6 +103,7 @@ cc_library(
         "//common:ostream",
         "//common:version",
         "//common:vlog",
+        "//toolchain/base:pretty_stack_trace_function",
         "//toolchain/base:value_store",
         "//toolchain/check",
         "//toolchain/codegen",
diff --git a/toolchain/driver/compile_subcommand.cpp b/toolchain/driver/compile_subcommand.cpp
index d944795cfa2ba..ddf2f93da8bce 100644
--- a/toolchain/driver/compile_subcommand.cpp
+++ b/toolchain/driver/compile_subcommand.cpp
@@ -6,6 +6,7 @@
 
 #include "common/vlog.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "toolchain/base/pretty_stack_trace_function.h"
 #include "toolchain/check/check.h"
 #include "toolchain/codegen/codegen.h"
 #include "toolchain/diagnostics/sorting_diagnostic_consumer.h"
@@ -485,6 +486,10 @@ class CompilationUnit {
     consumer_->Flush();
   }
 
+  // Flushes diagnostics, specifically as part of generating stack trace
+  // information.
+  auto FlushForStackTrace() -> void { consumer_->Flush(); }
+
   auto input_filename() -> llvm::StringRef { return input_filename_; }
   auto success() -> bool { return success_; }
   auto has_source() -> bool { return source_.has_value(); }
@@ -660,6 +665,25 @@ auto CompileSubcommand::Run(DriverEnv& driver_env) -> DriverResult {
     stream_consumer.Flush();
   });
 
+  PrettyStackTraceFunction flush_on_crash([&](llvm::raw_ostream& out) {
+    // When crashing, flush diagnostics. If sorting diagnostics, they can be
+    // redirected to the crash stream; if streaming, the original stream is
+    // flushed.
+    // TODO: Eventually we'll want to limit the count.
+    if (options_.stream_errors) {
+      out << "Flushing diagnostics\n";
+    } else {
+      out << "Pending diagnostics:\n";
+      stream_consumer.set_stream(&out);
+    }
+
+    for (auto& unit : units) {
+      unit->FlushForStackTrace();
+    }
+    stream_consumer.Flush();
+    stream_consumer.set_stream(&driver_env.error_stream);
+  });
+
   // Returns a DriverResult object. Called whenever Compile returns.
   auto make_result = [&]() {
     DriverResult result = {.success = true};