From 13269ba93b7453def7084b00eb4a34ad787a7c45 Mon Sep 17 00:00:00 2001
From: John Ericson <John.Ericson@Obsidian.Systems>
Date: Mon, 17 Apr 2023 12:14:15 -0400
Subject: [PATCH] Make `RemoteStore::ConnectionHandle` part of class and expose

Will need to do subclass-specific implementations in the next commit.
This isn't because there will be multiple variations of the daemon
protocol (whew!) but because different clients pick and choose different
parts to use.
---
 src/libstore/remote-store-connection.hh | 31 +++++++++++++++
 src/libstore/remote-store.cc            | 52 +++++++------------------
 src/libstore/remote-store.hh            |  4 +-
 3 files changed, 47 insertions(+), 40 deletions(-)

diff --git a/src/libstore/remote-store-connection.hh b/src/libstore/remote-store-connection.hh
index d32d91a60..ce4740a9c 100644
--- a/src/libstore/remote-store-connection.hh
+++ b/src/libstore/remote-store-connection.hh
@@ -1,5 +1,6 @@
 #include "remote-store.hh"
 #include "worker-protocol.hh"
+#include "pool.hh"
 
 namespace nix {
 
@@ -94,4 +95,34 @@ struct RemoteStore::Connection
     std::exception_ptr processStderr(Sink * sink = 0, Source * source = 0, bool flush = true);
 };
 
+/**
+ * A wrapper around Pool<RemoteStore::Connection>::Handle that marks
+ * the connection as bad (causing it to be closed) if a non-daemon
+ * exception is thrown before the handle is closed. Such an exception
+ * causes a deviation from the expected protocol and therefore a
+ * desynchronization between the client and daemon.
+ */
+struct RemoteStore::ConnectionHandle
+{
+    Pool<RemoteStore::Connection>::Handle handle;
+    bool daemonException = false;
+
+    ConnectionHandle(Pool<RemoteStore::Connection>::Handle && handle)
+        : handle(std::move(handle))
+    { }
+
+    ConnectionHandle(ConnectionHandle && h)
+        : handle(std::move(h.handle))
+    { }
+
+    ~ConnectionHandle();
+
+    RemoteStore::Connection & operator * () { return *handle; }
+    RemoteStore::Connection * operator -> () { return &*handle; }
+
+    void processStderr(Sink * sink = 0, Source * source = 0, bool flush = true);
+
+    void withFramedSink(std::function<void(Sink & sink)> fun);
+};
+
 }
diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index bfe2258a4..926ccd9d1 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -159,49 +159,25 @@ void RemoteStore::setOptions(Connection & conn)
 }
 
 
-/* A wrapper around Pool<RemoteStore::Connection>::Handle that marks
-   the connection as bad (causing it to be closed) if a non-daemon
-   exception is thrown before the handle is closed. Such an exception
-   causes a deviation from the expected protocol and therefore a
-   desynchronization between the client and daemon. */
-struct ConnectionHandle
+RemoteStore::ConnectionHandle::~ConnectionHandle()
 {
-    Pool<RemoteStore::Connection>::Handle handle;
-    bool daemonException = false;
-
-    ConnectionHandle(Pool<RemoteStore::Connection>::Handle && handle)
-        : handle(std::move(handle))
-    { }
-
-    ConnectionHandle(ConnectionHandle && h)
-        : handle(std::move(h.handle))
-    { }
-
-    ~ConnectionHandle()
-    {
-        if (!daemonException && std::uncaught_exceptions()) {
-            handle.markBad();
-            debug("closing daemon connection because of an exception");
-        }
+    if (!daemonException && std::uncaught_exceptions()) {
+        handle.markBad();
+        debug("closing daemon connection because of an exception");
     }
+}
 
-    RemoteStore::Connection * operator -> () { return &*handle; }
-    RemoteStore::Connection & operator * () { return *handle; }
-
-    void processStderr(Sink * sink = 0, Source * source = 0, bool flush = true)
-    {
-        auto ex = handle->processStderr(sink, source, flush);
-        if (ex) {
-            daemonException = true;
-            std::rethrow_exception(ex);
-        }
+void RemoteStore::ConnectionHandle::processStderr(Sink * sink, Source * source, bool flush)
+{
+    auto ex = handle->processStderr(sink, source, flush);
+    if (ex) {
+        daemonException = true;
+        std::rethrow_exception(ex);
     }
-
-    void withFramedSink(std::function<void(Sink & sink)> fun);
-};
+}
 
 
-ConnectionHandle RemoteStore::getConnection()
+RemoteStore::ConnectionHandle RemoteStore::getConnection()
 {
     return ConnectionHandle(connections->get());
 }
@@ -1099,7 +1075,7 @@ std::exception_ptr RemoteStore::Connection::processStderr(Sink * sink, Source *
     return nullptr;
 }
 
-void ConnectionHandle::withFramedSink(std::function<void(Sink & sink)> fun)
+void RemoteStore::ConnectionHandle::withFramedSink(std::function<void(Sink & sink)> fun)
 {
     (*this)->to.flush();
 
diff --git a/src/libstore/remote-store.hh b/src/libstore/remote-store.hh
index b12f5437f..68ec29e6c 100644
--- a/src/libstore/remote-store.hh
+++ b/src/libstore/remote-store.hh
@@ -17,7 +17,6 @@ class Pid;
 struct FdSink;
 struct FdSource;
 template<typename T> class Pool;
-struct ConnectionHandle;
 
 struct RemoteStoreConfig : virtual StoreConfig
 {
@@ -182,6 +181,8 @@ protected:
 
     void setOptions() override;
 
+    struct ConnectionHandle;
+
     ConnectionHandle getConnection();
 
     friend struct ConnectionHandle;
@@ -199,5 +200,4 @@ private:
         std::shared_ptr<Store> evalStore);
 };
 
-
 }