Make computeFSClosure() much faster over SSH

src/libstore/daemon.cc

@@ -21,6 +21,7 @@
 #include "nix/store/globals.hh"
 #include "nix/store/active-builds.hh"
 #include "nix/util/provenance.hh"
+#include "nix/util/async.hh"
 
 #ifndef _WIN32 // TODO need graceful async exit support on Windows?
 #  include "nix/util/monitor-fd.hh"
@@ -886,6 +887,37 @@ static void performOp(
         break;
     }
 
+    case WorkerProto::Op::QueryPathInfos: {
+        auto paths = WorkerProto::Serialise<StorePathSet>::read(*store, rconn);
+        logger->startWork();
+        std::vector<ValidPathInfo> infos;
+        {
+            asio::io_context ctx;
+            std::exception_ptr ex;
+            asio::co_spawn(
+                ctx,
+                [&]() -> asio::awaitable<void> {
+                    co_await store->queryPathInfos(
+                        paths, [&](std::vector<std::pair<StorePath, std::shared_ptr<const ValidPathInfo>>> results) {
+                            for (auto & [path, info] : results)
+                                if (info)
+                                    infos.push_back(*info);
+                        });
+                },
+                [&](std::exception_ptr e) { ex = e; });
+            ctx.run();
+            if (ex)
+                std::rethrow_exception(ex);
+        }
+        logger->stopWork();
+        /* Write the infos for the valid paths. Paths not reported are
+           invalid. */
+        conn.to << infos.size();
+        for (auto & info : infos)
+            WorkerProto::write(*store, wconn, info);
+        break;
+    }
+
     case WorkerProto::Op::OptimiseStore:
         logger->startWork();
         store->optimiseStore();

src/libstore/include/nix/store/remote-store.hh

@@ -62,6 +62,10 @@ struct RemoteStore : public virtual Store,
     void queryPathInfoUncached(
         const StorePath & path, Callback<std::shared_ptr<const ValidPathInfo>> callback) noexcept override;
 
+    asio::awaitable<void> queryPathInfos(
+        const std::set<StorePath> & paths,
+        fun<void(std::vector<std::pair<StorePath, std::shared_ptr<const ValidPathInfo>>>)> callback) override;
+
     void queryReferrers(const StorePath & path, StorePathSet & referrers) override;
 
     StorePathSet queryValidDerivers(const StorePath & path) override;

src/libstore/include/nix/store/store-api.hh

@@ -14,6 +14,8 @@
 #include "nix/store/store-dir-config.hh"
 #include "nix/store/store-reference.hh"
 #include "nix/util/source-path.hh"
+#include "nix/util/async.hh"
+#include "nix/util/fun.hh"
 
 #include <nlohmann/json_fwd.hpp>
 #include <atomic>
@@ -430,6 +432,19 @@ public:
      */
     void queryPathInfo(const StorePath & path, Callback<ref<const ValidPathInfo>> callback) noexcept;
 
+    /**
+     * Asynchronously query information about multiple store paths. As
+     * results arrive (possibly in batches from a remote server),
+     * `callback` is invoked one or more times with a vector of
+     * `(path, info)` pairs. A null `info` denotes that the path is
+     * not valid. Every requested path is reported exactly once across
+     * all invocations of `callback`. Unlike `queryPathInfo()`, an
+     * invalid path is not an error.
+     */
+    virtual asio::awaitable<void> queryPathInfos(
+        const std::set<StorePath> & paths,
+        fun<void(std::vector<std::pair<StorePath, std::shared_ptr<const ValidPathInfo>>>)> callback);
+
     /**
      * Version of queryPathInfo() that only queries the local narinfo cache and not
      * the actual store.

src/libstore/include/nix/store/worker-protocol.hh

@@ -114,6 +114,7 @@ struct WorkerProto
     static constexpr std::string_view featureProvenance = "provenance";
     static constexpr std::string_view featureVersionedAddToStoreMultiple = "versionedAddToStoreMultiple";
     static constexpr std::string_view featureAddTempRoots = "addTempRoots";
+    static constexpr std::string_view featureQueryPathInfos = "queryPathInfos";
 
     /**
      * A unidirectional read connection, to be used by the read half of the
@@ -240,6 +241,7 @@ enum struct WorkerProto::Op : uint64_t {
     AddPermRoot = 47,
     QueryActiveBuilds = 48,
     AddTempRoots = 49,
+    QueryPathInfos = 50,
 };
 
 struct WorkerProto::ClientHandshakeInfo

src/libstore/misc.cc

@@ -19,15 +19,11 @@
 namespace nix {
 
 void Store::computeFSClosure(
-    const StorePathSet & startPaths,
-    StorePathSet & paths_,
-    bool flipDirection,
-    bool includeOutputs,
-    bool includeDerivers)
+    const StorePathSet & startPaths, StorePathSet & out, bool flipDirection, bool includeOutputs, bool includeDerivers)
 {
-    std::function<asio::awaitable<StorePathSet>(const StorePath & path)> queryDeps;
-    if (flipDirection)
-        queryDeps = [this, includeOutputs, includeDerivers](const StorePath & path) -> asio::awaitable<StorePathSet> {
+    if (flipDirection) {
+        std::function<asio::awaitable<StorePathSet>(const StorePath & path)> queryDeps =
+            [this, includeOutputs, includeDerivers](const StorePath & path) -> asio::awaitable<StorePathSet> {
             StorePathSet res;
             StorePathSet referrers;
             queryReferrers(path, referrers);
@@ -45,27 +41,62 @@ void Store::computeFSClosure(
                         res.insert(*maybeOutPath);
             co_return res;
         };
-    else
-        queryDeps = [this, includeOutputs, includeDerivers](const StorePath & path) -> asio::awaitable<StorePathSet> {
-            StorePathSet res;
-            auto info = co_await callbackToAwaitable<ref<const ValidPathInfo>>(
-                [this, path](Callback<ref<const ValidPathInfo>> cb) { queryPathInfo(path, std::move(cb)); });
+        computeClosure<StorePath>(startPaths, out, GetEdgesAsync<StorePath>(queryDeps));
+    } else {
+
+        asio::io_context ctx;
+        std::exception_ptr ex;
+
+        asio::co_spawn(
+            ctx,
+            [&]() -> asio::awaitable<void> {
+                auto todo = startPaths;
+
+                StorePathSet required, done;
+
+                while (!todo.empty()) {
+                    StorePathSet batch;
+                    for (auto & path : std::exchange(todo, {}))
+                        if (done.insert(path).second)
+                            batch.insert(path);
+
+                    co_await queryPathInfos(
+                        batch, [&](std::vector<std::pair<StorePath, std::shared_ptr<const ValidPathInfo>>> infos) {
+                            for (auto & [path, info] : infos) {
+                                if (!info) {
+                                    if (required.contains(path))
+                                        throw InvalidPath("path '%s' is not valid", printStorePath(path));
+                                    continue;
+                                }
 
-            for (auto & ref : info->references)
-                if (ref != path)
-                    res.insert(ref);
+                                out.insert(path);
 
-            if (includeOutputs && path.isDerivation())
-                for (auto & [_, maybeOutPath] : queryPartialDerivationOutputMap(path))
-                    if (maybeOutPath && isValidPath(*maybeOutPath))
-                        res.insert(*maybeOutPath);
+                                for (auto & ref : info->references)
+                                    if (ref != path) {
+                                        required.insert(ref);
+                                        todo.insert(ref);
+                                    }
 
-            if (includeDerivers && info->deriver && isValidPath(*info->deriver))
-                res.insert(*info->deriver);
-            co_return res;
-        };
+                                if (includeOutputs && path.isDerivation())
+                                    // FIXME: need an async, multiple-path version of queryPartialDerivationOutputMap().
+                                    for (auto & [_, maybeOutPath] : queryPartialDerivationOutputMap(path))
+                                        if (maybeOutPath)
+                                            todo.insert(*maybeOutPath);
+
+                                if (includeDerivers && info->deriver)
+                                    todo.insert(*info->deriver);
+                            }
+                        });
+                }
+
+                co_return;
+            },
+            [&](std::exception_ptr e) { ex = e; });
 
-    computeClosure<StorePath>(startPaths, paths_, GetEdgesAsync<StorePath>(queryDeps));
+        ctx.run();
+        if (ex)
+            std::rethrow_exception(ex);
+    }
 }
 
 void Store::computeFSClosure(

src/libstore/remote-store.cc

@@ -246,6 +246,71 @@ void RemoteStore::queryPathInfoUncached(
     }
 }
 
+asio::awaitable<void> RemoteStore::queryPathInfos(
+    const std::set<StorePath> & paths,
+    fun<void(std::vector<std::pair<StorePath, std::shared_ptr<const ValidPathInfo>>>)> callback)
+{
+    /* Filter out paths that we already have cached. */
+    StorePathSet uncached;
+    {
+        std::vector<std::pair<StorePath, std::shared_ptr<const ValidPathInfo>>> cached;
+        for (auto & path : paths) {
+            if (auto r = queryPathInfoFromClientCache(path))
+                cached.emplace_back(path, *r);
+            else
+                uncached.insert(path);
+        }
+        if (!cached.empty())
+            callback(std::move(cached));
+    }
+
+    if (uncached.empty())
+        co_return;
+
+    {
+        auto conn(getConnection());
+
+        if (conn->protoVersion.features.contains(WorkerProto::featureQueryPathInfos)) {
+            auto cacheResult = [&](const StorePath & path, std::shared_ptr<const ValidPathInfo> info) {
+                pathInfoCache->lock()->upsert(path, PathInfoCacheValue{.value = info});
+            };
+
+            conn->to << WorkerProto::Op::QueryPathInfos;
+            WorkerProto::write(*this, *conn, uncached);
+            conn.processStderr();
+
+            /* Read the infos for the valid paths. Paths not reported
+               by the daemon are invalid. */
+            auto todo = uncached;
+            auto n = readNum<size_t>(conn->from);
+            std::vector<std::pair<StorePath, std::shared_ptr<const ValidPathInfo>>> results;
+            for (size_t i = 0; i < n; i++) {
+                auto info =
+                    std::make_shared<const ValidPathInfo>(WorkerProto::Serialise<ValidPathInfo>::read(*this, *conn));
+                if (!todo.erase(info->path))
+                    throw Error("daemon returned path info for unexpected path '%s'", printStorePath(info->path));
+                cacheResult(info->path, info);
+                results.emplace_back(info->path, info);
+            }
+
+            for (auto & path : todo) {
+                cacheResult(path, nullptr);
+                results.emplace_back(path, nullptr);
+            }
+
+            callback(std::move(results));
+
+            co_return;
+        }
+
+        /* Release the connection to prevent the fallback below from
+           deadlocking on the connection pool. */
+    }
+
+    /* Fallback for daemons that don't support the batched operation. */
+    co_await Store::queryPathInfos(uncached, std::move(callback));
+}
+
 void RemoteStore::queryReferrers(const StorePath & path, StorePathSet & referrers)
 {
     auto conn(getConnection());

src/libstore/store-api.cc

@@ -613,6 +613,25 @@ void Store::queryPathInfo(const StorePath & storePath, Callback<ref<const ValidP
         }});
 }
 
+asio::awaitable<void> Store::queryPathInfos(
+    const std::set<StorePath> & paths,
+    fun<void(std::vector<std::pair<StorePath, std::shared_ptr<const ValidPathInfo>>>)> callback)
+{
+    /* Default implementation: query each path individually, reporting
+       each result as it arrives. */
+    co_await forEachAsync(paths, [&](const StorePath & path) -> asio::awaitable<void> {
+        std::shared_ptr<const ValidPathInfo> info;
+        try {
+            auto i = co_await callbackToAwaitable<ref<const ValidPathInfo>>(
+                [&](Callback<ref<const ValidPathInfo>> cb) { queryPathInfo(path, std::move(cb)); });
+            info = i.get_ptr();
+        } catch (InvalidPath &) {
+        }
+        std::vector<std::pair<StorePath, std::shared_ptr<const ValidPathInfo>>> result{{path, info}};
+        callback(std::move(result));
+    });
+}
+
 void Store::queryRealisation(
     const DrvOutput & id, Callback<std::shared_ptr<const UnkeyedRealisation>> callback) noexcept
 {

src/libstore/worker-protocol.cc

@@ -27,6 +27,7 @@ const WorkerProto::Version WorkerProto::latest = {
             std::string{WorkerProto::featureProvenance},
             std::string{WorkerProto::featureVersionedAddToStoreMultiple},
             std::string{WorkerProto::featureAddTempRoots},
+            std::string{WorkerProto::featureQueryPathInfos},
         },
 };