1 files changed, 61 insertions, 24 deletions

diff --git a/src/libstore/nar-info-disk-cache.cc b/src/libstore/nar-info-disk-cache.cc
index f4ea739b0..2645f468b 100644
--- a/src/libstore/nar-info-disk-cache.cc
+++ b/src/libstore/nar-info-disk-cache.cc
@@ -84,11 +84,10 @@ public:
 
     Sync<State> _state;
 
-    NarInfoDiskCacheImpl()
+    NarInfoDiskCacheImpl(Path dbPath = getCacheDir() + "/nix/binary-cache-v6.sqlite")
     {
         auto state(_state.lock());
 
-        Path dbPath = getCacheDir() + "/nix/binary-cache-v6.sqlite";
         createDirs(dirOf(dbPath));
 
         state->db = SQLite(dbPath);
@@ -98,7 +97,7 @@ public:
         state->db.exec(schema);
 
         state->insertCache.create(state->db,
-            "insert or replace into BinaryCaches(url, timestamp, storeDir, wantMassQuery, priority) values (?, ?, ?, ?, ?)");
+            "insert into BinaryCaches(url, timestamp, storeDir, wantMassQuery, priority) values (?1, ?2, ?3, ?4, ?5) on conflict (url) do update set timestamp = ?2, storeDir = ?3, wantMassQuery = ?4, priority = ?5 returning id;");
 
         state->queryCache.create(state->db,
             "select id, storeDir, wantMassQuery, priority from BinaryCaches where url = ? and timestamp > ?");
@@ -166,38 +165,71 @@ public:
         return i->second;
     }
 
-    void createCache(const std::string & uri, const Path & storeDir, bool wantMassQuery, int priority) override
+private:
+
+    std::optional<Cache> queryCacheRaw(State & state, const std::string & uri)
     {
-        retrySQLite<void>([&]() {
+        auto i = state.caches.find(uri);
+        if (i == state.caches.end()) {
+            auto queryCache(state.queryCache.use()(uri)(time(0) - cacheInfoTtl));
+            if (!queryCache.next())
+                return std::nullopt;
+            auto cache = Cache {
+                .id = (int) queryCache.getInt(0),
+                .storeDir = queryCache.getStr(1),
+                .wantMassQuery = queryCache.getInt(2) != 0,
+                .priority = (int) queryCache.getInt(3),
+            };
+            state.caches.emplace(uri, cache);
+        }
+        return getCache(state, uri);
+    }
+
+public:
+    int createCache(const std::string & uri, const Path & storeDir, bool wantMassQuery, int priority) override
+    {
+        return retrySQLite<int>([&]() {
             auto state(_state.lock());
+            SQLiteTxn txn(state->db);
+
+            // To avoid the race, we have to check if maybe someone hasn't yet created
+            // the cache for this URI in the meantime.
+            auto cache(queryCacheRaw(*state, uri));
+
+            if (cache)
+                return cache->id;
+
+            Cache ret {
+                .id = -1, // set below
+                .storeDir = storeDir,
+                .wantMassQuery = wantMassQuery,
+                .priority = priority,
+            };
+
+            {
+                auto r(state->insertCache.use()(uri)(time(0))(storeDir)(wantMassQuery)(priority));
+                assert(r.next());
+                ret.id = (int) r.getInt(0);
+            }
 
-            // FIXME: race
+            state->caches[uri] = ret;
 
-            state->insertCache.use()(uri)(time(0))(storeDir)(wantMassQuery)(priority).exec();
-            assert(sqlite3_changes(state->db) == 1);
-            state->caches[uri] = Cache{(int) sqlite3_last_insert_rowid(state->db), storeDir, wantMassQuery, priority};
+            txn.commit();
+            return ret.id;
         });
     }
 
-    std::optional<CacheInfo> cacheExists(const std::string & uri) override
+    std::optional<CacheInfo> upToDateCacheExists(const std::string & uri) override
     {
         return retrySQLite<std::optional<CacheInfo>>([&]() -> std::optional<CacheInfo> {
             auto state(_state.lock());
-
-            auto i = state->caches.find(uri);
-            if (i == state->caches.end()) {
-                auto queryCache(state->queryCache.use()(uri)(time(0) - cacheInfoTtl));
-                if (!queryCache.next())
-                    return std::nullopt;
-                state->caches.emplace(uri,
-                    Cache{(int) queryCache.getInt(0), queryCache.getStr(1), queryCache.getInt(2) != 0, (int) queryCache.getInt(3)});
-            }
-
-            auto & cache(getCache(*state, uri));
-
+            auto cache(queryCacheRaw(*state, uri));
+            if (!cache)
+                return std::nullopt;
             return CacheInfo {
-                .wantMassQuery = cache.wantMassQuery,
-                .priority = cache.priority
+                .id = cache->id,
+                .wantMassQuery = cache->wantMassQuery,
+                .priority = cache->priority
             };
         });
     }
@@ -359,4 +391,9 @@ ref<NarInfoDiskCache> getNarInfoDiskCache()
     return cache;
 }
 
+ref<NarInfoDiskCache> getTestNarInfoDiskCache(Path dbPath)
+{
+    return make_ref<NarInfoDiskCacheImpl>(dbPath);
+}
+
 }