libstore: keep reading the daemon's messages while blocked sending to it

src/libstore/remote-store.cc

@@ -944,12 +944,48 @@ std::shared_ptr<SourceAccessor> RemoteStore::getFSAccessor(const StorePath & pat
     return getRemoteFSAccessor(requireValidPath)->accessObject(path);
 }
 
+/**
+ * Writes to the daemon connection, reading the daemon's stderr
+ * messages whenever a write would block.
+ *
+ * The daemon (or a thread of it, such as an auto-GC logging every
+ * path it deletes) can block on sending us a message while we block
+ * on sending it data. Once both socket buffers are full, reading only
+ * between frames, as `FramedSink` does, is a deadlock.
+ *
+ * A daemon error can't be consumed here mid-frame, since the daemon
+ * drains our frames before reporting it, so an exception from reading
+ * only ever means the connection is broken.
+ */
+struct DrainingSink : BufferedSink
+{
+    FdSink & to;
+    FdSource & from;
+    fun<void()> drain;
+
+    DrainingSink(FdSink & to, FdSource & from, fun<void()> drain)
+        : to(to)
+        , from(from)
+        , drain(std::move(drain))
+    {
+    }
+
+    void writeUnbuffered(std::string_view data) override
+    {
+        to.written += data.size();
+        writeFullWhileDraining(to.fd, data, from.fd, drain);
+    }
+};
+
 void RemoteStore::ConnectionHandle::withFramedSink(fun<void(Sink & sink)> sendData)
 {
     (*this)->to.flush();
 
+    /* Outlives the FramedSink, whose destructor sends the terminator through it. */
+    DrainingSink to((*this)->to, (*this)->from, [&]() { processStderr(nullptr, nullptr, false, false); });
+
     {
-        FramedSink sink((*this)->to, [&]() {
+        FramedSink sink(to, [&]() {
             /* Periodically process stderr messages and exceptions
                from the daemon. */
             processStderr(nullptr, nullptr, false, false);

src/libutil-tests/file-descriptor.cc

@@ -6,10 +6,12 @@
 #include "nix/util/signals.hh"
 
 #include <cstring>
+#include <thread>
 
 #ifndef _WIN32
 #  include <fcntl.h>
 #  include <stdlib.h>
+#  include <sys/socket.h>
 #endif
 
 namespace nix {
@@ -322,4 +324,61 @@ TEST(WriteFull, RespectsAllowInterrupts)
     EXPECT_EQ(source.readLine(/*eofOk=*/true), "hello");
 }
 
+#ifndef _WIN32
+TEST(WriteFullWhileDraining, DrainsPeerWhileBlocked)
+{
+    int fds[2];
+    ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0);
+    AutoCloseFD ours(fds[0]), theirs(fds[1]);
+
+    // Far more than fits in a socket buffer in either direction.
+    const size_t size = 8 * 1024 * 1024;
+    std::string toPeer(size, 'a'), fromPeer(size, 'b');
+
+    // Like a daemon stuck logging to us: the peer sends everything
+    // before it reads anything, so a plain writeFull() from our side
+    // would deadlock with it.
+    std::string peerGot;
+    std::thread peer([&]() {
+        writeFull(theirs.get(), fromPeer, /*allowInterrupts=*/false);
+        peerGot = drainFD(theirs.get());
+        theirs.close();
+    });
+
+    std::string drained;
+    writeFullWhileDraining(ours.get(), toPeer, ours.get(), [&]() {
+        std::array<char, 64 * 1024> buf;
+        auto n = nix::read(ours.get(), {reinterpret_cast<std::byte *>(buf.data()), buf.size()});
+        drained.append(buf.data(), n);
+    });
+
+    // The descriptor must be blocking again afterwards.
+    EXPECT_EQ(fcntl(ours.get(), F_GETFL) & O_NONBLOCK, 0);
+
+    shutdown(ours.get(), SHUT_WR);
+    drained += drainFD(ours.get());
+    peer.join();
+
+    EXPECT_EQ(peerGot.size(), size);
+    EXPECT_TRUE(peerGot == toPeer);
+    EXPECT_EQ(drained.size(), size);
+    EXPECT_TRUE(drained == fromPeer);
+}
+
+TEST(WriteFullWhileDraining, PropagatesDrainException)
+{
+    int fds[2];
+    ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0);
+    AutoCloseFD ours(fds[0]), theirs(fds[1]);
+
+    // Something to drain, and nobody reading our end, so the write blocks.
+    writeFull(theirs.get(), "x", /*allowInterrupts=*/false);
+
+    std::string big(8 * 1024 * 1024, 'a');
+    EXPECT_THROW(writeFullWhileDraining(ours.get(), big, ours.get(), []() { throw Error("drain failed"); }), Error);
+
+    EXPECT_EQ(fcntl(ours.get(), F_GETFL) & O_NONBLOCK, 0);
+}
+#endif
+
 } // namespace nix

src/libutil/file-descriptor.cc

@@ -1,6 +1,7 @@
 #include "nix/util/serialise.hh"
 #include "nix/util/util.hh"
 #include "nix/util/signals.hh"
+#include "nix/util/finally.hh"
 
 #include <span>
 #include <fcntl.h>
@@ -114,6 +115,60 @@ void writeFull(Descriptor fd, std::string_view s, bool allowInterrupts)
     }
 }
 
+void writeFullWhileDraining(Descriptor fd, std::string_view s, Descriptor drainFd, const fun<void()> & drain)
+{
+#ifdef _WIN32
+    writeFull(fd, s);
+#else
+    /* O_NONBLOCK lives on the open file description, which `drain`
+       may read from expecting to block, so only set it while we're
+       writing, and leave it alone if it was already set. */
+    auto flags = fcntl(fd, F_GETFL);
+    if (flags == -1)
+        throw SysError("getting file descriptor flags");
+    bool toggle = !(flags & O_NONBLOCK);
+    auto setNonBlocking = [&](bool nonBlocking) {
+        if (toggle && fcntl(fd, F_SETFL, nonBlocking ? flags | O_NONBLOCK : flags) == -1)
+            throw SysError("setting file descriptor flags");
+    };
+    setNonBlocking(true);
+    Finally restore([&]() { setNonBlocking(false); });
+
+    while (!s.empty()) {
+        checkInterrupt();
+
+        try {
+            s.remove_prefix(write(fd, {reinterpret_cast<const std::byte *>(s.data()), s.size()}, true));
+            continue;
+        } catch (SystemError & e) {
+            if (!e.is(std::errc::resource_unavailable_try_again) && !e.is(std::errc::operation_would_block))
+                throw;
+        }
+
+        /* The write would block. Wait until it can make progress, but
+           also wait for input from the peer, which may be blocked on
+           sending to us and so never get around to reading from us. */
+        struct pollfd pfds[2];
+        pfds[0].fd = fd;
+        pfds[0].events = POLLOUT;
+        pfds[0].revents = 0;
+        pfds[1].fd = drainFd;
+        pfds[1].events = POLLIN;
+        pfds[1].revents = 0;
+        if (poll(pfds, 2, -1) == -1) {
+            if (errno == EINTR)
+                continue;
+            throw SysError("poll on file descriptor failed");
+        }
+        if (pfds[1].revents & (POLLIN | POLLHUP | POLLERR)) {
+            setNonBlocking(false);
+            drain();
+            setNonBlocking(true);
+        }
+    }
+#endif
+}
+
 void writeLine(Descriptor fd, std::string s)
 {
     s += '\n';

src/libutil/include/nix/util/file-descriptor.hh

@@ -11,6 +11,7 @@
 
 #include "nix/util/canon-path.hh"
 #include "nix/util/error.hh"
+#include "nix/util/fun.hh"
 #include "nix/util/os-string.hh"
 
 #ifdef _WIN32
@@ -129,6 +130,27 @@ void readFull(Descriptor fd, char * buf, size_t count);
 
 void writeFull(Descriptor fd, std::string_view s, bool allowInterrupts = true);
 
+/**
+ * Write all of `s` to `fd`, like `writeFull`, but keep servicing
+ * `drainFd` while the write would block: whenever `fd` is not
+ * writable and `drainFd` becomes readable (or hangs up), `drain` is
+ * called before the write is retried.
+ *
+ * Use this when the peer behind `fd` may itself block on sending us
+ * data that we haven't read yet (typically `drainFd` is the other
+ * direction of the same connection), since a plain blocking write
+ * would then deadlock with it.
+ *
+ * `fd` is put in non-blocking mode only while writing, and `drain`
+ * is called with it blocking again, so `drain` can rely on blocking
+ * reads even when `drainFd` is the same open file description. An
+ * exception from `drain` propagates immediately, leaving an
+ * unspecified prefix of `s` written.
+ *
+ * On Windows this is just `writeFull`.
+ */
+void writeFullWhileDraining(Descriptor fd, std::string_view s, Descriptor drainFd, const fun<void()> & drain);
+
 /**
  * Read a line from an unbuffered file descriptor.
  * See BufferedSource::readLine for a buffered variant.

tests/functional/meson.build

@@ -62,6 +62,7 @@ suites = [
       'experimental-features.sh',
       'fetchMercurial.sh',
       'gc-auto.sh',
+      'nix-copy-auto-gc.sh',
       'user-envs.sh',
       'user-envs-migration.sh',
       'cli-characterisation.sh',

tests/functional/nix-copy-auto-gc.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+
+source common.sh
+
+# A `nix copy` into the daemon must not deadlock with an auto-GC that the
+# daemon starts while it is still receiving the paths. The GC logs every path
+# it deletes to the client over the same socket the client is sending on, so
+# once both socket buffers were full, each side used to wait on the other
+# forever.
+
+TODO_NixOS
+
+clearStore
+
+# Far more garbage than fits in a socket buffer once the daemon has logged
+# "deleting '<path>'" for each one.
+longName=$(printf 'g%.0s' {1..150})
+nix eval --expr "
+  builtins.deepSeq (builtins.genList (i: builtins.toFile \"$longName-\${toString i}\" (toString i)) 2000) true
+" > /dev/null
+[[ $(find "$NIX_STORE_DIR" -maxdepth 1 -name "*-$longName-*" | wc -l) -eq 2000 ]]
+
+# A path far bigger than a socket buffer, in a binary cache but not in the
+# store, so that it is still being streamed into the daemon when the GC starts.
+head -c 20000000 /dev/zero > "$TEST_ROOT/big"
+bigPath=$(nix store add-path --name big "$TEST_ROOT/big")
+nix copy --to "file://$TEST_ROOT/cache" "$bigPath"
+nix-store --delete "$bigPath"
+
+# Make the daemon think it is nearly out of space.
+export _NIX_TEST_FREE_SPACE_FILE=$TEST_ROOT/fake-free
+echo 100 > "$_NIX_TEST_FREE_SPACE_FILE"
+cat >> "$test_nix_conf" <<EOF
+min-free = 1M
+max-free = 1G
+min-free-check-interval = 1
+EOF
+
+startDaemon
+
+# A client stuck in write() ignores SIGTERM, so follow up with SIGKILL.
+timeout -k 5 120 nix copy --no-check-sigs --from "file://$TEST_ROOT/cache" "$bigPath"
+
+nix path-info "$bigPath"
+[[ $(find "$NIX_STORE_DIR" -maxdepth 1 -name "*-$longName-*" | wc -l) -eq 0 ]]