Convert batch Request type from tuple<> to a struct type with traits detection.

While working on other batch requests I found that I wanted to try a ReadRequest type with two string fields. Access to the fields was by type in a few places, and since ReadRequest was implemented as a tuple<...>, duplicate field types are problematic, so I've removed that aspect of ReadRequest.

PiperOrigin-RevId: 832726371
Change-Id: Id6ce5c31edb4f7e16fc79f47bd4a1d2e3534dda7

Author: laramiel

tensorstore/kvstore/batch_util.h

@@ -50,33 +50,65 @@
 namespace tensorstore {
 namespace internal_kvstore_batch {
 
-// Common portion of read request used with `BatchReadEntry`.
+// Type trait that checks if a type is a valid `ReadRequest` type as required
+// by `BatchReadEntry`.  Valid types include any structs with the following
+// members:
 //
-// This is combined in `ReadRequest` with additional fields only used for some
-// kvstore implementations.
+//    Promise<kvstore::ReadResult> promise;
+//    OptionalByteRangeRequest byte_range;
+//
+// Additional members may be present, such as
+// `kvstore::ReadGenerationConditions` or `kvstore::Key`, but at least those two
+// members must be present.
+//
+// See `ByteRangeReadRequest` and `ByteRangeGenerationReadRequest` below for
+// specific types that satisfy this requirement.
+template <typename T, typename = void>
+struct IsByteRangeReadRequestLike : std::false_type {};
+
+template <typename T>
+struct IsByteRangeReadRequestLike<
+    T, std::void_t<decltype(std::declval<T&>().promise),
+                   decltype(std::declval<T&>().byte_range)>>
+    : public std::conjunction<
+          std::is_same<decltype(std::declval<T&>().promise),
+                       Promise<kvstore::ReadResult>>,
+          std::is_same<decltype(std::declval<T&>().byte_range),
+                       OptionalByteRangeRequest>> {};
+
+template <typename T>
+constexpr bool IsByteRangeReadRequestLikeV =
+    IsByteRangeReadRequestLike<T>::value;
+
+// Common `BatchReadEntry` ReadRequest type.
 struct ByteRangeReadRequest {
   Promise<kvstore::ReadResult> promise;
   OptionalByteRangeRequest byte_range;
 };
 
-// Individual read request (entry in batch) used with `BatchReadEntry`.
-//
-// Possibly `Member` types include:
+// Common `BatchReadEntry` ReadRequest type request with generation conditions.
 //
-// - `kvstore::ReadGenerationConditions`
-// - `kvstore::Key`
-// - some derived key type, like `uint64_t`.
-template <typename... Member>
-using ReadRequest = std::tuple<ByteRangeReadRequest, Member...>;
+// This is used when generation conditions are not included in the batch entry
+// key.
+struct ByteRangeGenerationReadRequest {
+  Promise<kvstore::ReadResult> promise;
+  OptionalByteRangeRequest byte_range;
+  kvstore::ReadGenerationConditions generation_conditions;
+};
 
 // Batch of read requests with an aggregate staleness bound, used by
 // `BatchReadEntry`.
 //
 // The aggregate staleness bound is set to the maximum staleness bound of all
 // individual requests.
+//
+// \tparam RequestType Must satisfy `IsByteRangeReadRequestLikeV`.
 template <typename RequestType>
 struct RequestBatch {
+  static_assert(IsByteRangeReadRequestLikeV<RequestType>);
+
   using Request = RequestType;
+
   absl::Time staleness_bound = absl::InfinitePast();
   absl::InlinedVector<Request, 1> requests;
 
@@ -123,10 +155,13 @@ template <typename DerivedDriver, typename RequestType,
           typename... BatchEntryKeyMember>
 class BatchReadEntry : public Batch::Impl::Entry {
  public:
+  static_assert(IsByteRangeReadRequestLikeV<RequestType>);
+
   using Driver = DerivedDriver;
   using BatchEntryKey =
       std::tuple<internal::IntrusivePtr<DerivedDriver>, BatchEntryKeyMember...>;
   using Request = RequestType;
+
   using KeyParam =
       std::tuple<DerivedDriver*, KeyParamType<BatchEntryKeyMember>...>;
 
@@ -204,7 +239,7 @@ class BatchReadEntry : public Batch::Impl::Entry {
   absl::Mutex mutex_;
 
   void AddRequest(absl::Time staleness_bound, Request&& request) {
-    absl::MutexLock lock(&mutex_);
+    absl::MutexLock lock(mutex_);
     request_batch.AddRequest(staleness_bound, std::move(request));
   }
 };
@@ -217,10 +252,13 @@ void SetCommonResult(span<const Request> requests,
                      Result<kvstore::ReadResult>&& result) {
   if (requests.empty()) return;
   for (size_t i = 1; i < requests.size(); ++i) {
-    std::get<ByteRangeReadRequest>(requests[i]).promise.SetResult(result);
+    if (requests[i].promise.result_needed()) {
+      requests[i].promise.SetResult(result);
+    }
+  }
+  if (requests[0].promise.result_needed()) {
+    requests[0].promise.SetResult(std::move(result));
   }
-  std::get<ByteRangeReadRequest>(requests[0])
-      .promise.SetResult(std::move(result));
 }
 template <typename Requests>
 void SetCommonResult(const Requests& requests,
@@ -231,11 +269,10 @@ void SetCommonResult(const Requests& requests,
 
 template <typename Request>
 void SortRequestsByStartByte(span<Request> requests) {
-  std::sort(
-      requests.begin(), requests.end(), [](const Request& a, const Request& b) {
-        return std::get<ByteRangeReadRequest>(a).byte_range.inclusive_min <
-               std::get<ByteRangeReadRequest>(b).byte_range.inclusive_min;
-      });
+  std::sort(requests.begin(), requests.end(),
+            [](const Request& a, const Request& b) {
+              return a.byte_range.inclusive_min < b.byte_range.inclusive_min;
+            });
 }
 
 // Resolves coalesced requests with the appropriate cord subranges.
@@ -244,19 +281,18 @@ void ResolveCoalescedRequests(ByteRange coalesced_byte_range,
                               span<Request> coalesced_requests,
                               kvstore::ReadResult&& read_result) {
   for (auto& request : coalesced_requests) {
-    auto& byte_range_request = std::get<ByteRangeReadRequest>(request);
     kvstore::ReadResult sub_read_result;
     sub_read_result.stamp = read_result.stamp;
     sub_read_result.state = read_result.state;
     if (read_result.state == kvstore::ReadResult::kValue) {
       ABSL_DCHECK_EQ(coalesced_byte_range.size(), read_result.value.size());
-      int64_t request_start = byte_range_request.byte_range.inclusive_min -
-                              coalesced_byte_range.inclusive_min;
-      int64_t request_size = byte_range_request.byte_range.size();
+      int64_t request_start =
+          request.byte_range.inclusive_min - coalesced_byte_range.inclusive_min;
+      int64_t request_size = request.byte_range.size();
       sub_read_result.value =
           read_result.value.Subcord(request_start, request_size);
     }
-    byte_range_request.promise.SetResult(std::move(sub_read_result));
+    request.promise.SetResult(std::move(sub_read_result));
   }
 }
 
@@ -277,19 +313,17 @@ void ResolveCoalescedRequests(ByteRange coalesced_byte_range,
 template <typename Request, typename Predicate, typename Callback>
 void ForEachCoalescedRequest(span<Request> requests, Predicate predicate,
                              Callback callback) {
+  static_assert(IsByteRangeReadRequestLikeV<Request>);
+
   SortRequestsByStartByte(requests);
 
   size_t request_i = 0;
   while (request_i < requests.size()) {
-    auto coalesced_byte_range =
-        std::get<ByteRangeReadRequest>(requests[request_i])
-            .byte_range.AsByteRange();
+    auto coalesced_byte_range = requests[request_i].byte_range.AsByteRange();
     size_t end_request_i;
     for (end_request_i = request_i + 1; end_request_i < requests.size();
          ++end_request_i) {
-      auto next_byte_range =
-          std::get<ByteRangeReadRequest>(requests[end_request_i])
-              .byte_range.AsByteRange();
+      auto next_byte_range = requests[end_request_i].byte_range.AsByteRange();
       if (next_byte_range.inclusive_min < coalesced_byte_range.exclusive_max ||
           predicate(coalesced_byte_range, next_byte_range.inclusive_min)) {
         coalesced_byte_range.exclusive_max = std::max(
@@ -314,12 +348,9 @@ void ForEachCoalescedRequest(span<Request> requests, Predicate predicate,
 template <typename Request>
 bool ValidateRequestGeneration(Request& request,
                                const TimestampedStorageGeneration& stamp) {
-  auto& byte_range_request = std::get<ByteRangeReadRequest>(request);
-  if (!byte_range_request.promise.result_needed()) return false;
-  if (!std::get<kvstore::ReadGenerationConditions>(request).Matches(
-          stamp.generation)) {
-    byte_range_request.promise.SetResult(
-        kvstore::ReadResult::Unspecified(stamp));
+  if (!request.promise.result_needed()) return false;
+  if (!request.generation_conditions.Matches(stamp.generation)) {
+    request.promise.SetResult(kvstore::ReadResult::Unspecified(stamp));
     return false;
   }
   return true;
@@ -334,14 +365,15 @@ bool ValidateRequestGeneration(Request& request,
 template <typename Request>
 bool ValidateRequestGenerationAndByteRange(
     Request& request, const TimestampedStorageGeneration& stamp, int64_t size) {
+  static_assert(IsByteRangeReadRequestLikeV<Request>);
+  static_assert(
+      std::is_member_pointer<decltype(&Request::generation_conditions)>::value);
   if (!ValidateRequestGeneration(request, stamp)) {
     return false;
   }
-  auto& byte_range_request = std::get<ByteRangeReadRequest>(request);
   TENSORSTORE_ASSIGN_OR_RETURN(
-      byte_range_request.byte_range,
-      byte_range_request.byte_range.Validate(size),
-      (byte_range_request.promise.SetResult(std::move(_)), false));
+      request.byte_range, request.byte_range.Validate(size),
+      (request.promise.SetResult(std::move(_)), false));
   return true;
 }
 

tensorstore/kvstore/file/file_key_value_store.cc

@@ -364,7 +364,7 @@ Result<UniqueFileDescriptor> OpenValueFile(const std::string& path,
 class BatchReadTask;
 using BatchReadTaskBase = internal_kvstore_batch::BatchReadEntry<
     FileKeyValueStore,
-    internal_kvstore_batch::ReadRequest<kvstore::ReadGenerationConditions>,
+    /*ReadRequest=*/internal_kvstore_batch::ByteRangeGenerationReadRequest,
     // BatchEntryKey members:
     std::string /* file_path*/>;
 
@@ -448,12 +448,9 @@ class BatchReadTask final
     }
 
     if (requests.size() == 1) {
-      auto& byte_range_request =
-          std::get<internal_kvstore_batch::ByteRangeReadRequest>(requests[0]);
       // Perform single read immediately.
-      byte_range_request.promise.SetResult(
-          DoByteRangeRead(byte_range_request.byte_range.AsByteRange()));
-
+      requests[0].promise.SetResult(
+          DoByteRangeRead(requests[0].byte_range.AsByteRange()));
       return;
     }
 
@@ -480,9 +477,7 @@ class BatchReadTask final
     int64_t inclusive_min = std::numeric_limits<int64_t>::max();
     int64_t total_size = 0;
     for (const auto& req : requests) {
-      const auto byte_range =
-          std::get<internal_kvstore_batch::ByteRangeReadRequest>(req)
-              .byte_range.AsByteRange();
+      const auto byte_range = req.byte_range.AsByteRange();
       inclusive_min = std::min(inclusive_min, byte_range.inclusive_min);
       exclusive_max = std::max(exclusive_max, byte_range.exclusive_max);
       total_size += byte_range.size();
@@ -503,13 +498,11 @@ class BatchReadTask final
       } else if (mapped_result.ok()) {
         absl::Cord file_contents = std::move(mapped_result).value().as_cord();
         for (const auto& req : requests) {
-          auto& byte_range_request =
-              std::get<internal_kvstore_batch::ByteRangeReadRequest>(req);
-          ByteRange byte_range = byte_range_request.byte_range.AsByteRange();
+          ByteRange byte_range = req.byte_range.AsByteRange();
           assert(byte_range.inclusive_min >= inclusive_min);
           absl::Cord subcord = file_contents.Subcord(
               byte_range.inclusive_min - inclusive_min, byte_range.size());
-          byte_range_request.promise.SetResult(
+          req.promise.SetResult(
               kvstore::ReadResult::Value(std::move(subcord), stamp_));
         }
         return true;
@@ -530,9 +523,7 @@ class BatchReadTask final
     // Determine if it's possible to read entire blocks.
     int64_t exclusive_max = 0;
     for (const auto& req : requests) {
-      const auto byte_range =
-          std::get<internal_kvstore_batch::ByteRangeReadRequest>(req)
-              .byte_range.AsByteRange();
+      const auto byte_range = req.byte_range.AsByteRange();
       exclusive_max = std::max(exclusive_max, byte_range.exclusive_max);
     }
     exclusive_max = RoundUpTo(exclusive_max, block_alignment);
@@ -565,7 +556,7 @@ Future<ReadResult> FileKeyValueStore::Read(Key key, ReadOptions options) {
   auto [promise, future] = PromiseFuturePair<kvstore::ReadResult>::Make();
   BatchReadTask::MakeRequest<BatchReadTask>(
       *this, {std::move(key)}, options.batch, options.staleness_bound,
-      BatchReadTask::Request{{std::move(promise), options.byte_range},
+      BatchReadTask::Request{std::move(promise), options.byte_range,
                              std::move(options.generation_conditions)});
   return std::move(future);
 }

tensorstore/kvstore/generic_coalescing_batch_util.h

@@ -36,7 +36,7 @@ namespace internal_kvstore_batch {
 
 template <typename DerivedDriver>
 using GenericCoalescingBatchReadEntryBase =
-    BatchReadEntry<DerivedDriver, ReadRequest<>,
+    BatchReadEntry<DerivedDriver, /*ReadRequest=*/ByteRangeReadRequest,
                    // BatchEntryKey members:
                    kvstore::Key, kvstore::ReadGenerationConditions>;
 
@@ -135,7 +135,7 @@ Future<kvstore::ReadResult> HandleBatchRequestByGenericByteRangeCoalescing(
   Entry::template MakeRequest<Entry>(
       driver, std::move(key), std::move(options.generation_conditions),
       options.batch, options.staleness_bound,
-      typename Entry::Request{{std::move(promise), options.byte_range}});
+      typename Entry::Request{std::move(promise), options.byte_range});
   return std::move(future);
 }
 

tensorstore/kvstore/mock_kvstore.cc

@@ -89,13 +89,9 @@ void MockKeyValueStore::BatchReadRequest::operator()(
     kvstore::ReadOptions options;
     options.staleness_bound = request_batch.staleness_bound;
     options.batch = batch;
-    options.generation_conditions =
-        std::get<kvstore::ReadGenerationConditions>(request);
-    auto& byte_range_request =
-        std::get<internal_kvstore_batch::ByteRangeReadRequest>(request);
-    options.byte_range = byte_range_request.byte_range;
-    LinkResult(byte_range_request.promise,
-               target->Read(key, std::move(options)));
+    options.generation_conditions = request.generation_conditions;
+    options.byte_range = request.byte_range;
+    LinkResult(request.promise, target->Read(key, std::move(options)));
   }
 }
 
@@ -104,7 +100,9 @@ Future<kvstore::ReadResult> MockKeyValueStore::Read(Key key,
   if (handle_batch_requests && options.batch) {
     class BatchEntry;
     using BatchEntryBase = internal_kvstore_batch::BatchReadEntry<
-        MockKeyValueStore, BatchReadRequest::Request, kvstore::Key>;
+        MockKeyValueStore, /*ReadRequest=*/BatchReadRequest::Request,
+        // BatchEntryKey members:
+        kvstore::Key>;
     class BatchEntry : public BatchEntryBase {
      public:
       using BatchEntryBase::BatchEntryBase;
@@ -120,13 +118,9 @@ Future<kvstore::ReadResult> MockKeyValueStore::Read(Key key,
           ::nlohmann::json::array_t requests_log;
           for (const auto& request : request_batch.requests) {
             ::nlohmann::json::object_t request_log;
-            AddByteRangeToLogEntry(
-                request_log,
-                std::get<internal_kvstore_batch::ByteRangeReadRequest>(request)
-                    .byte_range);
-            AddGenerationConditionsToLogEntry(
-                request_log,
-                std::get<kvstore::ReadGenerationConditions>(request));
+            AddByteRangeToLogEntry(request_log, request.byte_range);
+            AddGenerationConditionsToLogEntry(request_log,
+                                              request.generation_conditions);
             requests_log.push_back(std::move(request_log));
           }
           log_entry.emplace("requests", std::move(requests_log));
@@ -147,7 +141,7 @@ Future<kvstore::ReadResult> MockKeyValueStore::Read(Key key,
     auto [promise, future] = PromiseFuturePair<kvstore::ReadResult>::Make();
     BatchEntry::MakeRequest<BatchEntry>(
         *this, std::move(key), options.batch, options.staleness_bound,
-        BatchEntry::Request{{std::move(promise), options.byte_range},
+        BatchEntry::Request{std::move(promise), options.byte_range,
                             std::move(options.generation_conditions)});
     return std::move(future);
   }

tensorstore/kvstore/mock_kvstore.h

@@ -26,7 +26,6 @@
 #include "tensorstore/kvstore/driver.h"
 #include "tensorstore/kvstore/generation.h"
 #include "tensorstore/kvstore/key_range.h"
-#include "tensorstore/kvstore/operations.h"
 #include "tensorstore/kvstore/spec.h"
 #include "tensorstore/kvstore/supported_features.h"
 #include "tensorstore/util/future.h"
@@ -58,9 +57,9 @@ class MockKeyValueStore : public kvstore::Driver {
   };
 
   struct BatchReadRequest {
+    using Request = internal_kvstore_batch::ByteRangeGenerationReadRequest;
+
     Key key;
-    using Request =
-        internal_kvstore_batch::ReadRequest<kvstore::ReadGenerationConditions>;
     using RequestBatch = internal_kvstore_batch::RequestBatch<Request>;
     RequestBatch request_batch;
     void operator()(kvstore::DriverPtr target) const;