Expose a bunch of things in the Legacy SSH Store for Hydra

src/libstore/legacy-ssh-store.cc

@@ -72,7 +72,7 @@ ref<LegacySSHStore::Connection> LegacySSHStore::openConnection()
         command.push_back("--store");
         command.push_back(remoteStore.get());
     }
-    conn->sshConn = master.startCommand(std::move(command));
+    conn->sshConn = master.startCommand(std::move(command), std::list{extraSshArgs});
     conn->to = FdSink(conn->sshConn->in.get());
     conn->from = FdSource(conn->sshConn->out.get());
 
@@ -103,29 +103,40 @@ std::string LegacySSHStore::getUri()
     return *uriSchemes().begin() + "://" + host;
 }
 
+std::map<StorePath, UnkeyedValidPathInfo> LegacySSHStore::queryPathInfosUncached(
+    const StorePathSet & paths)
+{
+    auto conn(connections->get());
+
+    /* No longer support missing NAR hash */
+    assert(GET_PROTOCOL_MINOR(conn->remoteVersion) >= 4);
+
+    debug("querying remote host '%s' for info on '%s'", host, concatStringsSep(", ", printStorePathSet(paths)));
+
+    auto infos = conn->queryPathInfos(*this, paths);
+
+    for (const auto & [_, info] : infos) {
+        if (info.narHash == Hash::dummy)
+            throw Error("NAR hash is now mandatory");
+    }
+
+    return infos;
+}
 
 void LegacySSHStore::queryPathInfoUncached(const StorePath & path,
     Callback<std::shared_ptr<const ValidPathInfo>> callback) noexcept
 {
     try {
-        auto conn(connections->get());
+        auto infos = queryPathInfosUncached({path});
-
-        /* No longer support missing NAR hash */
-        assert(GET_PROTOCOL_MINOR(conn->remoteVersion) >= 4);
 
         debug("querying remote host '%s' for info on '%s'", host, printStorePath(path));
 
-        auto infos = conn->queryPathInfos(*this, {path});
-
         switch (infos.size()) {
         case 0:
             return callback(nullptr);
         case 1: {
             auto & [path2, info] = *infos.begin();
 
-            if (info.narHash == Hash::dummy)
-                throw Error("NAR hash is now mandatory");
-
             assert(path == path2);
             return callback(std::make_shared<ValidPathInfo>(
                 std::move(path),
@@ -198,11 +209,19 @@ void LegacySSHStore::addToStore(const ValidPathInfo & info, Source & source,
 
 void LegacySSHStore::narFromPath(const StorePath & path, Sink & sink)
 {
-    auto conn(connections->get());
+    narFromPath(path, [&](auto & source) {
+        copyNAR(source, sink);
+    });
+}
 
+
+void LegacySSHStore::narFromPath(const StorePath & path, std::function<void(Source &)> fun)
+{
+    auto conn(connections->get());
     conn->to << ServeProto::Command::DumpStorePath << printStorePath(path);
     conn->to.flush();
-    copyNAR(conn->from, sink);
+
+    fun(conn->from);
 }
 
 
@@ -229,6 +248,25 @@ BuildResult LegacySSHStore::buildDerivation(const StorePath & drvPath, const Bas
     return ServeProto::Serialise<BuildResult>::read(*this, *conn);
 }
 
+std::function<BuildResult()> LegacySSHStore::buildDerivationAsync(
+    const StorePath & drvPath, const BasicDerivation & drv,
+    const ServeProto::BuildOptions & options)
+{
+    {
+        assert(maxConnections <= 1);
+        auto conn(connections->get());
+        conn->putBuildDerivationRequest(*this, drvPath, drv, options);
+    }
+
+    return [this]() -> BuildResult {
+        // TODO close over connection instead.
+        assert(maxConnections <= 1);
+        auto conn(connections->get());
+
+        return ServeProto::Serialise<BuildResult>::read(*this, *conn);
+    };
+}
+
 
 void LegacySSHStore::buildPaths(const std::vector<DerivedPath> & drvPaths, BuildMode buildMode, std::shared_ptr<Store> evalStore)
 {
@@ -300,6 +338,32 @@ StorePathSet LegacySSHStore::queryValidPaths(const StorePathSet & paths,
 }
 
 
+StorePathSet LegacySSHStore::queryValidPaths(const StorePathSet & paths,
+    bool lock, SubstituteFlag maybeSubstitute)
+{
+    auto conn(connections->get());
+    return conn->queryValidPaths(*this,
+        lock, paths, maybeSubstitute);
+}
+
+
+void LegacySSHStore::addMultipleToStoreLegacy(Store & srcStore, const StorePathSet & paths)
+{
+    auto conn(connections->get());
+    conn->to << ServeProto::Command::ImportPaths;
+    try {
+        srcStore.exportPaths(paths, conn->to);
+    } catch (...) {
+        conn->good = false;
+        throw;
+    }
+    conn->to.flush();
+
+    if (readInt(conn->from) != 1)
+        throw Error("remote machine failed to import closure");
+}
+
+
 void LegacySSHStore::connect()
 {
     auto conn(connections->get());
@@ -313,6 +377,23 @@ unsigned int LegacySSHStore::getProtocol()
 }
 
 
+pid_t LegacySSHStore::getConnectionPid()
+{
+    auto conn(connections->get());
+    return conn->sshConn->sshPid;
+}
+
+
+LegacySSHStore::ConnectionStats LegacySSHStore::getConnectionStats()
+{
+    auto conn(connections->get());
+    return {
+        .bytesReceived = conn->from.read,
+        .bytesSent = conn->to.written,
+    };
+}
+
+
 /**
  * The legacy ssh protocol doesn't support checking for trusted-user.
  * Try using ssh-ng:// instead if you want to know.

src/libstore/legacy-ssh-store.hh

@@ -6,6 +6,7 @@
 #include "ssh.hh"
 #include "callback.hh"
 #include "pool.hh"
+#include "serve-protocol-impl.hh"
 
 namespace nix {
 
@@ -19,6 +20,11 @@ struct LegacySSHStoreConfig : virtual CommonSSHStoreConfig
     const Setting<int> maxConnections{this, 1, "max-connections",
         "Maximum number of concurrent SSH connections."};
 
+    /**
+     * Hack for hydra
+     */
+    Strings extraSshArgs = {};
+
     const std::string name() override { return "SSH Store"; }
 
     std::string doc() override;
@@ -60,11 +66,21 @@ public:
     void queryPathInfoUncached(const StorePath & path,
         Callback<std::shared_ptr<const ValidPathInfo>> callback) noexcept override;
 
+    std::map<StorePath, UnkeyedValidPathInfo> queryPathInfosUncached(
+        const StorePathSet & paths);
+
     void addToStore(const ValidPathInfo & info, Source & source,
         RepairFlag repair, CheckSigsFlag checkSigs) override;
 
     void narFromPath(const StorePath & path, Sink & sink) override;
 
+    /**
+     * Gives Hands over the connection temporarily as source.
+     *
+     * Caller must be sure to not consume more than the NAR.
+     */
+    void narFromPath(const StorePath & path, std::function<void(Source &)> fun);
+
     std::optional<StorePath> queryPathFromHashPart(const std::string & hashPart) override
     { unsupported("queryPathFromHashPart"); }
 
@@ -93,6 +109,15 @@ public:
     BuildResult buildDerivation(const StorePath & drvPath, const BasicDerivation & drv,
         BuildMode buildMode) override;
 
+    /**
+     * Unsafe if connection in pool is greater than 1!
+     *
+     * Can make safe once we have C++23 `std::move_only_function`.
+     */
+    std::function<BuildResult()> buildDerivationAsync(
+        const StorePath & drvPath, const BasicDerivation & drv,
+        const ServeProto::BuildOptions & options);
+
     void buildPaths(const std::vector<DerivedPath> & drvPaths, BuildMode buildMode, std::shared_ptr<Store> evalStore) override;
 
     void ensurePath(const StorePath & path) override
@@ -119,10 +144,36 @@ public:
     StorePathSet queryValidPaths(const StorePathSet & paths,
         SubstituteFlag maybeSubstitute = NoSubstitute) override;
 
+    /**
+     * Custom variation that atomically creates temp locks on the remote
+     * side.
+     *
+     * This exists to prevent a race where the remote host
+     * garbage-collects paths that are already there. Optionally, ask
+     * the remote host to substitute missing paths.
+     */
+    StorePathSet queryValidPaths(const StorePathSet & paths,
+        bool lock,
+        SubstituteFlag maybeSubstitute = NoSubstitute);
+
+    /**
+     * Just exists because this is exactly what Hydra was doing, and we
+     * don't yet want an algorithmic change.
+     */
+    void addMultipleToStoreLegacy(Store & srcStore, const StorePathSet & paths);
+
     void connect() override;
 
     unsigned int getProtocol() override;
 
+    struct ConnectionStats {
+        size_t bytesReceived, bytesSent;
+    };
+
+    ConnectionStats getConnectionStats();
+
+    pid_t getConnectionPid();
+
     /**
      * The legacy ssh protocol doesn't support checking for trusted-user.
      * Try using ssh-ng:// instead if you want to know.