From 40f6b283be40f10893b2424b7f3397f98435709a Mon Sep 17 00:00:00 2001
From: Ian Roddis <gitlab@ie2r.com>
Date: Tue, 15 Jun 2021 14:43:47 -0300
Subject: [PATCH] Fixing things for programs with very large output.

---
 .../daggy/executors/ForkingExecutor.hpp       |  1 +
 daggy/src/executors/ForkingExecutor.cpp       | 36 ++++++++++---------
 tests/unit_executor_forkingexecutor.cpp       | 12 +++++++
 3 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/daggy/include/daggy/executors/ForkingExecutor.hpp b/daggy/include/daggy/executors/ForkingExecutor.hpp
index 9e14c67..de10d34 100644
--- a/daggy/include/daggy/executors/ForkingExecutor.hpp
+++ b/daggy/include/daggy/executors/ForkingExecutor.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <iostream>
 #include "../Executor.hpp"
 
 namespace daggy {
diff --git a/daggy/src/executors/ForkingExecutor.cpp b/daggy/src/executors/ForkingExecutor.cpp
index b814612..b86bef9 100644
--- a/daggy/src/executors/ForkingExecutor.cpp
+++ b/daggy/src/executors/ForkingExecutor.cpp
@@ -1,5 +1,9 @@
 #include <daggy/executors/ForkingExecutor.hpp>
 
+#include <array>
+#include <utility>
+
+#include <fcntl.h>
 #include <unistd.h>
 #include <wait.h>
 #include <poll.h>
@@ -9,27 +13,21 @@ using namespace daggy::executor;
 std::string slurp(int fd)  {
   std::string result;
 
-  const size_t BUFFER_SIZE = 4096;
+  const ssize_t BUFFER_SIZE = 4096;
   char buffer[BUFFER_SIZE];
 
   struct pollfd pfd{ .fd = fd, .events = POLLIN, .revents = 0 };
-  poll(&pfd, 1, 0);
+  poll(&pfd, 1, 1);
 
   while (pfd.revents & POLLIN) {
     ssize_t bytes = read(fd, buffer, BUFFER_SIZE);
-    if (bytes == -1) {
-      if (errno == EINTR) {
-        continue;
-      } else {
-        perror("read");
-        exit(1);
-      }
-    } else if (bytes == 0) {
+    if (bytes == 0) {
       break;
     } else {
-      result += buffer;
-      if (bytes < BUFFER_SIZE) break;
+      result.append(buffer, bytes);
     }
+    pfd.revents = 0;
+    poll(&pfd, 1, 1);
   }
 
   return result;
@@ -51,8 +49,8 @@ daggy::AttemptRecord
   argv.push_back(nullptr);
 
   // Create the pipe
-  int stdoutPipe[2]; pipe(stdoutPipe);
-  int stderrPipe[2]; pipe(stderrPipe);
+  int stdoutPipe[2]; pipe2(stdoutPipe, O_DIRECT);
+  int stderrPipe[2]; pipe2(stderrPipe, O_DIRECT);
 
   pid_t child = fork();
   if (child < 0) {
@@ -66,8 +64,13 @@ daggy::AttemptRecord
     exit(-1);
   }
 
+  std::atomic<bool> running = true;
+  std::thread stdoutReader([&]() { while(running) rec.output.append(slurp(stdoutPipe[0])); });
+  std::thread stderrReader([&]() { while(running) rec.error.append(slurp(stderrPipe[0])); });
+
   int rc = 0;
   waitpid(child, &rc, 0);
+  running = false;
 
   rec.stopTime = Clock::now();
   if (WIFEXITED(rc)) {
@@ -76,9 +79,8 @@ daggy::AttemptRecord
     rec.rc = -1;
   }
 
-
-  rec.output = slurp(stdoutPipe[0]);
-  rec.error = slurp(stderrPipe[0]);
+  stdoutReader.join();
+  stderrReader.join();
 
   close(stdoutPipe[0]);
   close(stderrPipe[0]);
diff --git a/tests/unit_executor_forkingexecutor.cpp b/tests/unit_executor_forkingexecutor.cpp
index 36d5fda..0f42989 100644
--- a/tests/unit_executor_forkingexecutor.cpp
+++ b/tests/unit_executor_forkingexecutor.cpp
@@ -1,4 +1,5 @@
 #include <iostream>
+#include <filesystem>
 
 #include "daggy/executors/ForkingExecutor.hpp"
 
@@ -26,4 +27,15 @@ TEST_CASE("Basic Execution", "[forking_executor]") {
     REQUIRE(rec.error == "/usr/bin/expr: syntax error: missing argument after ‘+’\n");
     REQUIRE(rec.output.empty());
   }
+
+  SECTION("Large Output") {
+    const std::string BIG_FILE{"/usr/share/dict/cracklib-small"};
+    std::vector<std::string> cmd{"/usr/bin/cat", BIG_FILE};
+
+    auto rec = ex.runCommand(cmd);
+
+    REQUIRE(rec.rc == 0);
+    REQUIRE(rec.output.size() == std::filesystem::file_size(BIG_FILE));
+    REQUIRE(rec.error.empty());
+  }
 }