[Tooling] Restore working dir in ClangTool.

[ilya-biryukov]

Summary

And add an option to disable this behavior. The option is only used in
AllTUsExecutor to avoid races when running concurrently on multiple
threads.

This fixes PR38869 introduced by r340937.

[ioeric]

lg

include/clang/Tooling/Tooling.h

public:
   /// \returns 0 on success; 1 if any error occurred; 2 if there is no error but
   /// some files are skipped due to missing compile commands.
   int run(ToolAction *Action);
 
   /// Create an AST for each file specified in the command line and
   /// append them to ASTs.
   int buildASTs(std::vector<std::unique_ptr<ASTUnit>> &ASTs);
 
+  /// Sets whether working directory should be restored after calling run(). By
+  /// default, working directory is restored. However, it could be useful to
+  /// turn this off when running on multiple threads to avoid the raciness.
+  void setRestoreWorkingDir(bool RestoreCWD);
+
   /// Returns the file manager used in the tool.
   ///
   /// The file manager is shared between all translation units.
   FileManager &getFiles() { return *Files; }
 
   llvm::ArrayRef<std::string> getSourcePaths() const { return SourcePaths; }
 
 private:
   const CompilationDatabase &Compilations;
   std::vector<std::string> SourcePaths;
   std::shared_ptr<PCHContainerOperations> PCHContainerOps;
 
   llvm::IntrusiveRefCntPtr<vfs::OverlayFileSystem> OverlayFileSystem;
   llvm::IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem;
   llvm::IntrusiveRefCntPtr<FileManager> Files;
 
   // Contains a list of pairs (<file name>, <file content>).
   std::vector<std::pair<StringRef, StringRef>> MappedFileContents;
 
   llvm::StringSet<> SeenWorkingDirectories;
 
   ArgumentsAdjuster ArgsAdjuster;
 
   DiagnosticConsumer *DiagConsumer = nullptr;
+
+  bool RestoreCWD = true;
 };
 
 template <typename T>
 std::unique_ptr<FrontendActionFactory> newFrontendActionFactory() {
   class SimpleFrontendActionFactory : public FrontendActionFactory {
   public:
     FrontendAction *create() override { return new T; }
   };

lib/Tooling/AllTUsExecution.cpp

auto Count = [&]() {
     return ++Counter;
   };
 
   auto &Action = Actions.front();
 
   {
     llvm::ThreadPool Pool(ThreadCount == 0 ? llvm::hardware_concurrency()
                                            : ThreadCount);
-
+    llvm::SmallString<128> InitialWorkingDir;
+    if (auto EC = llvm::sys::fs::current_path(InitialWorkingDir)) {
+      InitialWorkingDir = "";
+      llvm::errs() << "Error while getting current working directory: "
+                   << EC.message() << "\n";
+    }
     for (std::string File : Files) {
       Pool.async(
           [&](std::string Path) {
             Log("[" + std::to_string(Count()) + "/" + TotalNumStr +
                 "] Processing file " + Path);
             ClangTool Tool(Compilations, {Path});
             Tool.appendArgumentsAdjuster(Action.second);
             Tool.appendArgumentsAdjuster(getDefaultArgumentsAdjusters());
             for (const auto &FileAndContent : OverlayFiles)
               Tool.mapVirtualFile(FileAndContent.first(),
                                   FileAndContent.second);
+            // Do not restore working dir from multiple threads to avoid races.
+            Tool.setRestoreWorkingDir(false);
             if (Tool.run(Action.first.get()))
               AppendError(llvm::Twine("Failed to run action on ") + Path +
                           "\n");
           },
           File);
     }
+    if (!InitialWorkingDir.empty()) {
+      if (auto EC = llvm::sys::fs::set_current_path(InitialWorkingDir))
+        llvm::errs() << "Error while restoring working directory: "
+                     << EC.message() << "\n";
+    }
   }
 
   if (!ErrorMsg.empty())
     return make_string_error(ErrorMsg);
 
   return llvm::Error::success();
 }
 

lib/Tooling/Tooling.cpp

if (!AbsPath) {
       llvm::errs() << "Skipping " << SourcePath
                    << ". Error while getting an absolute path: "
                    << llvm::toString(AbsPath.takeError()) << "\n";
       continue;
     }
     AbsolutePaths.push_back(std::move(*AbsPath));
   }
 
+  // Remember the working directory in case we need to restore it.
+  std::string InitialWorkingDir;
+  if (RestoreCWD) {
+    if (auto CWD = OverlayFileSystem->getCurrentWorkingDirectory()) {
+      InitialWorkingDir = std::move(*CWD);
+    } else {
+      llvm::errs() << "Could not get working directory: "
+                   << CWD.getError().message() << "\n";
+    }
+  }
+
   for (llvm::StringRef File : AbsolutePaths) {
     // Currently implementations of CompilationDatabase::getCompileCommands can
     // change the state of the file system (e.g.  prepare generated headers), so
     // this method needs to run right before we invoke the tool, as the next
     // file may require a different (incompatible) state of the file system.
     //
     // FIXME: Make the compilation database interface more explicit about the
     // requirements to the order of invocation of its members.
for (CompileCommand &CompileCommand : CompileCommandsForFile) {
 
       if (!Invocation.run()) {
         // FIXME: Diagnostics should be used instead.
         llvm::errs() << "Error while processing " << File << ".\n";
         ProcessingFailed = true;
       }
     }
   }
+
+  if (!InitialWorkingDir.empty()) {
+    if (auto EC =
+            OverlayFileSystem->setCurrentWorkingDirectory(InitialWorkingDir))
+      llvm::errs() << "Error when trying to restore working dir: "
+                   << EC.message() << "\n";
+  }
   return ProcessingFailed ? 1 : (FileSkipped ? 2 : 0);
 }
 
 namespace {
 
 class ASTBuilderAction : public ToolAction {
   std::vector<std::unique_ptr<ASTUnit>> &ASTs;
 

 
 } // namespace
 
 int ClangTool::buildASTs(std::vector<std::unique_ptr<ASTUnit>> &ASTs) {
   ASTBuilderAction Action(ASTs);
   return run(&Action);
 }
 
+void ClangTool::setRestoreWorkingDir(bool RestoreCWD) {
+  this->RestoreCWD = RestoreCWD;
+}
+
 namespace clang {
 namespace tooling {
 
 std::unique_ptr<ASTUnit>
 buildASTFromCode(const Twine &Code, const Twine &FileName,
                  std::shared_ptr<PCHContainerOperations> PCHContainerOps) {
   return buildASTFromCodeWithArgs(Code, std::vector<std::string>(), FileName,
                                   "clang-tool", std::move(PCHContainerOps));