[clang][deps] Ensure DiagnosticConsumer::finish is always called (#127110)

When using the clang dependency scanner with an arbitrary
DiagnosticConsumer, it is important that we always call finish().
Previously, if there was an error preventing us from reaching the
scanning action, or if the command line contained no scannable actions
we would fail to finish(), which would break some consumers (e.g.
serialized diag consumer).

Author: benlangmuir

clang/lib/Tooling/DependencyScanning/DependencyScanningWorker.cpp

@@ -324,15 +324,12 @@ class DependencyScanningAction : public tooling::ToolAction {
 
     // Create the compiler's actual diagnostics engine.
     sanitizeDiagOpts(ScanInstance.getDiagnosticOpts());
+    assert(!DiagConsumerFinished && "attempt to reuse finished consumer");
     ScanInstance.createDiagnostics(DriverFileMgr->getVirtualFileSystem(),
                                    DiagConsumer, /*ShouldOwnClient=*/false);
     if (!ScanInstance.hasDiagnostics())
       return false;
 
-    // Some DiagnosticConsumers require that finish() is called.
-    auto DiagConsumerFinisher =
-        llvm::make_scope_exit([DiagConsumer]() { DiagConsumer->finish(); });
-
     ScanInstance.getPreprocessorOpts().AllowPCHWithDifferentModulesCachePath =
         true;
 
@@ -446,10 +443,11 @@ class DependencyScanningAction : public tooling::ToolAction {
     if (ScanInstance.getDiagnostics().hasErrorOccurred())
       return false;
 
-    // Each action is responsible for calling finish.
-    DiagConsumerFinisher.release();
     const bool Result = ScanInstance.ExecuteAction(*Action);
 
+    // ExecuteAction is responsible for calling finish.
+    DiagConsumerFinished = true;
+
     if (Result)
       setLastCC1Arguments(std::move(OriginalInvocation));
 
@@ -460,6 +458,7 @@ class DependencyScanningAction : public tooling::ToolAction {
   }
 
   bool hasScanned() const { return Scanned; }
+  bool hasDiagConsumerFinished() const { return DiagConsumerFinished; }
 
   /// Take the cc1 arguments corresponding to the most recent invocation used
   /// with this action. Any modifications implied by the discovered dependencies
@@ -491,6 +490,7 @@ class DependencyScanningAction : public tooling::ToolAction {
   std::shared_ptr<ModuleDepCollector> MDC;
   std::vector<std::string> LastCC1Arguments;
   bool Scanned = false;
+  bool DiagConsumerFinished = false;
 };
 
 } // end anonymous namespace
@@ -693,6 +693,11 @@ bool DependencyScanningWorker::scanDependencies(
   if (Success && !Action.hasScanned())
     Diags->Report(diag::err_fe_expected_compiler_job)
         << llvm::join(CommandLine, " ");
+
+  // Ensure finish() is called even if we never reached ExecuteAction().
+  if (!Action.hasDiagConsumerFinished())
+    DC.finish();
+
   return Success && Action.hasScanned();
 }
 

clang/unittests/Tooling/DependencyScanning/DependencyScannerTest.cpp

@@ -15,6 +15,7 @@
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/DependencyScanning/DependencyScanningTool.h"
+#include "clang/Tooling/DependencyScanning/DependencyScanningWorker.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/MC/TargetRegistry.h"
@@ -302,3 +303,84 @@ TEST(DependencyScanner, ScanDepsWithModuleLookup) {
               InterceptFS->StatPaths.end());
   EXPECT_EQ(InterceptFS->ReadFiles, std::vector<std::string>{"test.m"});
 }
+
+TEST(DependencyScanner, ScanDepsWithDiagConsumer) {
+  StringRef CWD = "/root";
+
+  auto VFS = new llvm::vfs::InMemoryFileSystem();
+  VFS->setCurrentWorkingDirectory(CWD);
+  auto Sept = llvm::sys::path::get_separator();
+  std::string HeaderPath =
+      std::string(llvm::formatv("{0}root{0}header.h", Sept));
+  std::string TestPath = std::string(llvm::formatv("{0}root{0}test.cpp", Sept));
+  std::string AsmPath = std::string(llvm::formatv("{0}root{0}test.s", Sept));
+
+  VFS->addFile(HeaderPath, 0, llvm::MemoryBuffer::getMemBuffer("\n"));
+  VFS->addFile(TestPath, 0,
+               llvm::MemoryBuffer::getMemBuffer("#include \"header.h\"\n"));
+  VFS->addFile(AsmPath, 0, llvm::MemoryBuffer::getMemBuffer(""));
+
+  DependencyScanningService Service(ScanningMode::DependencyDirectivesScan,
+                                    ScanningOutputFormat::Make);
+  DependencyScanningWorker Worker(Service, VFS);
+
+  llvm::DenseSet<ModuleID> AlreadySeen;
+  FullDependencyConsumer DC(AlreadySeen);
+  CallbackActionController AC(nullptr);
+
+  struct EnsureFinishedConsumer : public DiagnosticConsumer {
+    bool Finished = false;
+    void finish() override { Finished = true; }
+  };
+
+  {
+    // Check that a successful scan calls DiagConsumer.finish().
+    std::vector<std::string> Args = {"clang",
+                                     "-target",
+                                     "x86_64-apple-macosx10.7",
+                                     "-c",
+                                     "test.cpp",
+                                     "-o"
+                                     "test.cpp.o"};
+
+    EnsureFinishedConsumer DiagConsumer;
+    bool Success = Worker.computeDependencies(CWD, Args, DC, AC, DiagConsumer);
+
+    EXPECT_TRUE(Success);
+    EXPECT_EQ(DiagConsumer.getNumErrors(), 0u);
+    EXPECT_TRUE(DiagConsumer.Finished);
+  }
+
+  {
+    // Check that an invalid command-line, which never enters the scanning
+    // action calls DiagConsumer.finish().
+    std::vector<std::string> Args = {"clang", "-invalid-arg"};
+    EnsureFinishedConsumer DiagConsumer;
+    bool Success = Worker.computeDependencies(CWD, Args, DC, AC, DiagConsumer);
+
+    EXPECT_FALSE(Success);
+    EXPECT_GE(DiagConsumer.getNumErrors(), 1u);
+    EXPECT_TRUE(DiagConsumer.Finished);
+  }
+
+  {
+    // Check that a valid command line that produces no scanning jobs calls
+    // DiagConsumer.finish().
+    std::vector<std::string> Args = {"clang",
+                                     "-target",
+                                     "x86_64-apple-macosx10.7",
+                                     "-c",
+                                     "-x",
+                                     "assembler",
+                                     "test.s",
+                                     "-o"
+                                     "test.cpp.o"};
+
+    EnsureFinishedConsumer DiagConsumer;
+    bool Success = Worker.computeDependencies(CWD, Args, DC, AC, DiagConsumer);
+
+    EXPECT_FALSE(Success);
+    EXPECT_EQ(DiagConsumer.getNumErrors(), 1u);
+    EXPECT_TRUE(DiagConsumer.Finished);
+  }
+}