Narrow the CRC scope in file_test

explorer/file_test.cpp

@@ -35,7 +35,7 @@ class ExplorerFileTest : public FileTestBase {
   auto Run(const llvm::SmallVector<llvm::StringRef>& test_args,
            llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem>& fs,
            FILE* /*input_stream*/, llvm::raw_pwrite_stream& output_stream,
-           llvm::raw_pwrite_stream& error_stream)
+           llvm::raw_pwrite_stream& error_stream) const
       -> ErrorOr<RunResult> override {
     // Add the prelude.
     llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> prelude =
@@ -58,21 +58,22 @@ class ExplorerFileTest : public FileTestBase {
       args.push_back(arg.data());
     }
 
+    RawStringOstream trace_stream;
     int exit_code =
         ExplorerMain(args.size(), args.data(), /*install_path=*/"", PreludePath,
                      output_stream, error_stream,
-                     check_trace_output() ? output_stream : trace_stream_, *fs);
+                     check_trace_output() ? output_stream : trace_stream, *fs);
 
     // Skip trace test check as they use stdout stream instead of
     // trace_stream_ostream
-    if (absl::GetFlag(FLAGS_trace) && trace_stream_.TakeStr().empty()) {
+    if (absl::GetFlag(FLAGS_trace) && trace_stream.TakeStr().empty()) {
       return Error("Tracing should always do something");
     }
 
     return {{.success = exit_code == EXIT_SUCCESS}};
   }
 
-  auto GetDefaultArgs() -> llvm::SmallVector<std::string> override {
+  auto GetDefaultArgs() const -> llvm::SmallVector<std::string> override {
     llvm::SmallVector<std::string> args;
     if (absl::GetFlag(FLAGS_trace)) {
       args.push_back("--trace_file=-");
@@ -83,14 +84,15 @@ class ExplorerFileTest : public FileTestBase {
   }
 
   auto GetLineNumberReplacements(llvm::ArrayRef<llvm::StringRef> filenames)
-      -> llvm::SmallVector<LineNumberReplacement> override {
+      const -> llvm::SmallVector<LineNumberReplacement> override {
     if (check_trace_output()) {
       return {};
     }
     return FileTestBase::GetLineNumberReplacements(filenames);
   }
 
-  auto DoExtraCheckReplacements(std::string& check_line) -> void override {
+  auto DoExtraCheckReplacements(std::string& check_line) const
+      -> void override {
     // Ignore the resulting column of EndOfFile because it's often the end of
     // the CHECK comment.
     RE2::GlobalReplace(&check_line, prelude_line_re_,
@@ -106,11 +108,10 @@ class ExplorerFileTest : public FileTestBase {
 
  private:
   // Trace output is directly checked for a few tests.
-  auto check_trace_output() -> bool {
+  auto check_trace_output() const -> bool {
     return test_name().find("/trace/") != std::string::npos;
   }
 
-  RawStringOstream trace_stream_;
   RE2 prelude_line_re_;
   RE2 timing_re_;
 };

testing/file_test/file_test_base.cpp

@@ -268,7 +268,7 @@ auto FileTestCase::TestBody() -> void {
 }
 
 auto FileTestBase::GetLineNumberReplacements(
-    llvm::ArrayRef<llvm::StringRef> filenames)
+    llvm::ArrayRef<llvm::StringRef> filenames) const
     -> llvm::SmallVector<LineNumberReplacement> {
   return {{.has_file = true,
            .re = std::make_shared<RE2>(
@@ -329,88 +329,147 @@ class FileTestEventListener : public testing::EmptyTestEventListener {
   llvm::MutableArrayRef<FileTestInfo> tests_;
 };
 
+// Returns true if the main thread should be used to run tests. This is if
+// either --dump_output is specified, or only 1 thread is needed to run tests.
+static auto SingleThreaded(llvm::ArrayRef<FileTestInfo> tests) -> bool {
+  if (absl::GetFlag(FLAGS_dump_output) || absl::GetFlag(FLAGS_threads) == 1) {
+    return true;
+  }
+
+  bool found_test_to_run = false;
+  for (const auto& test : tests) {
+    if (!test.registered_test->should_run()) {
+      continue;
+    }
+    if (found_test_to_run) {
+      // At least two tests will run, so multi-threaded.
+      return false;
+    }
+    // Found the first test to run.
+    found_test_to_run = true;
+  }
+  // 0 or 1 test will be run, so single-threaded.
+  return false;
+}
+
+// Runs the test in the section that would be inside a lock, possibly inside a
+// CrashRecoveryContext.
+static auto RunSingleTestHelper(FileTestInfo& test, FileTestBase& test_instance)
+    -> void {
+  // Add a crash trace entry with the single-file test command.
+  std::string test_command = GetBazelCommand(BazelMode::Test, test.test_name);
+  llvm::PrettyStackTraceString stack_trace_entry(test_command.c_str());
+
+  if (auto err = RunTestFile(test_instance, absl::GetFlag(FLAGS_dump_output),
+                             **test.test_result);
+      !err.ok()) {
+    test.test_result = std::move(err).error();
+  }
+}
+
+// Runs a single test. Uses a CrashRecoveryContext, and returns false on a
+// crash.
+static auto RunSingleTest(FileTestInfo& test, bool single_threaded,
+                          std::mutex& output_mutex) -> bool {
+  std::unique_ptr<FileTestBase> test_instance(test.factory_fn());
+
+  if (absl::GetFlag(FLAGS_dump_output)) {
+    std::unique_lock<std::mutex> lock(output_mutex);
+    llvm::errs() << "\n--- Dumping: " << test.test_name << "\n\n";
+  }
+
+  // Load expected output.
+  test.test_result = ProcessTestFile(test_instance->test_name(),
+                                     absl::GetFlag(FLAGS_autoupdate));
+  if (test.test_result->ok()) {
+    // Execution must be serialized for either serial tests or console
+    // output.
+    std::unique_lock<std::mutex> output_lock;
+
+    if ((*test.test_result)->capture_console_output ||
+        !test_instance->AllowParallelRun()) {
+      output_lock = std::unique_lock<std::mutex>(output_mutex);
+    }
+
+    if (single_threaded) {
+      RunSingleTestHelper(test, *test_instance);
+    } else {
+      // Use a crash recovery context to try to get a stack trace when
+      // multiple threads may crash in parallel, which otherwise leads to the
+      // program aborting without printing a stack trace.
+      llvm::CrashRecoveryContext crc;
+      crc.DumpStackAndCleanupOnFailure = true;
+      if (!crc.RunSafely([&] { RunSingleTestHelper(test, *test_instance); })) {
+        return false;
+      }
+    }
+  }
+
+  if (!test.test_result->ok()) {
+    std::unique_lock<std::mutex> lock(output_mutex);
+    llvm::errs() << "\n" << test.test_result->error().message() << "\n";
+    return true;
+  }
+
+  test.autoupdate_differs =
+      RunAutoupdater(test_instance.get(), **test.test_result,
+                     /*dry_run=*/!absl::GetFlag(FLAGS_autoupdate));
+
+  std::unique_lock<std::mutex> lock(output_mutex);
+  if (absl::GetFlag(FLAGS_dump_output)) {
+    llvm::outs().flush();
+    const TestFile& test_file = **test.test_result;
+    llvm::errs() << "\n--- Exit with success: "
+                 << (test_file.run_result.success ? "true" : "false")
+                 << "\n--- Autoupdate differs: "
+                 << (test.autoupdate_differs ? "true" : "false") << "\n";
+  } else {
+    llvm::errs() << (test.autoupdate_differs ? "!" : ".");
+  }
+
+  return true;
+}
+
 auto FileTestEventListener::OnTestProgramStart(
     const testing::UnitTest& /*unit_test*/) -> void {
-  llvm::CrashRecoveryContext::Enable();
-  llvm::DefaultThreadPool pool(
-      {.ThreadsRequested = absl::GetFlag(FLAGS_dump_output)
-                               ? 1
-                               : absl::GetFlag(FLAGS_threads)});
+  bool single_threaded = SingleThreaded(tests_);
+
+  std::unique_ptr<llvm::ThreadPoolInterface> pool;
+  if (single_threaded) {
+    pool = std::make_unique<llvm::SingleThreadExecutor>();
+  } else {
+    // Enable the CRC for use in `RunSingleTest`.
+    llvm::CrashRecoveryContext::Enable();
+    pool = std::make_unique<llvm::DefaultThreadPool>(llvm::ThreadPoolStrategy{
+        .ThreadsRequested = absl::GetFlag(FLAGS_threads)});
+  }
   if (!absl::GetFlag(FLAGS_dump_output)) {
-    llvm::errs() << "Running tests with " << pool.getMaxConcurrency()
+    llvm::errs() << "Running tests with " << pool->getMaxConcurrency()
                  << " thread(s)\n";
   }
 
-  // Guard access to both `llvm::errs` and `crashed`.
-  bool crashed = false;
+  // Guard access to output (stdout and stderr).
   std::mutex output_mutex;
+  std::atomic<bool> crashed = false;
 
   for (auto& test : tests_) {
     if (!test.registered_test->should_run()) {
       continue;
     }
 
-    pool.async([&output_mutex, &crashed, &test] {
+    pool->async([&] {
       // If any thread crashed, don't try running more.
-      {
-        std::unique_lock<std::mutex> lock(output_mutex);
-        if (crashed) {
-          return;
-        }
+      if (crashed) {
+        return;
       }
 
-      // Use a crash recovery context to try to get a stack trace when
-      // multiple threads may crash in parallel, which otherwise leads to the
-      // program aborting without printing a stack trace.
-      llvm::CrashRecoveryContext crc;
-      crc.DumpStackAndCleanupOnFailure = true;
-      bool thread_crashed = !crc.RunSafely([&] {
-        std::unique_ptr<FileTestBase> test_instance(test.factory_fn());
-
-        // Add a crash trace entry with the single-file test command.
-        std::string test_command =
-            GetBazelCommand(BazelMode::Test, test.test_name);
-        llvm::PrettyStackTraceString stack_trace_entry(test_command.c_str());
-
-        if (absl::GetFlag(FLAGS_dump_output)) {
-          std::unique_lock<std::mutex> lock(output_mutex);
-          llvm::errs() << "\n--- Dumping: " << test.test_name << "\n\n";
-        }
-
-        test.test_result = ProcessTestFileAndRun(
-            test_instance.get(), &output_mutex,
-            absl::GetFlag(FLAGS_dump_output), absl::GetFlag(FLAGS_autoupdate));
-
-        if (!test.test_result->ok()) {
-          std::unique_lock<std::mutex> lock(output_mutex);
-          llvm::errs() << "\n" << test.test_result->error().message() << "\n";
-          return;
-        }
-
-        test.autoupdate_differs =
-            RunAutoupdater(test_instance.get(), **test.test_result,
-                           /*dry_run=*/!absl::GetFlag(FLAGS_autoupdate));
-
-        std::unique_lock<std::mutex> lock(output_mutex);
-        if (absl::GetFlag(FLAGS_dump_output)) {
-          llvm::outs().flush();
-          const TestFile& test_file = **test.test_result;
-          llvm::errs() << "\n--- Exit with success: "
-                       << (test_file.run_result.success ? "true" : "false")
-                       << "\n--- Autoupdate differs: "
-                       << (test.autoupdate_differs ? "true" : "false") << "\n";
-        } else {
-          llvm::errs() << (test.autoupdate_differs ? "!" : ".");
-        }
-      });
-      if (thread_crashed) {
-        std::unique_lock<std::mutex> lock(output_mutex);
+      if (!RunSingleTest(test, single_threaded, output_mutex)) {
         crashed = true;
       }
     });
   }
 
-  pool.wait();
+  pool->wait();
   if (crashed) {
     // Abort rather than returning so that we don't get a LeakSanitizer report.
     // We expect to have leaked memory if one or more of our tests crashed.

testing/file_test/file_test_base.h

@@ -71,34 +71,35 @@ class FileTestBase {
   virtual auto Run(const llvm::SmallVector<llvm::StringRef>& test_args,
                    llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem>& fs,
                    FILE* input_stream, llvm::raw_pwrite_stream& output_stream,
-                   llvm::raw_pwrite_stream& error_stream)
+                   llvm::raw_pwrite_stream& error_stream) const
       -> ErrorOr<RunResult> = 0;
 
   // Returns default arguments. Only called when a file doesn't set ARGS.
-  virtual auto GetDefaultArgs() -> llvm::SmallVector<std::string> = 0;
+  virtual auto GetDefaultArgs() const -> llvm::SmallVector<std::string> = 0;
 
   // Returns a map of string replacements to implement `%{key}` -> `value` in
   // arguments.
-  virtual auto GetArgReplacements() -> llvm::StringMap<std::string> {
+  virtual auto GetArgReplacements() const -> llvm::StringMap<std::string> {
     return {};
   }
 
   // Returns a regex to match the default file when a line may not be present.
   // May return nullptr if unused. If GetLineNumberReplacements returns an entry
   // with has_file=false, this is required.
   virtual auto GetDefaultFileRE(llvm::ArrayRef<llvm::StringRef> /*filenames*/)
-      -> std::optional<RE2> {
+      const -> std::optional<RE2> {
     return std::nullopt;
   }
 
   // Returns replacement information for line numbers. See LineReplacement for
   // construction.
   virtual auto GetLineNumberReplacements(
-      llvm::ArrayRef<llvm::StringRef> filenames)
+      llvm::ArrayRef<llvm::StringRef> filenames) const
       -> llvm::SmallVector<LineNumberReplacement>;
 
   // Optionally allows children to provide extra replacements for autoupdate.
-  virtual auto DoExtraCheckReplacements(std::string& /*check_line*/) -> void {}
+  virtual auto DoExtraCheckReplacements(std::string& /*check_line*/) const
+      -> void {}
 
   // Whether to allow running the test in parallel, particularly for autoupdate.
   // This can be overridden to force some tests to be run serially. At any given

testing/file_test/file_test_base_test.cpp

@@ -23,25 +23,25 @@ class FileTestBaseTest : public FileTestBase {
   auto Run(const llvm::SmallVector<llvm::StringRef>& test_args,
            llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem>& fs,
            FILE* input_stream, llvm::raw_pwrite_stream& output_stream,
-           llvm::raw_pwrite_stream& error_stream)
+           llvm::raw_pwrite_stream& error_stream) const
       -> ErrorOr<RunResult> override;
 
-  auto GetArgReplacements() -> llvm::StringMap<std::string> override {
+  auto GetArgReplacements() const -> llvm::StringMap<std::string> override {
     return {{"replacement", "replaced"}};
   }
 
-  auto GetDefaultArgs() -> llvm::SmallVector<std::string> override {
+  auto GetDefaultArgs() const -> llvm::SmallVector<std::string> override {
     return {"default_args", "%s"};
   }
 
-  auto GetDefaultFileRE(llvm::ArrayRef<llvm::StringRef> filenames)
+  auto GetDefaultFileRE(llvm::ArrayRef<llvm::StringRef> filenames) const
       -> std::optional<RE2> override {
     return std::make_optional<RE2>(
         llvm::formatv(R"(file: ({0}))", llvm::join(filenames, "|")));
   }
 
   auto GetLineNumberReplacements(llvm::ArrayRef<llvm::StringRef> filenames)
-      -> llvm::SmallVector<LineNumberReplacement> override {
+      const -> llvm::SmallVector<LineNumberReplacement> override {
     auto replacements = FileTestBase::GetLineNumberReplacements(filenames);
     auto filename = std::filesystem::path(test_name().str()).filename();
     if (llvm::StringRef(filename).starts_with("file_only_re_")) {
@@ -260,7 +260,7 @@ auto FileTestBaseTest::Run(
     const llvm::SmallVector<llvm::StringRef>& test_args,
     llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem>& fs,
     FILE* input_stream, llvm::raw_pwrite_stream& output_stream,
-    llvm::raw_pwrite_stream& error_stream) -> ErrorOr<RunResult> {
+    llvm::raw_pwrite_stream& error_stream) const -> ErrorOr<RunResult> {
   PrintArgs(test_args, output_stream);
 
   auto filename = std::filesystem::path(test_name().str()).filename();