[Driver] <rdar://39504759> Continue demultiplexing subprocess output after initial read.

Author: graydon

include/swift/Basic/Statistics.def

@@ -41,6 +41,12 @@ DRIVER_STATISTIC(NumDriverJobsSkipped)
 /// EXIT_SUCCESS.
 DRIVER_STATISTIC(NumProcessFailures)
 
+/// Total number of driver poll() calls on subprocess pipes.
+DRIVER_STATISTIC(NumDriverPipePolls)
+
+/// Total number of driver read() calls on subprocess pipes.
+DRIVER_STATISTIC(NumDriverPipeReads)
+
 /// Next 10 statistics count dirtying-events in the driver's dependency graph,
 /// which it uses to decide which files are invalid (and thus which files to
 /// build). There are two dimensions to each dirtying event:

include/swift/Basic/TaskQueue.h

@@ -23,6 +23,7 @@
 #include <queue>
 
 namespace swift {
+class UnifiedStatsReporter;
 namespace sys {
 
 class Task; // forward declared to allow for platform-specific implementations
@@ -46,13 +47,18 @@ class TaskQueue {
   /// The number of tasks to execute in parallel.
   unsigned NumberOfParallelTasks;
 
+  /// Optional place to count I/O and subprocess events.
+  UnifiedStatsReporter *Stats;
+
 public:
   /// \brief Create a new TaskQueue instance.
   ///
   /// \param NumberOfParallelTasks indicates the number of tasks which should
   /// be run in parallel. If 0, the TaskQueue will choose the most appropriate
   /// number of parallel tasks for the current system.
-  TaskQueue(unsigned NumberOfParallelTasks = 0);
+  /// \param Optional stats reporter to count I/O and subprocess events.
+  TaskQueue(unsigned NumberOfParallelTasks = 0,
+            UnifiedStatsReporter *USR = nullptr);
   virtual ~TaskQueue();
 
   // TODO: remove once -Wdocumentation stops warning for \param, \returns on

lib/Basic/TaskQueue.cpp

@@ -29,8 +29,10 @@ using namespace swift::sys;
 #include "Default/TaskQueue.inc"
 #endif
 
-TaskQueue::TaskQueue(unsigned NumberOfParallelTasks)
-  : NumberOfParallelTasks(NumberOfParallelTasks) {}
+TaskQueue::TaskQueue(unsigned NumberOfParallelTasks,
+                     UnifiedStatsReporter *USR)
+  : NumberOfParallelTasks(NumberOfParallelTasks),
+    Stats(USR){}
 
 TaskQueue::~TaskQueue() = default;
 

lib/Basic/Unix/TaskQueue.inc

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/Basic/TaskQueue.h"
+#include "swift/Basic/Statistic.h"
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/DenseMap.h"
@@ -84,12 +85,16 @@ class Task {
   /// from the Task.
   std::string Errors;
 
+  /// Optional place to count I/O and subprocess events.
+  UnifiedStatsReporter *Stats;
+
 public:
   Task(const char *ExecPath, ArrayRef<const char *> Args,
-       ArrayRef<const char *> Env, void *Context, bool SeparateErrors)
+       ArrayRef<const char *> Env, void *Context, bool SeparateErrors,
+       UnifiedStatsReporter *USR)
       : ExecPath(ExecPath), Args(Args), Env(Env), Context(Context),
         SeparateErrors(SeparateErrors), Pid(-1), Pipe(-1), ErrorPipe(-1),
-        State(Preparing) {
+        State(Preparing), Stats(USR) {
     assert((Env.empty() || Env.back() == nullptr) &&
            "Env must either be empty or null-terminated!");
   }
@@ -108,8 +113,11 @@ public:
   bool execute();
 
   /// \brief Reads data from the pipes, if any is available.
-  /// \returns true on error, false on success
-  bool readFromPipes();
+  ///
+  /// If \p UntilEnd is true, reads until the end of the stream; otherwise reads
+  /// once (possibly with a retry on EINTR), and returns.
+  /// \returns true on error, false on success.
+  bool readFromPipes(bool UntilEnd = true);
 
   /// \brief Performs any post-execution work for this Task, such as reading
   /// piped output and closing the pipe.
@@ -241,7 +249,9 @@ bool Task::execute() {
   return false;
 }
 
-static bool readFromAPipe(int Pipe, std::string &Output) {
+static bool readFromAPipe(int Pipe, std::string &Output,
+                          UnifiedStatsReporter *Stats,
+                          bool UntilEnd) {
   char outputBuffer[1024];
   ssize_t readBytes = 0;
   while ((readBytes = read(Pipe, outputBuffer, sizeof(outputBuffer))) != 0) {
@@ -253,15 +263,19 @@ static bool readFromAPipe(int Pipe, std::string &Output) {
     }
 
     Output.append(outputBuffer, readBytes);
+    if (Stats)
+      Stats->getDriverCounters().NumDriverPipeReads++;
+    if (!UntilEnd)
+      break;
   }
 
   return false;
 }
 
-bool Task::readFromPipes() {
-  bool Ret = readFromAPipe(Pipe, Output);
+bool Task::readFromPipes(bool UntilEnd) {
+  bool Ret = readFromAPipe(Pipe, Output, Stats, UntilEnd);
   if (SeparateErrors) {
-    Ret |= readFromAPipe(ErrorPipe, Errors);
+    Ret |= readFromAPipe(ErrorPipe, Errors, Stats, UntilEnd);
   }
   return Ret;
 }
@@ -302,7 +316,7 @@ void TaskQueue::addTask(const char *ExecPath, ArrayRef<const char *> Args,
                         ArrayRef<const char *> Env, void *Context,
                         bool SeparateErrors) {
   std::unique_ptr<Task> T(
-      new Task(ExecPath, Args, Env, Context, SeparateErrors));
+      new Task(ExecPath, Args, Env, Context, SeparateErrors, Stats));
   QueuedTasks.push(std::move(T));
 }
 
@@ -355,6 +369,8 @@ bool TaskQueue::execute(TaskBeganCallback Began, TaskFinishedCallback Finished,
         continue;
       return true;
     }
+    if (Stats)
+      Stats->getDriverCounters().NumDriverPipePolls++;
 
     // Holds all fds which have finished during this loop iteration.
     std::vector<int> FinishedFds;
@@ -373,8 +389,15 @@ bool TaskQueue::execute(TaskBeganCallback Began, TaskFinishedCallback Finished,
                "All outstanding fds must be associated with an executing Task");
         Task &T = *iter->second;
         if (fd.revents & POLLIN || fd.revents & POLLPRI) {
-          // There's data available to read.
-          T.readFromPipes();
+          // There's data available to read. Read _some_ of it here, but not
+          // necessarily _all_, since the pipe is in blocking mode and we might
+          // have other input pending (or soon -- before this subprocess is done
+          // writing) from other subprocesses.
+          //
+          // FIXME: longer term, this should probably either be restructured to
+          // use O_NONBLOCK, or at very least poll the stderr file descriptor as
+          // well; the whole loop here is a bit of a mess.
+          T.readFromPipes(/*UntilEnd = */false);
         }
 
         if (fd.revents & POLLHUP || fd.revents & POLLERR) {

lib/Driver/Compilation.cpp

@@ -593,7 +593,8 @@ namespace driver {
       if (Comp.SkipTaskExecution)
         TQ.reset(new DummyTaskQueue(Comp.NumberOfParallelCommands));
       else
-        TQ.reset(new TaskQueue(Comp.NumberOfParallelCommands));
+        TQ.reset(new TaskQueue(Comp.NumberOfParallelCommands,
+                               Comp.Stats.get()));
       if (Comp.ShowIncrementalBuildDecisions || Comp.Stats)
         IncrementalTracer = &ActualIncrementalTracer;
     }

test/Driver/pipe_round_robin.swift.gyb

@@ -0,0 +1,33 @@
+// RUN: %empty-directory(%t/manyfuncs)
+// RUN: %empty-directory(%t/stats)
+//
+// This test is looking at behaviour of the driver's task queue, checking
+// to make sure that it drains the output streams of multiple subprocesses
+// evenly, rather than one subprocess all-at-once before the next
+// all-at-once. This isn't batch-mode specific but it emerges somewhat
+// vividly there when combined with lots of files that do
+// -debug-time-function-bodies.
+//
+// Here we do a non-batch-mode variant that has lots of functions in each
+// of 4 files.
+//
+// RUN: %gyb -D N=1 %s -o %t/manyfuncs/file1.swift
+// RUN: %gyb -D N=2 %s -o %t/manyfuncs/file2.swift
+// RUN: %gyb -D N=3 %s -o %t/manyfuncs/file3.swift
+// RUN: %gyb -D N=4 %s -o %t/manyfuncs/file4.swift
+//
+// We calculate the ratio of poll() calls to read() calls; these should be
+// nearly equal (we test abs(read/poll) < 2.0) if we're doing interleaved
+// reading. If we're doing non-interleaved reading, they become radically
+// different (eg. thousands of reads per poll).
+//
+// RUN: %target-build-swift -j 4 -module-name manyfuncs -typecheck -stats-output-dir %t/stats -Xfrontend -debug-time-function-bodies %t/manyfuncs/*.swift
+// RUN: %utils/process-stats-dir.py --evaluate 'abs(float(NumDriverPipeReads) / float(NumDriverPipePolls)) < 2.0' %t/stats
+
+% for i in range(1,1000):
+func process_${N}_function_${i}(_ x: Int) -> Int {
+  let v = (1 + 2 * 3 + x * 5 + x + 6)
+  let a = [v, 1, ${i}, 3, ${N}, 4]
+  return a.reduce(0, {$0 + $1})
+}
+% end