Merge pull request #10748 from NixOS/legacy-ssh-extensions-for-hydra

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

Author: Ericson2314

src/libstore/legacy-ssh-store.cc

@@ -69,7 +69,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());
 
@@ -100,29 +100,38 @@ 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());
-
-        /* 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});
+        auto infos = queryPathInfosUncached({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),
@@ -193,13 +202,19 @@ void LegacySSHStore::addToStore(const ValidPathInfo & info, Source & source,
 
 void LegacySSHStore::narFromPath(const StorePath & path, Sink & sink)
 {
-    auto conn(connections->get());
-    conn->narFromPath(*this, path, [&](auto & source) {
+    narFromPath(path, [&](auto & source) {
         copyNAR(source, sink);
     });
 }
 
 
+void LegacySSHStore::narFromPath(const StorePath & path, std::function<void(Source &)> fun)
+{
+    auto conn(connections->get());
+    conn->narFromPath(*this, path, fun);
+}
+
+
 static ServeProto::BuildOptions buildSettings()
 {
     return {
@@ -223,6 +238,19 @@ BuildResult LegacySSHStore::buildDerivation(const StorePath & drvPath, const Bas
     return conn->getBuildDerivationResponse(*this);
 }
 
+std::function<BuildResult()> LegacySSHStore::buildDerivationAsync(
+    const StorePath & drvPath, const BasicDerivation & drv,
+    const ServeProto::BuildOptions & options)
+{
+    // Until we have C++23 std::move_only_function
+    auto conn = std::make_shared<Pool<Connection>::Handle>(connections->get());
+    (*conn)->putBuildDerivationRequest(*this, drvPath, drv, options);
+
+    return [this,conn]() -> BuildResult {
+        return (*conn)->getBuildDerivationResponse(*this);
+    };
+}
+
 
 void LegacySSHStore::buildPaths(const std::vector<DerivedPath> & drvPaths, BuildMode buildMode, std::shared_ptr<Store> evalStore)
 {
@@ -294,6 +322,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());
@@ -307,6 +361,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.hh"
 
 namespace nix {
 
@@ -24,6 +25,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"; }
 
     static std::set<std::string> uriSchemes() { return {"ssh"}; }
@@ -60,11 +66,24 @@ struct LegacySSHStore : public virtual LegacySSHStoreConfig, public virtual Stor
     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;
 
+    /**
+     * Hands over the connection temporarily as source to the given
+     * function. The function must not consume beyond the NAR; it can
+     * not just blindly try to always read more bytes until it is
+     * cut-off.
+     *
+     * This is exposed for sake of Hydra.
+     */
+    void narFromPath(const StorePath & path, std::function<void(Source &)> fun);
+
     std::optional<StorePath> queryPathFromHashPart(const std::string & hashPart) override
     { unsupported("queryPathFromHashPart"); }
 
@@ -93,6 +112,16 @@ public:
     BuildResult buildDerivation(const StorePath & drvPath, const BasicDerivation & drv,
         BuildMode buildMode) override;
 
+    /**
+     * Note, the returned function must only be called once, or we'll
+     * try to read from the connection twice.
+     *
+     * @todo Use 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 +148,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.