Refactor cache locking and eliminate queued_for_writeback_bytes_limit

The refactoring cache locking is intended to facilitate later
concurrent improvements.

Non-transactional writes are now always committed immediately.  To
defer committing, a transaction can be used instead.

PiperOrigin-RevId: 599336232
Change-Id: I9b88690e8106bf898bbadd25996a628b044228b2

Author: jbms

README.md

@@ -21,7 +21,7 @@ storage and manipulation of large multi-dimensional arrays that:
 * Offers an asynchronous API to enable high-throughput access even to
   high-latency remote storage.
 
-* Supports read/writeback caching and transactions, with strong atomicity,
+* Supports read caching and transactions, with strong atomicity,
   isolation, consistency, and durability (ACID) guarantees.
 
 * Supports safe, efficient access from multiple processes and machines via

docs/context_schema.yml

@@ -61,14 +61,6 @@ definitions:
           least-recently used data that is not in use is evicted from the cache
           when this limit is reached.
         default: 0
-      queued_for_writeback_bytes_limit:
-        type: integer
-        minimum: 0
-        description: |-
-          Soft limit on the total number of bytes of data pending writeback.
-          Writeback is initated on the least-recently used data that is pending
-          writeback when this limit is reached.  Defaults to half of
-          `.total_bytes_limit`.
   data_copy_concurrency:
     $id: Context.data_copy_concurrency
     description: |-

tensorstore/driver/image/driver_impl.h

@@ -347,7 +347,7 @@ Future<internal::DriverHandle> ImageDriverSpec<Specialization>::Open(
   internal::EncodeCacheKey(&cache_identifier, store.driver,
                            data_copy_concurrency, store.path);
   auto cache = internal::GetOrCreateAsyncInitializedCache<CacheType>(
-      **cache_pool, cache_identifier,
+      cache_pool->get(), cache_identifier,
       [&] {
         auto cache = std::make_unique<CacheType>();
         cache->data_copy_concurrency_ = data_copy_concurrency;

tensorstore/driver/json/driver.cc

@@ -357,7 +357,7 @@ Future<internal::Driver::Handle> JsonDriverSpec::Open(
   internal::EncodeCacheKey(&cache_identifier, store.driver,
                            data_copy_concurrency);
   auto cache = internal::GetOrCreateAsyncInitializedCache<JsonCache>(
-      **cache_pool, cache_identifier,
+      cache_pool->get(), cache_identifier,
       [&] {
         auto cache = std::make_unique<JsonCache>();
         cache->data_copy_concurrency_ = data_copy_concurrency;

tensorstore/driver/json/driver_test.cc

@@ -249,6 +249,7 @@ TEST(JsonDriverTest, ReadError) {
   {
     auto write_future =
         tensorstore::Write(MakeScalarArray<::nlohmann::json>(42), store);
+    write_future.Force();
     TENSORSTORE_EXPECT_OK(write_future.copy_future);
     mock_key_value_store->read_requests.pop().promise.SetResult(
         absl::UnknownError("read error2"));
@@ -274,6 +275,7 @@ TEST(JsonDriverTest, ConditionalWriteback) {
   {
     auto write_future =
         tensorstore::Write(MakeScalarArray<::nlohmann::json>(42), store);
+    write_future.Force();
     TENSORSTORE_EXPECT_OK(write_future.copy_future);
     mock_key_value_store->read_requests.pop()(memory_store);
     mock_key_value_store->write_requests.pop()(memory_store);
@@ -284,6 +286,7 @@ TEST(JsonDriverTest, ConditionalWriteback) {
   {
     auto write_future =
         tensorstore::Write(MakeScalarArray<::nlohmann::json>(42), store);
+    write_future.Force();
     TENSORSTORE_EXPECT_OK(write_future.copy_future);
     mock_key_value_store->read_requests.pop()(memory_store);
     // No write request, since value is unchanged.
@@ -304,6 +307,7 @@ TEST(JsonDriverTest, UnconditionalWriteback) {
                                    tensorstore::Open(spec, context).result());
   auto write_future =
       tensorstore::Write(MakeScalarArray<::nlohmann::json>(42), store);
+  write_future.Force();
   {
     auto write_req = mock_key_value_store->write_requests.pop();
     EXPECT_EQ(tensorstore::StorageGeneration::Unknown(),
@@ -315,29 +319,21 @@ TEST(JsonDriverTest, UnconditionalWriteback) {
 
 TEST(JsonDriverTest, ZeroElementWrite) {
   auto json_spec = GetSpec("");
-  json_spec["cache_pool"] = {{"total_bytes_limit", 10000000}};
-  TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store,
-                                   tensorstore::Open(json_spec).result());
-  // Confirm that a one-element write is not immediately committed due to cache.
-  {
-    auto write_future =
-        tensorstore::Write(MakeScalarArray<::nlohmann::json>(42), store);
-    TENSORSTORE_EXPECT_OK(write_future.copy_future);
-    absl::SleepFor(absl::Milliseconds(10));
-    EXPECT_FALSE(write_future.commit_future.ready());
-    // When forced, future becomes ready.
-    TENSORSTORE_EXPECT_OK(write_future.commit_future);
-  }
-
+  json_spec["kvstore"] = {{"driver", "mock_key_value_store"},
+                          {"path", GetPath()}};
+  auto context = tensorstore::Context::Default();
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(
+      auto mock_key_value_store_resource,
+      context.GetResource<tensorstore::internal::MockKeyValueStoreResource>());
+  auto mock_key_value_store = *mock_key_value_store_resource;
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(
+      auto store, tensorstore::Open(json_spec, context).result());
   // Test that a write to zero elements is detected as a non-modification, and
   // leads to an immediately-ready future.
   {
     auto write_future = tensorstore::Write(
         tensorstore::AllocateArray<::nlohmann::json>({0}),
         store | tensorstore::Dims(0).AddNew().SizedInterval(0, 0));
-    TENSORSTORE_EXPECT_OK(write_future.copy_future);
-    absl::SleepFor(absl::Milliseconds(10));
-    EXPECT_TRUE(write_future.commit_future.ready());
     TENSORSTORE_EXPECT_OK(write_future.commit_future);
   }
 }

tensorstore/driver/kvs_backed_chunk_driver.cc

@@ -825,25 +825,22 @@ Result<internal::Driver::Handle> CreateTensorStoreFromMetadata(
   }
   absl::Status data_key_value_store_status;
   const auto& state_ref = *state;
-  auto data_cache =
-      (*state->cache_pool())
-          ->GetCache<DataCacheBase>(
-              typeid(state_ref), chunk_cache_identifier,
-              [&]() -> std::unique_ptr<DataCacheBase> {
-                auto store_result = state->GetDataKeyValueStore(
-                    GetOwningCache(*base.metadata_cache_entry_).base_store_,
-                    metadata.get());
-                if (!store_result) {
-                  data_key_value_store_status =
-                      std::move(store_result).status();
-                  return nullptr;
-                }
-                DataCacheInitializer initializer;
-                initializer.store = std::move(*store_result);
-                initializer.metadata_cache_entry = base.metadata_cache_entry_;
-                initializer.metadata = metadata;
-                return state->GetDataCache(std::move(initializer));
-              });
+  auto data_cache = internal::GetCacheWithExplicitTypeInfo<DataCacheBase>(
+      state->cache_pool()->get(), typeid(state_ref), chunk_cache_identifier,
+      [&]() -> std::unique_ptr<DataCacheBase> {
+        auto store_result = state->GetDataKeyValueStore(
+            GetOwningCache(*base.metadata_cache_entry_).base_store_,
+            metadata.get());
+        if (!store_result) {
+          data_key_value_store_status = std::move(store_result).status();
+          return nullptr;
+        }
+        DataCacheInitializer initializer;
+        initializer.store = std::move(*store_result);
+        initializer.metadata_cache_entry = base.metadata_cache_entry_;
+        initializer.metadata = metadata;
+        return state->GetDataCache(std::move(initializer));
+      });
   TENSORSTORE_RETURN_IF_ERROR(data_key_value_store_status);
   TENSORSTORE_ASSIGN_OR_RETURN(
       auto new_transform,
@@ -1225,7 +1222,7 @@ internal::CachePtr<MetadataCache> GetOrCreateMetadataCache(
   internal::EncodeCacheKey(&base.metadata_cache_key_, spec.store.driver,
                            typeid(*state), state->GetMetadataCacheKey());
   return internal::GetOrCreateAsyncInitializedCache<MetadataCache>(
-      **state->cache_pool(), base.metadata_cache_key_,
+      state->cache_pool()->get(), base.metadata_cache_key_,
       [&] {
         ABSL_LOG_IF(INFO, TENSORSTORE_KVS_DRIVER_DEBUG)
             << "Creating metadata cache: open_state=" << state;

tensorstore/driver/neuroglancer_precomputed/BUILD

@@ -82,6 +82,7 @@ tensorstore_cc_test(
         "//tensorstore:staleness_bound",
         "//tensorstore:static_cast",
         "//tensorstore:strided_layout",
+        "//tensorstore:transaction",
         "//tensorstore/driver:driver_testutil",
         "//tensorstore/driver/zarr",
         "//tensorstore/index_space:dim_expression",

tensorstore/driver/neuroglancer_precomputed/driver_test.cc

@@ -61,6 +61,7 @@
 #include "tensorstore/static_cast.h"
 #include "tensorstore/strided_layout.h"
 #include "tensorstore/tensorstore.h"
+#include "tensorstore/transaction.h"
 #include "tensorstore/util/dimension_set.h"
 #include "tensorstore/util/result.h"
 #include "tensorstore/util/status.h"
@@ -1501,9 +1502,7 @@ TEST(ShardedWriteTest, Basic) {
   EXPECT_THAT(tensorstore::Read(store).result(), ::testing::Optional(array));
 }
 
-// Disable due to race condition whereby writeback of a shard may start while
-// some chunks that have been modified are still being written back to it.
-TEST(FullShardWriteTest, Basic) {
+TEST(FullShardWriteTest, WithTransaction) {
   auto context = Context::Default();
 
   TENSORSTORE_ASSERT_OK_AND_ASSIGN(
@@ -1512,8 +1511,6 @@ TEST(FullShardWriteTest, Basic) {
   auto mock_key_value_store = *mock_key_value_store_resource;
 
   ::nlohmann::json json_spec{
-      // Use a cache to avoid early writeback of partial shard.
-      {"context", {{"cache_pool", {{"total_bytes_limit", 10'000'000}}}}},
       {"driver", "neuroglancer_precomputed"},
       {"kvstore",
        {
@@ -1556,51 +1553,46 @@ TEST(FullShardWriteTest, Basic) {
   // Shard 4 origin: {0, 0, 8}
   // Shard 5 origin: {0, 4, 8}
 
-  // Repeat the test to try to detect errors due to possible timing-dependent
-  // behavior differences.
-  for (int i = 0; i < 100; ++i) {
-    auto store_future = tensorstore::Open(json_spec, context);
-    store_future.Force();
+  tensorstore::Transaction txn(tensorstore::isolated);
+  auto store_future = tensorstore::Open(json_spec, context);
+  store_future.Force();
 
-    {
-      auto req = mock_key_value_store->read_requests.pop();
-      EXPECT_EQ("prefix/info", req.key);
-      req.promise.SetResult(kvstore::ReadResult::Missing(absl::Now()));
-    }
-
-    {
-      auto req = mock_key_value_store->write_requests.pop();
-      EXPECT_EQ("prefix/info", req.key);
-      EXPECT_EQ(StorageGeneration::NoValue(), req.options.if_equal);
-      req.promise.SetResult(TimestampedStorageGeneration{
-          StorageGeneration::FromString("g0"), absl::Now()});
-    }
+  {
+    auto req = mock_key_value_store->read_requests.pop();
+    EXPECT_EQ("prefix/info", req.key);
+    req.promise.SetResult(kvstore::ReadResult::Missing(absl::Now()));
+  }
 
-    TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, store_future.result());
+  {
+    auto req = mock_key_value_store->write_requests.pop();
+    EXPECT_EQ("prefix/info", req.key);
+    EXPECT_EQ(StorageGeneration::NoValue(), req.options.if_equal);
+    req.promise.SetResult(TimestampedStorageGeneration{
+        StorageGeneration::FromString("g0"), absl::Now()});
+  }
 
-    auto future = tensorstore::Write(
-        tensorstore::MakeScalarArray<uint16_t>(42),
-        tensorstore::ChainResult(
-            store,
-            tensorstore::Dims(0, 1, 2).SizedInterval({0, 4, 8}, {4, 2, 2})));
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, store_future.result());
 
-    // Ensure copying finishes before writeback starts.
-    TENSORSTORE_ASSERT_OK(future.copy_future.result());
-    ASSERT_FALSE(future.commit_future.ready());
+  auto future = tensorstore::Write(
+      tensorstore::MakeScalarArray<uint16_t>(42),
+      store | txn |
+          tensorstore::Dims(0, 1, 2).SizedInterval({0, 4, 8}, {4, 2, 2}));
 
-    future.Force();
+  // Ensure copying finishes before writeback starts.
+  TENSORSTORE_ASSERT_OK(future);
 
-    {
-      auto req = mock_key_value_store->write_requests.pop();
-      ASSERT_EQ("prefix/1_1_1/5.shard", req.key);
-      // Writeback is unconditional because the entire shard is being written.
-      ASSERT_EQ(StorageGeneration::Unknown(), req.options.if_equal);
-      req.promise.SetResult(TimestampedStorageGeneration{
-          StorageGeneration::FromString("g0"), absl::Now()});
-    }
+  txn.CommitAsync().IgnoreFuture();
 
-    TENSORSTORE_ASSERT_OK(future.result());
+  {
+    auto req = mock_key_value_store->write_requests.pop();
+    ASSERT_EQ("prefix/1_1_1/5.shard", req.key);
+    // Writeback is unconditional because the entire shard is being written.
+    ASSERT_EQ(StorageGeneration::Unknown(), req.options.if_equal);
+    req.promise.SetResult(TimestampedStorageGeneration{
+        StorageGeneration::FromString("g0"), absl::Now()});
   }
+
+  TENSORSTORE_ASSERT_OK(txn.future());
 }
 
 // Tests that an empty path is handled correctly.

tensorstore/driver/stack/driver_test.cc

@@ -85,7 +85,7 @@ ::nlohmann::json GetRank1Length4N5Driver(int inclusive_min,
   if (exclusive_max == -1) {
     exclusive_max = inclusive_min + 4;
   }
-  auto path = absl::StrCat("p", inclusive_min, exclusive_max, "/");
+  auto path = absl::StrCat("p", inclusive_min, "_", exclusive_max, "/");
   auto result = ::nlohmann::json{
       {"driver", "n5"},
       {"kvstore",

tensorstore/driver/virtual_chunked/BUILD

@@ -24,6 +24,7 @@ tensorstore_cc_library(
         "//tensorstore/driver:chunk_cache_driver",
         "//tensorstore/index_space:index_transform",
         "//tensorstore/internal:data_copy_concurrency_resource",
+        "//tensorstore/internal/cache",
         "//tensorstore/internal/cache:cache_pool_resource",
         "//tensorstore/internal/cache:chunk_cache",
         "//tensorstore/kvstore:generation",