feat(common): add StartWithCorkedInitialMetadata.

google/cloud/async_streaming_read_write_rpc.h

@@ -51,14 +51,36 @@ class AsyncStreamingReadWriteRpc {
   /**
    * Start the streaming RPC.
    *
-   * Applications should call `Start()` and wait for its result before calling
-   * `Read()` and/or `Write()`. If `Start()` completes with `false` the stream
-   * has completed with an error.  The application should not call `Read()` or
-   * `Write()` in this case. On errors, the application should call`Finish()` to
-   * determine the status of the streaming RPC.
+   * Applications should call one of `Start()` or
+   * `StartWithCorkedInitialMetadata()` and wait for its result before calling
+   * `Read()` and/or `Write()`.
+   *
+   * If `Start()` completes with `false` the stream has completed with an error.
+   * The application should not call `Read()` or `Write()` in this case. On
+   * errors, the application should call `Finish()` to determine the status of
+   * the streaming RPC.
    */
   virtual future<bool> Start() = 0;
 
+  /**
+   * Start the streaming RPC with corked initial metadata.
+   *
+   * Applications should call one of `Start()` or
+   * `StartWithCorkedInitialMetadata()` and wait for its result before calling
+   * `Read()` and/or `Write()`.
+   *
+   * If `StartWithCorkedInitialMetadata()` completes with `false` the stream has
+   * completed with an error. The application should not call `Read()` or
+   * `Write()` in this case. On errors, the application should call `Finish()`
+   * to determine the status of the streaming RPC.
+   *
+   * This method differs from `Start()` by not sending the initial metadata
+   * until the first `Write()` call on the stream. That saves a bounce through
+   * the gRPC CompletionQueue at stream startup if the caller does not need the
+   * initial metadata in order to send the first `Write()`.
+   */
+  virtual future<bool> StartWithCorkedInitialMetadata() = 0;
+
   /**
    * Read one response from the streaming RPC.
    *

google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc

@@ -87,6 +87,9 @@ class AsyncStreamingReadWriteRpcTracking
 
   void Cancel() override { child_->Cancel(); }
   future<bool> Start() override { return child_->Start(); }
+  future<bool> StartWithCorkedInitialMetadata() override {
+    return child_->StartWithCorkedInitialMetadata();
+  }
   future<std::optional<Response>> Read() override { return child_->Read(); }
   future<bool> Write(Request const& r, grpc::WriteOptions o) override {
     return child_->Write(r, std::move(o));

google/cloud/internal/async_read_write_stream_auth.h

@@ -46,13 +46,11 @@ class AsyncStreamingReadWriteRpcAuth
   void Cancel() override { state_->Cancel(); }
 
   future<bool> Start() override {
-    using Result = StatusOr<std::shared_ptr<grpc::ClientContext>>;
-    auto weak = std::weak_ptr<SharedState>(state_);
-    return auth_->AsyncConfigureContext(state_->ReleaseInitialContext())
-        .then([weak](future<Result> f) mutable {
-          if (auto state = weak.lock()) return state->OnStart(f.get());
-          return make_ready_future(false);
-        });
+    return StartInternal(/*initial_metadata_corked=*/false);
+  }
+
+  future<bool> StartWithCorkedInitialMetadata() override {
+    return StartInternal(/*initial_metadata_corked=*/true);
   }
 
   future<std::optional<Response>> Read() override {
@@ -93,8 +91,8 @@ class AsyncStreamingReadWriteRpcAuth
       return std::move(initial_context);
     }
 
-    future<bool> OnStart(
-        StatusOr<std::shared_ptr<grpc::ClientContext>> context) {
+    future<bool> OnStart(StatusOr<std::shared_ptr<grpc::ClientContext>> context,
+                         bool initial_metadata_corked) {
       std::lock_guard<std::mutex> g{mu};
       if (cancelled) return make_ready_future(false);
       if (context) {
@@ -104,7 +102,8 @@ class AsyncStreamingReadWriteRpcAuth
             AsyncStreamingReadWriteRpcError<Request, Response>>(
             std::move(context).status());
       }
-      return stream->Start();
+      return initial_metadata_corked ? stream->StartWithCorkedInitialMetadata()
+                                     : stream->Start();
     }
 
     future<Status> Finish() {
@@ -135,6 +134,18 @@ class AsyncStreamingReadWriteRpcAuth
     bool cancelled = false;  // ABSL_GUARDED_BY(mu)
   };
 
+  future<bool> StartInternal(bool initial_metadata_corked) {
+    using Result = StatusOr<std::shared_ptr<grpc::ClientContext>>;
+    auto weak = std::weak_ptr<SharedState>(state_);
+    return auth_->AsyncConfigureContext(state_->ReleaseInitialContext())
+        .then([weak, initial_metadata_corked](future<Result> f) mutable {
+          if (auto state = weak.lock()) {
+            return state->OnStart(f.get(), initial_metadata_corked);
+          }
+          return make_ready_future(false);
+        });
+  }
+
   std::shared_ptr<GrpcAuthenticationStrategy> const auth_;
   std::shared_ptr<SharedState> const state_;
 };

google/cloud/internal/async_read_write_stream_auth_test.cc

@@ -92,6 +92,23 @@ TEST(AsyncStreamReadWriteAuth, Start) {
               UnorderedElementsAre(Pair("tk0", "v0"), Pair("tk1", "v1")));
 }
 
+TEST(AsyncStreamReadWriteAuth, StartWithCorkedInitialMetadata) {
+  auto factory = AuthStream::StreamFactory([](auto) {
+    auto mock = std::make_unique<StrictMock<MockStream>>();
+    EXPECT_CALL(*mock, StartWithCorkedInitialMetadata).WillOnce([] {
+      return make_ready_future(true);
+    });
+    return std::unique_ptr<BaseStream>(std::move(mock));
+  });
+  auto strategy = std::make_shared<StrictMock<MockAuthenticationStrategy>>();
+  EXPECT_CALL(*strategy, AsyncConfigureContext).WillOnce([](auto context) {
+    return make_ready_future(make_status_or(std::move(context)));
+  });
+  auto uut = std::make_unique<AuthStream>(
+      std::make_shared<grpc::ClientContext>(), strategy, factory);
+  EXPECT_TRUE(uut->StartWithCorkedInitialMetadata().get());
+}
+
 }  // namespace
 }  // namespace internal
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

google/cloud/internal/async_read_write_stream_impl.h

@@ -76,6 +76,12 @@ class AsyncStreamingReadWriteRpcImpl
     return op->p.get_future();
   }
 
+  future<bool> StartWithCorkedInitialMetadata() override {
+    context_->set_initial_metadata_corked(true);
+    stream_->StartCall(nullptr);
+    return make_ready_future(true);
+  }
+
   future<std::optional<Response>> Read() override {
     struct OnRead : public AsyncGrpcOperation {
       explicit OnRead(ImmutableOptions o) : call_context(std::move(o)) {}
@@ -223,6 +229,9 @@ class AsyncStreamingReadWriteRpcError
 
   void Cancel() override {}
   future<bool> Start() override { return make_ready_future(false); }
+  future<bool> StartWithCorkedInitialMetadata() override {
+    return make_ready_future(false);
+  }
   future<std::optional<Response>> Read() override {
     return make_ready_future<std::optional<Response>>(std::nullopt);
   }

google/cloud/internal/async_read_write_stream_impl_test.cc

@@ -199,12 +199,39 @@ TEST(AsyncReadWriteStreamingRpcTest, Basic) {
   EXPECT_THAT(finish.get(), IsOk());
 }
 
+TEST(AsyncReadWriteStreamingRpcTest, StartWithCorkedInitialMetadata) {
+  MockStub mock;
+  EXPECT_CALL(mock, FakeRpc)
+      .WillOnce([](grpc::ClientContext*, grpc::CompletionQueue*) {
+        auto stream = std::make_unique<MockReaderWriter>();
+        EXPECT_CALL(*stream, StartCall(nullptr)).Times(1);
+        return stream;
+      });
+
+  auto mock_cq = std::make_shared<MockCompletionQueueImpl>();
+  EXPECT_CALL(*mock_cq, cq).WillRepeatedly(Return(nullptr));
+  EXPECT_CALL(*mock_cq, StartOperation).Times(0);
+
+  google::cloud::CompletionQueue cq(mock_cq);
+  auto stream = MakeStreamingReadWriteRpc<FakeRequest, FakeResponse>(
+      cq, std::make_shared<grpc::ClientContext>(), MakeImmutableOptions({}),
+      [&mock](grpc::ClientContext* context, grpc::CompletionQueue* cq) {
+        return mock.FakeRpc(context, cq);
+      });
+
+  auto start = stream->StartWithCorkedInitialMetadata();
+  // If metadata is corked, the start future should be ready immediately.
+  EXPECT_TRUE(start.is_ready());
+  EXPECT_TRUE(start.get());
+}
+
 TEST(AsyncReadWriteStreamingRpcTest, Error) {
   auto stream = AsyncStreamingReadWriteRpcError<FakeRequest, FakeResponse>(
       Status(StatusCode::kPermissionDenied, "uh-oh"));
 
   stream.Cancel();
   EXPECT_FALSE(stream.Start().get());
+  EXPECT_FALSE(stream.StartWithCorkedInitialMetadata().get());
   EXPECT_FALSE(stream.Read().get().has_value());
   EXPECT_FALSE(stream.Write(FakeRequest{}, grpc::WriteOptions()).get());
   EXPECT_FALSE(stream.WritesDone().get());

google/cloud/internal/async_read_write_stream_logging.h

@@ -51,6 +51,17 @@ class AsyncStreamingReadWriteRpcLogging
     });
   }
 
+  future<bool> StartWithCorkedInitialMetadata() override {
+    auto prefix = std::string(__func__) + "(" + request_id_ + ")";
+    GCP_LOG(DEBUG) << prefix << " <<";
+    return child_->StartWithCorkedInitialMetadata().then(
+        [prefix](future<bool> f) {
+          auto r = f.get();
+          GCP_LOG(DEBUG) << prefix << " >> " << (r ? "true" : "false");
+          return r;
+        });
+  }
+
   future<std::optional<Response>> Read() override {
     auto prefix = std::string(__func__) + "(" + request_id_ + ")";
     auto const& opt = tracing_options_;

google/cloud/internal/async_read_write_stream_logging_test.cc

@@ -64,6 +64,18 @@ TEST_F(StreamingReadRpcLoggingTest, Start) {
               Contains(AllOf(HasSubstr("Start"), HasSubstr("test-id"))));
 }
 
+TEST_F(StreamingReadRpcLoggingTest, StartWithCorkedInitialMetadata) {
+  auto mock = std::make_unique<MockStream>();
+  EXPECT_CALL(*mock, StartWithCorkedInitialMetadata).WillOnce([] {
+    return make_ready_future(true);
+  });
+  TestedStream stream(std::move(mock), TracingOptions{}, "test-id");
+  EXPECT_TRUE(stream.StartWithCorkedInitialMetadata().get());
+  EXPECT_THAT(log_.ExtractLines(),
+              Contains(AllOf(HasSubstr("StartWithCorkedInitialMetadata"),
+                             HasSubstr("test-id"))));
+}
+
 TEST_F(StreamingReadRpcLoggingTest, ReadWithValue) {
   auto mock = std::make_unique<MockStream>();
   EXPECT_CALL(*mock, Read).WillOnce([] {

google/cloud/internal/async_read_write_stream_timeout.h

@@ -66,6 +66,10 @@ class AsyncStreamingReadWriteRpcTimeout
 
   future<bool> Start() override { return state_->Start(); }
 
+  future<bool> StartWithCorkedInitialMetadata() override {
+    return state_->StartWithCorkedInitialMetadata();
+  }
+
   future<std::optional<Response>> Read() override { return state_->Read(); }
 
   future<bool> Write(Request const& request,
@@ -109,6 +113,17 @@ class AsyncStreamingReadWriteRpcTimeout
           });
     }
 
+    future<bool> StartWithCorkedInitialMetadata() {
+      auto watchdog = CreateWatchdog(start_timeout);
+      return child->StartWithCorkedInitialMetadata().then(
+          [watchdog = std::move(watchdog), w = WeakFromThis()](auto f) mutable {
+            if (auto self = w.lock()) {
+              return self->OnStart(std::move(watchdog), f.get());
+            }
+            return make_ready_future(false);
+          });
+    }
+
     future<bool> OnStart(future<bool> watchdog, bool start) {
       watchdog.cancel();
       return watchdog.then([w = WeakFromThis(), start](auto f) {

google/cloud/internal/async_read_write_stream_timeout_test.cc

@@ -126,6 +126,64 @@ TEST(AsyncStreamingReadWriteRpcTimeout, StartTimeout) {
   EXPECT_FALSE(result.get());
 }
 
+TEST(AsyncStreamingReadWriteRpcTimeout, StartWithCorkedInitialMetadataSuccess) {
+  AsyncSequencer<bool> sequencer;
+  auto mock = std::make_unique<MockStream>();
+  EXPECT_CALL(*mock, StartWithCorkedInitialMetadata).WillOnce([&] {
+    return sequencer.PushBack("StartWithCorkedInitialMetadata");
+  });
+  EXPECT_CALL(*mock, Cancel).Times(0);
+
+  auto mock_cq = std::make_shared<MockCompletionQueueImpl>();
+  EXPECT_CALL(*mock_cq,
+              MakeRelativeTimer(std::chrono::nanoseconds(kStartTimeout)))
+      .WillOnce([&](auto) {
+        return sequencer.PushBack("MakeRelativeTimer").then(MakeTimerStatus);
+      });
+
+  CompletionQueue cq(mock_cq);
+
+  TestedStream uut(cq, kStartTimeout, kReadTimeout, kWriteTimeout,
+                   std::move(mock));
+  auto result = uut.StartWithCorkedInitialMetadata();
+  auto timer = sequencer.PopFrontWithName();
+  auto start = sequencer.PopFrontWithName();
+  ASSERT_THAT(timer.second, "MakeRelativeTimer");
+  ASSERT_THAT(start.second, "StartWithCorkedInitialMetadata");
+  start.first.set_value(true);
+  timer.first.set_value(false);
+  EXPECT_TRUE(result.get());
+}
+
+TEST(AsyncStreamingReadWriteRpcTimeout, StartWithCorkedInitialMetadataTimeout) {
+  AsyncSequencer<bool> sequencer;
+  auto mock = std::make_unique<MockStream>();
+  EXPECT_CALL(*mock, StartWithCorkedInitialMetadata).WillOnce([&] {
+    return sequencer.PushBack("StartWithCorkedInitialMetadata");
+  });
+  EXPECT_CALL(*mock, Cancel).Times(1);
+
+  auto mock_cq = std::make_shared<MockCompletionQueueImpl>();
+  EXPECT_CALL(*mock_cq,
+              MakeRelativeTimer(std::chrono::nanoseconds(kStartTimeout)))
+      .WillOnce([&](auto) {
+        return sequencer.PushBack("MakeRelativeTimer").then(MakeTimerStatus);
+      });
+
+  CompletionQueue cq(mock_cq);
+
+  TestedStream uut(cq, kStartTimeout, kReadTimeout, kWriteTimeout,
+                   std::move(mock));
+  auto result = uut.StartWithCorkedInitialMetadata();
+  auto timer = sequencer.PopFrontWithName();
+  auto start = sequencer.PopFrontWithName();
+  ASSERT_THAT(timer.second, "MakeRelativeTimer");
+  ASSERT_THAT(start.second, "StartWithCorkedInitialMetadata");
+  timer.first.set_value(true);
+  start.first.set_value(true);
+  EXPECT_FALSE(result.get());
+}
+
 TEST(AsyncStreamingReadWriteRpcTimeout, ReadSuccess) {
   AsyncSequencer<bool> sequencer;
   auto mock = std::make_unique<MockStream>();

google/cloud/internal/async_read_write_stream_tracing.h

@@ -43,19 +43,14 @@ class AsyncStreamingReadWriteRpcTracing
   }
 
   future<bool> Start() override {
-    // It is sufficient to set `span_` as the parent of `start_span`, because
-    // the lower levels do not create any spans.
-    opentelemetry::trace::StartSpanOptions options;
-    options.parent = span_->GetContext();
-    auto start_span = internal::MakeSpan("Start", options);
-    return impl_->Start().then(
-        [this, ss = std::move(start_span)](future<bool> f) {
-          EndSpan(*ss);
-          auto started = f.get();
-          span_->SetAttribute("gl-cpp.stream_started", started);
-          started_ = started;
-          return started;
-        });
+    return StartInternal<&AsyncStreamingReadWriteRpc<Request, Response>::Start>(
+        "Start");
+  }
+
+  future<bool> StartWithCorkedInitialMetadata() override {
+    return StartInternal<&AsyncStreamingReadWriteRpc<
+        Request, Response>::StartWithCorkedInitialMetadata>(
+        "StartWithCorkedInitialMetadata");
   }
 
   future<std::optional<Response>> Read() override {
@@ -119,6 +114,23 @@ class AsyncStreamingReadWriteRpcTracing
     return EndSpan(*std::move(span_), std::move(status));
   }
 
+  template <auto StartFunc>
+  future<bool> StartInternal(std::string_view span_name) {
+    // It is sufficient to set `span_` as the parent of `start_span`, because
+    // the lower levels do not create any spans.
+    opentelemetry::trace::StartSpanOptions options;
+    options.parent = span_->GetContext();
+    auto start_span = internal::MakeSpan(span_name, options);
+    return std::invoke(StartFunc, *impl_)
+        .then([this, ss = std::move(start_span)](future<bool> f) {
+          EndSpan(*ss);
+          auto started = f.get();
+          span_->SetAttribute("gl-cpp.stream_started", started);
+          started_ = started;
+          return started;
+        });
+  }
+
   std::shared_ptr<grpc::ClientContext> context_;
   std::unique_ptr<AsyncStreamingReadWriteRpc<Request, Response>> impl_;
   opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> span_;