Merge pull request #24008 from allevato/batch-mode-response-files (#24829)

Fix response file support for batch jobs.

Author: allevato

include/swift/Driver/Job.h

@@ -244,6 +244,20 @@ class Job {
     NewlyAdded
   };
 
+  /// Packs together information about response file usage for a job.
+  ///
+  /// The strings in this struct must be kept alive as long as the Job is alive
+  /// (e.g., by calling MakeArgString on the arg list associated with the
+  /// Compilation).
+  struct ResponseFileInfo {
+    /// The path to the response file that a job should use.
+    const char *path;
+
+    /// The '@'-prefixed argument string that should be passed to the tool to
+    /// use the response file.
+    const char *argString;
+  };
+
   using EnvironmentVector = std::vector<std::pair<const char *, const char *>>;
 
   /// If positive, contains llvm::ProcessID for a real Job on the host OS. If
@@ -279,34 +293,25 @@ class Job {
   /// Whether the job wants a list of input or output files created.
   std::vector<FilelistInfo> FilelistFileInfos;
 
-  /// Response file path
-  const char *ResponseFilePath;
-
-  /// This contains a single argument pointing to the response file path with
-  /// the '@' prefix.
-  /// The argument string must be kept alive as long as the Job is alive.
-  const char *ResponseFileArg;
+  /// The path and argument string to use for the response file if the job's
+  /// arguments should be passed using one.
+  Optional<ResponseFileInfo> ResponseFile;
 
   /// The modification time of the main input file, if any.
   llvm::sys::TimePoint<> InputModTime = llvm::sys::TimePoint<>::max();
 
 public:
-  Job(const JobAction &Source,
-      SmallVectorImpl<const Job *> &&Inputs,
-      std::unique_ptr<CommandOutput> Output,
-      const char *Executable,
+  Job(const JobAction &Source, SmallVectorImpl<const Job *> &&Inputs,
+      std::unique_ptr<CommandOutput> Output, const char *Executable,
       llvm::opt::ArgStringList Arguments,
       EnvironmentVector ExtraEnvironment = {},
       std::vector<FilelistInfo> Infos = {},
-      const char *ResponseFilePath = nullptr,
-      const char *ResponseFileArg = nullptr)
+      Optional<ResponseFileInfo> ResponseFile = None)
       : SourceAndCondition(&Source, Condition::Always),
         Inputs(std::move(Inputs)), Output(std::move(Output)),
         Executable(Executable), Arguments(std::move(Arguments)),
         ExtraEnvironment(std::move(ExtraEnvironment)),
-        FilelistFileInfos(std::move(Infos)),
-        ResponseFilePath(ResponseFilePath),
-        ResponseFileArg(ResponseFileArg) {}
+        FilelistFileInfos(std::move(Infos)), ResponseFile(ResponseFile) {}
 
   virtual ~Job();
 
@@ -316,7 +321,10 @@ class Job {
 
   const char *getExecutable() const { return Executable; }
   const llvm::opt::ArgStringList &getArguments() const { return Arguments; }
-  ArrayRef<const char *> getResponseFileArg() const { return ResponseFileArg; }
+  ArrayRef<const char *> getResponseFileArg() const {
+    assert(hasResponseFile());
+    return ResponseFile->argString;
+  }
   ArrayRef<FilelistInfo> getFilelistInfos() const { return FilelistFileInfos; }
   ArrayRef<const char *> getArgumentsForTaskExecution() const;
 
@@ -370,7 +378,7 @@ class Job {
   static void printArguments(raw_ostream &Stream,
                              const llvm::opt::ArgStringList &Args);
 
-  bool hasResponseFile() const { return ResponseFilePath != nullptr; }
+  bool hasResponseFile() const { return ResponseFile.hasValue(); }
 
   bool writeArgsToResponseFile() const;
 };
@@ -398,9 +406,9 @@ class BatchJob : public Job {
   BatchJob(const JobAction &Source, SmallVectorImpl<const Job *> &&Inputs,
            std::unique_ptr<CommandOutput> Output, const char *Executable,
            llvm::opt::ArgStringList Arguments,
-           EnvironmentVector ExtraEnvironment,
-           std::vector<FilelistInfo> Infos,
-           ArrayRef<const Job *> Combined, Job::PID &NextQuasiPID);
+           EnvironmentVector ExtraEnvironment, std::vector<FilelistInfo> Infos,
+           ArrayRef<const Job *> Combined, Job::PID &NextQuasiPID,
+           Optional<ResponseFileInfo> ResponseFile = None);
 
   ArrayRef<const Job*> getCombinedJobs() const {
     return CombinedJobs;

include/swift/Driver/ToolChain.h

@@ -209,6 +209,14 @@ class ToolChain {
   /// set to match the behavior of Clang.
   virtual bool shouldStoreInvocationInDebugInfo() const { return false; }
 
+  /// Gets the response file path and command line argument for an invocation
+  /// if the tool supports response files and if the command line length would
+  /// exceed system limits.
+  Optional<Job::ResponseFileInfo>
+  getResponseFileInfo(const Compilation &C, const char *executablePath,
+                      const InvocationInfo &invocationInfo,
+                      const JobContext &context) const;
+
 public:
   virtual ~ToolChain() = default;
 

lib/Driver/Job.cpp

@@ -363,7 +363,7 @@ void Job::printCommandLine(raw_ostream &os, StringRef Terminator) const {
   escapeAndPrintString(os, Executable);
   os << ' ';
   if (hasResponseFile()) {
-    printArguments(os, {ResponseFileArg});
+    printArguments(os, {ResponseFile->argString});
     os << " # ";
   }
   printArguments(os, Arguments);
@@ -419,8 +419,9 @@ void Job::printSummary(raw_ostream &os) const {
 }
 
 bool Job::writeArgsToResponseFile() const {
+  assert(hasResponseFile());
   std::error_code EC;
-  llvm::raw_fd_ostream OS(ResponseFilePath, EC, llvm::sys::fs::F_None);
+  llvm::raw_fd_ostream OS(ResponseFile->path, EC, llvm::sys::fs::F_None);
   if (EC) {
     return true;
   }
@@ -438,9 +439,10 @@ BatchJob::BatchJob(const JobAction &Source,
                    const char *Executable, llvm::opt::ArgStringList Arguments,
                    EnvironmentVector ExtraEnvironment,
                    std::vector<FilelistInfo> Infos,
-                   ArrayRef<const Job *> Combined, int64_t &NextQuasiPID)
+                   ArrayRef<const Job *> Combined, int64_t &NextQuasiPID,
+                   Optional<ResponseFileInfo> ResponseFile)
     : Job(Source, std::move(Inputs), std::move(Output), Executable, Arguments,
-          ExtraEnvironment, Infos),
+          ExtraEnvironment, Infos, ResponseFile),
       CombinedJobs(Combined.begin(), Combined.end()),
       QuasiPIDBase(NextQuasiPID) {
 

lib/Driver/ToolChain.cpp

@@ -62,6 +62,27 @@ ToolChain::JobContext::getTemporaryFilePath(const llvm::Twine &name,
   return C.getArgs().MakeArgString(buffer.str());
 }
 
+Optional<Job::ResponseFileInfo>
+ToolChain::getResponseFileInfo(const Compilation &C, const char *executablePath,
+                               const ToolChain::InvocationInfo &invocationInfo,
+                               const ToolChain::JobContext &context) const {
+  const bool forceResponseFiles =
+      C.getArgs().hasArg(options::OPT_driver_force_response_files);
+  assert((invocationInfo.allowsResponseFiles || !forceResponseFiles) &&
+         "Cannot force response file if platform does not allow it");
+
+  if (forceResponseFiles || (invocationInfo.allowsResponseFiles &&
+                             !llvm::sys::commandLineFitsWithinSystemLimits(
+                                 executablePath, invocationInfo.Arguments))) {
+    const char *responseFilePath =
+        context.getTemporaryFilePath("arguments", "resp");
+    const char *responseFileArg =
+        C.getArgs().MakeArgString(Twine("@") + responseFilePath);
+    return {{responseFilePath, responseFileArg}};
+  }
+  return None;
+}
+
 std::unique_ptr<Job> ToolChain::constructJob(
     const JobAction &JA, Compilation &C, SmallVectorImpl<const Job *> &&inputs,
     ArrayRef<const Action *> inputActions,
@@ -114,28 +135,16 @@ std::unique_ptr<Job> ToolChain::constructJob(
     }
   }
 
-  const char *responseFilePath = nullptr;
-  const char *responseFileArg = nullptr;
-
-  const bool forceResponseFiles =
-      C.getArgs().hasArg(options::OPT_driver_force_response_files);
-  assert((invocationInfo.allowsResponseFiles || !forceResponseFiles) &&
-         "Cannot force response file if platform does not allow it");
-
-  if (forceResponseFiles || (invocationInfo.allowsResponseFiles &&
-                             !llvm::sys::commandLineFitsWithinSystemLimits(
-                                 executablePath, invocationInfo.Arguments))) {
-    responseFilePath = context.getTemporaryFilePath("arguments", "resp");
-    responseFileArg = C.getArgs().MakeArgString(Twine("@") + responseFilePath);
-  }
+  // Determine if the argument list is so long that it needs to be written into
+  // a response file.
+  auto responseFileInfo =
+      getResponseFileInfo(C, executablePath, invocationInfo, context);
 
-  return llvm::make_unique<Job>(JA, std::move(inputs), std::move(output),
-                                executablePath,
-                                std::move(invocationInfo.Arguments),
-                                std::move(invocationInfo.ExtraEnvironment),
-                                std::move(invocationInfo.FilelistInfos),
-                                responseFilePath,
-                                responseFileArg);
+  return llvm::make_unique<Job>(
+      JA, std::move(inputs), std::move(output), executablePath,
+      std::move(invocationInfo.Arguments),
+      std::move(invocationInfo.ExtraEnvironment),
+      std::move(invocationInfo.FilelistInfos), responseFileInfo);
 }
 
 std::string
@@ -345,20 +354,26 @@ ToolChain::constructBatchJob(ArrayRef<const Job *> unsortedJobs,
                      *output, OI};
   auto invocationInfo = constructInvocation(*batchCJA, context);
   // Batch mode can produce quite long command lines; in almost every case these
-  // will trigger use of supplementary output file maps, but in some rare corner
-  // cases (very few files, very long paths) they might not. However, in those
-  // cases we _should_ degrade to using response files to pass arguments to the
-  // frontend, which is done automatically by code elsewhere.
+  // will trigger use of supplementary output file maps. However, if the driver
+  // command line is long for reasons unrelated to the number of input files,
+  // such as passing a large number of flags, then the individual batch jobs are
+  // also likely to overflow. We have to check for that explicitly here, because
+  // the BatchJob created here does not go through the same code path in
+  // constructJob above.
   //
   // The `allowsResponseFiles` flag on the `invocationInfo` we have here exists
   // only to model external tools that don't know about response files, such as
   // platform linkers; when talking to the frontend (which we control!) it
   // should always be true. But double check with an assert here in case someone
   // failed to set it in `constructInvocation`.
   assert(invocationInfo.allowsResponseFiles);
+  auto responseFileInfo =
+      getResponseFileInfo(C, executablePath, invocationInfo, context);
+
   return llvm::make_unique<BatchJob>(
       *batchCJA, inputJobs.takeVector(), std::move(output), executablePath,
       std::move(invocationInfo.Arguments),
       std::move(invocationInfo.ExtraEnvironment),
-      std::move(invocationInfo.FilelistInfos), sortedJobs, NextQuasiPID);
+      std::move(invocationInfo.FilelistInfos), sortedJobs, NextQuasiPID,
+      responseFileInfo);
 }

test/Driver/batch_mode_response_files.swift

@@ -0,0 +1,10 @@
+// Batch jobs go through a different code path than other jobs, so make sure
+// that they also use response files correctly when their argument lists are
+// too long.
+
+// RUN: %{python} -c 'for i in range(500001): print "-DTEST_" + str(i)' > %t.resp
+// RUN: %swiftc_driver -driver-print-jobs -module-name batch -enable-batch-mode -j 1 -c %S/Inputs/main.swift %S/Inputs/lib.swift @%t.resp 2>&1 > %t.jobs.txt
+// RUN: %FileCheck %s < %t.jobs.txt -check-prefix=BATCH
+
+// BATCH: bin{{/|\\\\}}swift{{c?}}
+// BATCH: @{{[^ ]*}}arguments-{{[0-9a-zA-Z]+}}.resp{{"?}} # -frontend -c -primary-file {{[^ ]+}}/Inputs/main.swift{{"?}} -primary-file {{[^ ]+}}/Inputs/lib.swift

test/Driver/force-response-files.swift

@@ -4,3 +4,7 @@
 // RUN: %swiftc_driver -driver-force-response-files -typecheck %S/../Inputs/empty.swift -### 2>&1 | %FileCheck %s
 // CHECK: @
 // CHECK: .resp
+
+// RUN: %swiftc_driver -enable-batch-mode -driver-force-response-files -typecheck %S/../Inputs/empty.swift -### 2>&1 | %FileCheck %s -check-prefix=BATCH
+// BATCH: @
+// BATCH: .resp