Fix data race when releasing weak pool reference held by cache

PiperOrigin-RevId: 805961624
Change-Id: Id86ba503a3f3b256363037be831201c0fbdd0f26

Author: jbms

tensorstore/internal/cache/cache.cc

@@ -557,16 +557,16 @@ PinnedCacheEntry<Cache> GetCacheEntryInternal(internal::Cache* cache,
 }
 
 void StrongPtrTraitsCache::decrement_impl(CacheImpl* cache) noexcept {
+  // Note: `pool_` must be accessed here because `cache` may have been
+  // concurrently destroyed after calling `DecrementCacheReferenceCount` unless
+  // `should_delete==true`.
+  CachePoolImpl* pool = cache->pool_;
   auto decrement_result =
       DecrementCacheReferenceCount(cache, CacheImpl::kStrongReferenceIncrement);
-  CachePoolImpl* pool = nullptr;
-  if (decrement_result.should_release_cache_pool_weak_reference()) {
-    pool = cache->pool_;
-  }
   if (decrement_result.should_delete()) {
-    DestroyCache(cache->pool_, cache);
+    DestroyCache(pool, cache);
   }
-  if (pool) {
+  if (decrement_result.should_release_cache_pool_weak_reference() && pool) {
     ReleaseWeakReference(pool);
   }
 }

tensorstore/internal/cache/cache_test.cc

@@ -1113,7 +1113,9 @@ TEST(CacheTest, WeakRefOwnedByEntry) {
     auto entry_b = GetCacheEntry(cache1, "b");
     entry_a->weak_ref = entry_b->AcquireWeakReference();
   }
-  { auto entry_c = GetCacheEntry(cache2, "c"); }
+  {
+    auto entry_c = GetCacheEntry(cache2, "c");
+  }
 
   EXPECT_THAT(log->entry_destroy_log, ElementsAre());
   pool = {};
@@ -1127,4 +1129,29 @@ TEST(CacheTest, WeakRefOwnedByEntry) {
                                    Pair("cache2", "c")));
 }
 
+TEST(CacheTest, ConcurrentReleaseStrongCachePoolEvict) {
+  CachePool::StrongPtr pool;
+  CachePtr<TestCache> cache1, cache2;
+  TestConcurrent(
+      kDefaultIterations,
+      /*initialize=*/
+      [&] {
+        CachePool::Limits limits;
+        limits.total_bytes_limit = 1;
+        pool = CachePool::Make(limits);
+        cache1 = GetTestCache(pool.get(), "x");
+        cache2 = GetTestCache(pool.get(), "y");
+        GetCacheEntry(cache1, "a");
+      },
+      /*finalize=*/
+      [&] {},
+      // Concurrent operations:
+      [&] { pool = {}; },    // release strong reference to pool
+      [&] { cache1 = {}; },  // release strong reference to cache1
+      [&] {
+        GetCacheEntry(cache2, "b");
+      }  // create new cache entry in cache2, evicting only entry in cache1
+  );
+}
+
 }  // namespace