libstore: Optimize BinaryCacheStore::queryValidPaths

src/libstore-tests/binary-cache-store.cc

@@ -0,0 +1,92 @@
+#include <gtest/gtest.h>
+
+#include "nix/store/binary-cache-store.hh"
+#include "nix/store/globals.hh"
+
+namespace nix {
+
+struct TestBinaryCacheStoreConfig : virtual Store::Config, BinaryCacheStoreConfig
+{
+    using Params = StoreReference::Params;
+
+    TestBinaryCacheStoreConfig(const Params & params)
+        : Store::Config(params)
+        , BinaryCacheStoreConfig(params)
+    {
+    }
+
+    ref<Store> openStore() const override
+    {
+        throw Unsupported("openStore");
+    }
+};
+
+struct TestBinaryCacheStore : virtual BinaryCacheStore
+{
+    using Config = TestBinaryCacheStoreConfig;
+
+    ref<Config> config;
+    std::atomic<size_t> fileExistsCalls = 0;
+    std::atomic<size_t> getFileCalls = 0;
+    std::set<std::string> existingFiles;
+
+    TestBinaryCacheStore(ref<Config> config)
+        : Store{*config}
+        , BinaryCacheStore{*config}
+        , config(config)
+    {
+    }
+
+    std::optional<TrustedFlag> isTrustedClient() override
+    {
+        return std::nullopt;
+    }
+
+    bool fileExists(const std::string & path) override
+    {
+        fileExistsCalls++;
+        return existingFiles.contains(path);
+    }
+
+    void upsertFile(
+        const std::string & path, RestartableSource & source, const std::string & mimeType, uint64_t sizeHint) override
+    {
+        throw Unsupported("upsertFile");
+    }
+
+    void getFile(const std::string & path, Sink & sink) override
+    {
+        getFileCalls++;
+        throw NoSuchBinaryCacheFile("file '%s' does not exist in binary cache", path);
+    }
+};
+
+TEST(BinaryCacheStore, queryValidPathsUsesExistenceChecks)
+{
+    initLibStore(false);
+
+    auto config = make_ref<TestBinaryCacheStoreConfig>(StoreReference::Params{});
+    auto store = std::make_shared<TestBinaryCacheStore>(config);
+
+    StorePathSet paths{
+        StorePath{"g1w7hy3qg1w7hy3qg1w7hy3qg1w7hy3q-foo"},
+        StorePath{"g1w7hy3qg1w7hy3qg1w7hy3qg1w7hy3r-bar"},
+        StorePath{"g1w7hy3qg1w7hy3qg1w7hy3qg1w7hy3s-baz"},
+    };
+
+    store->existingFiles = {
+        "g1w7hy3qg1w7hy3qg1w7hy3qg1w7hy3q.narinfo",
+        "g1w7hy3qg1w7hy3qg1w7hy3qg1w7hy3s.narinfo",
+    };
+
+    auto valid = store->queryValidPaths(paths);
+
+    EXPECT_EQ(valid.size(), 2u);
+    EXPECT_EQ(valid.count(StorePath{"g1w7hy3qg1w7hy3qg1w7hy3qg1w7hy3q-foo"}), 1u);
+    EXPECT_EQ(valid.count(StorePath{"g1w7hy3qg1w7hy3qg1w7hy3qg1w7hy3s-baz"}), 1u);
+
+    EXPECT_EQ(store->fileExistsCalls, paths.size());
+    EXPECT_EQ(store->getFileCalls, 0u);
+}
+
+} // namespace nix

src/libstore-tests/meson.build

@@ -54,6 +54,7 @@ deps_private += gtest
 subdir('nix-meson-build-support/common')
 
 sources = files(
+  'binary-cache-store.cc',
   'build-result.cc',
   'common-protocol.cc',
   'content-address.cc',

src/libstore/binary-cache-store.cc

@@ -8,6 +8,7 @@
 #include "nix/util/sync.hh"
 #include "nix/store/remote-fs-accessor.hh"
 #include "nix/store/nar-info-disk-cache.hh"
+#include "nix/store/filetransfer.hh"
 #include "nix/util/nar-accessor.hh"
 #include "nix/util/thread-pool.hh"
 #include "nix/util/callback.hh"
@@ -383,6 +384,61 @@ StorePath BinaryCacheStore::addToStoreFromDump(
         ->path;
 }
 
+StorePathSet BinaryCacheStore::queryValidPaths(const StorePathSet & paths, SubstituteFlag maybeSubstitute)
+{
+    struct State
+    {
+        size_t left;
+        StorePathSet valid;
+        std::exception_ptr exc;
+    };
+
+    Sync<State> state_(State{paths.size(), StorePathSet()});
+
+    std::condition_variable wakeup;
+    ThreadPool pool(fileTransferSettings.httpConnections);
+
+    auto doQuery = [&](const StorePath & path) {
+        checkInterrupt();
+
+        bool exists = false;
+        std::exception_ptr newExc{};
+
+        try {
+            exists = isValidPath(path);
+        } catch (...) {
+            newExc = std::current_exception();
+        }
+
+        auto state(state_.lock());
+
+        if (exists)
+            state->valid.insert(path);
+
+        if (newExc)
+            state->exc = newExc;
+
+        assert(state->left);
+        if (!--state->left)
+            wakeup.notify_one();
+    };
+
+    for (auto & path : paths)
+        pool.enqueue(std::bind(doQuery, path));
+
+    pool.process();
+
+    while (true) {
+        auto state(state_.lock());
+        if (!state->left) {
+            if (state->exc)
+                std::rethrow_exception(state->exc);
+            return std::move(state->valid);
+        }
+        state.wait(wakeup);
+    }
+}
+
 bool BinaryCacheStore::isValidPathUncached(const StorePath & storePath)
 {
     // FIXME: this only checks whether a .narinfo with a matching hash

src/libstore/include/nix/store/binary-cache-store.hh

@@ -168,6 +168,8 @@ public:
 
     bool isValidPathUncached(const StorePath & path) override;
 
+    StorePathSet queryValidPaths(const StorePathSet & paths, SubstituteFlag maybeSubstitute = NoSubstitute) override;
+
     void queryPathInfoUncached(
         const StorePath & path, Callback<std::shared_ptr<const ValidPathInfo>> callback) noexcept override;