kvstore/gcs_grpc updates:

* Minor consistency and other changes for gcs_grpc ops.
* Add additional tests for gcs_grpc read, with a few tensorstore::Batch cases.
* Cleanup tsan issue in gcs_grpc_testbench hooks.

PiperOrigin-RevId: 833977592
Change-Id: I1ea1003ae0ce61418afe8837444ad587d9e4360a

Author: laramiel

tensorstore/internal/grpc/grpc_mock.h

@@ -25,7 +25,7 @@
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "grpc/grpc_security_constants.h"
-#include "grpc/support/log.h"
+#include "grpc/support/log.h"  // IWYU pragma: keep
 #include "grpcpp/channel.h"  // third_party
 #include "grpcpp/completion_queue.h"  // third_party
 #include "grpcpp/create_channel.h"  // third_party
@@ -35,6 +35,7 @@
 #include "grpcpp/server.h"  // third_party
 #include "grpcpp/server_builder.h"  // third_party
 #include "grpcpp/server_context.h"  // third_party
+#include "grpcpp/support/status.h"  // third_party  // IWYU pragma: keep
 
 namespace tensorstore {
 namespace grpc_mocker {
@@ -117,7 +118,7 @@ class MockGrpcServer {
     builder.SetSyncServerOption(::grpc::ServerBuilder::NUM_CQS, 2)
         .SetSyncServerOption(::grpc::ServerBuilder::MIN_POLLERS, 1)
         .SetSyncServerOption(::grpc::ServerBuilder::MAX_POLLERS, 2)
-        .SetSyncServerOption(::grpc::ServerBuilder::CQ_TIMEOUT_MSEC, 10000);
+        .SetSyncServerOption(::grpc::ServerBuilder::CQ_TIMEOUT_MSEC, 1000);
 
     builder.AddListeningPort(
         "localhost:0", grpc::experimental::LocalServerCredentials(LOCAL_TCP),

tensorstore/kvstore/batch_util.h

@@ -221,7 +221,7 @@ class BatchReadEntry : public Batch::Impl::Entry {
   // This may be useful for implementing retry logic.
   template <typename DerivedEntry>
   std::unique_ptr<DerivedEntry> MigrateExistingRequestsToNewEntry(
-      span<Request> existing_requests) {
+      tensorstore::span<Request> existing_requests) {
     using Self =
         BatchReadEntry<DerivedDriver, RequestType, BatchEntryKeyMember...>;
     static_assert(std::is_base_of_v<Self, DerivedEntry>);
@@ -248,7 +248,7 @@ class BatchReadEntry : public Batch::Impl::Entry {
 //
 // This may be useful for setting error results.
 template <typename Request>
-void SetCommonResult(span<const Request> requests,
+void SetCommonResult(tensorstore::span<const Request> requests,
                      Result<kvstore::ReadResult>&& result) {
   if (requests.empty()) return;
   for (size_t i = 1; i < requests.size(); ++i) {
@@ -263,12 +263,13 @@ void SetCommonResult(span<const Request> requests,
 template <typename Requests>
 void SetCommonResult(const Requests& requests,
                      Result<kvstore::ReadResult>&& result) {
-  SetCommonResult(span<const typename Requests::value_type>(requests),
-                  std::move(result));
+  SetCommonResult(
+      tensorstore::span<const typename Requests::value_type>(requests),
+      std::move(result));
 }
 
 template <typename Request>
-void SortRequestsByStartByte(span<Request> requests) {
+void SortRequestsByStartByte(tensorstore::span<Request> requests) {
   std::sort(requests.begin(), requests.end(),
             [](const Request& a, const Request& b) {
               return a.byte_range.inclusive_min < b.byte_range.inclusive_min;
@@ -278,7 +279,7 @@ void SortRequestsByStartByte(span<Request> requests) {
 // Resolves coalesced requests with the appropriate cord subranges.
 template <typename Request>
 void ResolveCoalescedRequests(ByteRange coalesced_byte_range,
-                              span<Request> coalesced_requests,
+                              tensorstore::span<Request> coalesced_requests,
                               kvstore::ReadResult&& read_result) {
   for (auto& request : coalesced_requests) {
     kvstore::ReadResult sub_read_result;
@@ -311,8 +312,8 @@ void ResolveCoalescedRequests(ByteRange coalesced_byte_range,
 //     coalesced_byte_range, span<Request> coalesced_requests)` to be invoked
 //     for each coalesced set of requests.
 template <typename Request, typename Predicate, typename Callback>
-void ForEachCoalescedRequest(span<Request> requests, Predicate predicate,
-                             Callback callback) {
+void ForEachCoalescedRequest(tensorstore::span<Request> requests,
+                             Predicate predicate, Callback callback) {
   static_assert(IsByteRangeReadRequestLikeV<Request>);
 
   SortRequestsByStartByte(requests);

tensorstore/kvstore/gcs_grpc/BUILD

@@ -94,6 +94,7 @@ tensorstore_cc_test(
     deps = [
         ":gcs_grpc",
         ":mock_storage_service",
+        "//tensorstore:batch",
         "//tensorstore:context",
         "//tensorstore/internal:flat_cord_builder",
         "//tensorstore/internal/grpc:grpc_mock",

tensorstore/kvstore/gcs_grpc/gcs_grpc.cc

@@ -96,7 +96,6 @@
 
 // protos
 #include "google/protobuf/empty.pb.h"
-#include "google/storage/v2/storage.grpc.pb.h"
 #include "google/storage/v2/storage.pb.h"
 
 namespace jb = tensorstore::internal_json_binding;
@@ -114,7 +113,6 @@ using ::tensorstore::kvstore::ListEntry;
 using ::tensorstore::kvstore::ListReceiver;
 using ::tensorstore::kvstore::SupportedFeatures;
 
-using ::google::storage::v2::Storage;
 
 namespace tensorstore {
 namespace internal_gcs_grpc {