libstore: handle log size limiting in only one place

extracting the size-specific log handling into a new sink deduplicates
the size-limiting code and lets us clean up the derivation goal a bit:
a number of log-related fields we kept only because previous ownership
semantics were a lot more ad-hoc can be moved into the new sink class.

Change-Id: Ie770a276597c84f9d928a9624c0f0de86adfe815

lix/libstore/build/derivation-goal.cc

@@ -29,7 +29,10 @@
 #include <kj/async-unix.h>
 #include <kj/async.h>
 #include <kj/debug.h>
+#include <kj/exception.h>
 #include <kj/vector.h>
+#include <limits>
+#include <memory>
 #include <optional>
 #include <ranges>
 #include <sys/types.h>
@@ -1141,6 +1144,37 @@ kj::Promise<Result<SingleDrvOutputs>> DerivationGoal::registerOutputs()
     return assertPathValidity();
 }
 
+DerivationGoal::LogSink::LogSink(AutoCloseFD fd, ref<BufferedSink> file, bool compress, uint64_t limit)
+    : fd(std::move(fd))
+    , file(file)
+    , target(compress ? makeCompressionSink("bzip2", *file) : file)
+    , limit(limit)
+{
+}
+
+DerivationGoal::LogSink::~LogSink() noexcept
+{
+    signal.fulfiller->fulfill(false);
+}
+
+void DerivationGoal::LogSink::operator()(std::string_view data)
+{
+    writtenSoFar += data.size();
+    if (writtenSoFar <= limit) {
+        (*target)(data);
+    } else {
+        signal.fulfiller->fulfill(true);
+    }
+}
+
+void DerivationGoal::LogSink::finish()
+{
+    if (auto inner2 = dynamic_cast<FinishSink *>(&*target)) {
+        inner2->finish();
+    }
+    file->flush();
+}
+
 Path DerivationGoal::openLogFile()
 {
     if (!settings.keepLog) return "";
@@ -1159,26 +1193,26 @@ Path DerivationGoal::openLogFile()
     Path logFileName = fmt("%s/%s%s", dir, baseName.substr(2),
         settings.compressLog ? ".bz2" : "");
 
-    fdLogFile = sys::open(logFileName, O_CREAT | O_WRONLY | O_TRUNC | O_CLOEXEC, 0666);
+    auto fdLogFile = sys::open(logFileName, O_CREAT | O_WRONLY | O_TRUNC | O_CLOEXEC, 0666);
     if (!fdLogFile) throw SysError("creating log file '%1%'", logFileName);
 
-    logFileSink = std::make_shared<FdSink>(fdLogFile.get());
+    auto logFileSink = make_ref<FdSink>(fdLogFile.get());
+    const auto logLimit =
+        settings.maxLogSize ? settings.maxLogSize.get() : std::numeric_limits<uint64_t>::max();
 
-    if (settings.compressLog)
-        logSink = std::shared_ptr<CompressionSink>(makeCompressionSink("bzip2", *logFileSink));
-    else
-        logSink = logFileSink;
+    logSink = std::make_unique<LogSink>(
+        std::move(fdLogFile), logFileSink, settings.compressLog, logLimit
+    );
 
     return logFileName;
 }
 
 void DerivationGoal::closeLogFile()
 {
-    auto logSink2 = std::dynamic_pointer_cast<CompressionSink>(logSink);
-    if (logSink2) logSink2->finish();
-    if (logFileSink) logFileSink->flush();
-    logSink = logFileSink = 0;
-    fdLogFile.reset();
+    if (logSink) {
+        logSink->finish();
+    }
+    logSink = 0;
 }
 
 
@@ -1195,7 +1229,6 @@ kj::Promise<Result<std::optional<Goal::WorkResult>>> DerivationGoal::handleRawCh
 try {
     assert(hook);
 
-    uint64_t logSize = 0;
     std::string currentHookLine;
 
     AsyncFdIoStream in(AsyncFdIoStream::shared_fd{}, hook->fromHook.get());
@@ -1225,10 +1258,6 @@ try {
                         if (type == resBuildLogLine) {
                             const std::string logLine =
                                 (fields.size() > 0 ? fields[0].get<std::string>() : "") + "\n";
-                            logSize += logLine.size();
-                            if (settings.maxLogSize && logSize > settings.maxLogSize) {
-                                co_return tooMuchLogs();
-                            }
                             (*logSink)(logLine);
                         } else if (type == resSetPhase && ! fields.is_null()) {
                             const auto phase = fields[0];
@@ -1275,6 +1304,22 @@ DerivationGoal::wrapChildHandler(kj::Promise<Result<std::optional<WorkResult>>>
         );
     }
 
+    if (logSink) {
+        handler = handler.exclusiveJoin(logSink->signal.promise.then(
+            [&](bool limitReached) -> kj::Promise<Result<std::optional<WorkResult>>> {
+                try {
+                    if (limitReached) {
+                        return {tooMuchLogs()};
+                    } else {
+                        return kj::NEVER_DONE;
+                    }
+                } catch (...) {
+                    return {result::current_exception()};
+                }
+            }
+        ));
+    }
+
     return handler;
 }
 

lix/libstore/build/derivation-goal.hh

@@ -74,6 +74,25 @@ struct InitialOutput {
  */
 struct DerivationGoal : public Goal
 {
+    /**
+     * A sink that only passes a limited amount of data and signals exhaustion through a promise.
+     */
+    struct LogSink : FinishSink
+    {
+        AutoCloseFD fd;
+        ref<BufferedSink> file, target;
+        const uint64_t limit;
+        uint64_t writtenSoFar = 0;
+
+        kj::PromiseFulfillerPair<bool> signal = kj::newPromiseAndFulfiller<bool>();
+
+        LogSink(AutoCloseFD fd, ref<BufferedSink> file, bool compress, uint64_t limit);
+        ~LogSink() noexcept;
+
+        void operator()(std::string_view data) override;
+        void finish() override;
+    };
+
     /**
       * Whether this goal has completed. Completed goals can not be
       * asked for more outputs, a new goal must be created instead.
@@ -186,10 +205,9 @@ struct DerivationGoal : public Goal
     BuildResult buildResult;
 
     /**
-     * File descriptor for the log file.
+     * Sink for the log file.
      */
-    AutoCloseFD fdLogFile;
-    std::shared_ptr<BufferedSink> logFileSink, logSink;
+    std::shared_ptr<LogSink> logSink;
 
     /**
      * The most recent log lines.

lix/libstore/build/local-derivation-goal.cc

@@ -2695,7 +2695,6 @@ try {
 
     std::map<ActivityId, Activity> builderActivities;
     LogLineSplitter splitter;
-    uint64_t logSize = 0;
 
     auto act = logger->startActivity(
         lvlInfo,
@@ -2748,11 +2747,6 @@ try {
             co_return std::nullopt;
         }
 
-        logSize += data.size();
-        if (settings.maxLogSize && logSize > settings.maxLogSize) {
-            co_return tooMuchLogs();
-        }
-
         if (logSink) {
             (*logSink)(data);
         }