feat(common): GUAC for async read-write streams

google/cloud/CMakeLists.txt

@@ -402,6 +402,7 @@ if (GOOGLE_CLOUD_CPP_ENABLE_GRPC)
         internal/async_polling_loop.cc
         internal/async_polling_loop.h
         internal/async_read_stream_impl.h
+        internal/async_read_write_stream_auth.h
         internal/async_read_write_stream_impl.h
         internal/async_retry_loop.h
         internal/async_retry_unary_rpc.h
@@ -541,6 +542,7 @@ if (GOOGLE_CLOUD_CPP_ENABLE_GRPC)
             internal/async_connection_ready_test.cc
             internal/async_long_running_operation_test.cc
             internal/async_polling_loop_test.cc
+            internal/async_read_write_stream_auth_test.cc
             internal/async_read_write_stream_impl_test.cc
             internal/async_retry_loop_test.cc
             internal/async_retry_unary_rpc_test.cc

google/cloud/google_cloud_cpp_grpc_utils.bzl

@@ -32,6 +32,7 @@ google_cloud_cpp_grpc_utils_hdrs = [
     "internal/async_long_running_operation.h",
     "internal/async_polling_loop.h",
     "internal/async_read_stream_impl.h",
+    "internal/async_read_write_stream_auth.h",
     "internal/async_read_write_stream_impl.h",
     "internal/async_retry_loop.h",
     "internal/async_retry_unary_rpc.h",

google/cloud/google_cloud_cpp_grpc_utils_unit_tests.bzl

@@ -24,6 +24,7 @@ google_cloud_cpp_grpc_utils_unit_tests = [
     "internal/async_connection_ready_test.cc",
     "internal/async_long_running_operation_test.cc",
     "internal/async_polling_loop_test.cc",
+    "internal/async_read_write_stream_auth_test.cc",
     "internal/async_read_write_stream_impl_test.cc",
     "internal/async_retry_loop_test.cc",
     "internal/async_retry_unary_rpc_test.cc",

google/cloud/internal/async_read_write_stream_auth.h

@@ -0,0 +1,111 @@
+// Copyright 2021 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_INTERNAL_ASYNC_READ_WRITE_STREAM_AUTH_H
+#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_INTERNAL_ASYNC_READ_WRITE_STREAM_AUTH_H
+
+#include "google/cloud/internal/async_read_write_stream_impl.h"
+#include "google/cloud/internal/unified_grpc_credentials.h"
+#include "google/cloud/version.h"
+#include <functional>
+#include <memory>
+
+namespace google {
+namespace cloud {
+inline namespace GOOGLE_CLOUD_CPP_NS {
+namespace internal {
+
+template <typename Request, typename Response>
+class AsyncStreamingReadWriteRpcAuth
+    : public AsyncStreamingReadWriteRpc<Request, Response>,
+      public std::enable_shared_from_this<
+          AsyncStreamingReadWriteRpcAuth<Request, Response>> {
+ public:
+  using StreamFactory = std::function<
+      std::shared_ptr<AsyncStreamingReadWriteRpc<Request, Response>>(
+          std::unique_ptr<grpc::ClientContext>)>;
+
+  AsyncStreamingReadWriteRpcAuth(
+      std::unique_ptr<grpc::ClientContext> context,
+      std::shared_ptr<GrpcAuthenticationStrategy> auth, StreamFactory factory)
+      : context_(std::move(context)),
+        auth_(std::move(auth)),
+        factory_(std::move(factory)) {}
+
+  void Cancel() override {
+    if (context_) return context_->TryCancel();
+    if (stream_) return stream_->Cancel();
+  }
+
+  future<bool> Start() override {
+    using Result = StatusOr<std::unique_ptr<grpc::ClientContext>>;
+
+    auto weak =
+        std::weak_ptr<AsyncStreamingReadWriteRpcAuth>(this->shared_from_this());
+    return auth_->AsyncConfigureContext(std::move(context_))
+        .then([weak](future<Result> f) mutable {
+          if (auto self = weak.lock()) return self->OnStart(f.get());
+          return make_ready_future(false);
+        });
+  }
+
+  future<absl::optional<Response>> Read() override {
+    if (!stream_) return make_ready_future(absl::optional<Response>{});
+    return stream_->Read();
+  }
+
+  future<bool> Write(Request const& request,
+                     grpc::WriteOptions options) override {
+    if (!stream_) return make_ready_future(false);
+    return stream_->Write(request, std::move(options));
+  }
+
+  future<bool> WritesDone() override {
+    if (!stream_) return make_ready_future(false);
+    return stream_->WritesDone();
+  }
+
+  future<Status> Finish() override {
+    if (!stream_) {
+      return make_ready_future(
+          Status(StatusCode::kInvalidArgument,
+                 "uninitialized GrpcReadWriteStreamAuth<>"));
+    }
+    return stream_->Finish();
+  }
+
+ private:
+  future<bool> OnStart(StatusOr<std::unique_ptr<grpc::ClientContext>> context) {
+    if (!context) {
+      stream_ =
+          absl::make_unique<AsyncStreamingReadWriteRpcError<Request, Response>>(
+              std::move(context).status());
+      return make_ready_future(false);
+    }
+    stream_ = factory_(*std::move(context));
+    return stream_->Start();
+  }
+
+  std::unique_ptr<grpc::ClientContext> context_;
+  std::shared_ptr<GrpcAuthenticationStrategy> auth_;
+  StreamFactory factory_;
+  std::shared_ptr<AsyncStreamingReadWriteRpc<Request, Response>> stream_;
+};
+
+}  // namespace internal
+}  // namespace GOOGLE_CLOUD_CPP_NS
+}  // namespace cloud
+}  // namespace google
+
+#endif  // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_INTERNAL_ASYNC_READ_WRITE_STREAM_AUTH_H

google/cloud/internal/async_read_write_stream_auth_test.cc

@@ -0,0 +1,96 @@
+// Copyright 2021 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "google/cloud/internal/async_read_write_stream_auth.h"
+#include "google/cloud/completion_queue.h"
+#include "google/cloud/testing_util/mock_grpc_authentication_strategy.h"
+#include "google/cloud/testing_util/status_matchers.h"
+#include "absl/memory/memory.h"
+#include <gmock/gmock.h>
+#include <memory>
+#include <string>
+
+namespace google {
+namespace cloud {
+inline namespace GOOGLE_CLOUD_CPP_NS {
+namespace internal {
+namespace {
+
+using ::google::cloud::testing_util::IsOk;
+using ::google::cloud::testing_util::MockAuthenticationStrategy;
+
+struct FakeRequest {
+  std::string key;
+};
+
+struct FakeResponse {
+  std::string key;
+  std::string value;
+};
+
+using BaseStream = AsyncStreamingReadWriteRpc<FakeRequest, FakeResponse>;
+using AuthStream = AsyncStreamingReadWriteRpcAuth<FakeRequest, FakeResponse>;
+
+class MockStream : public BaseStream {
+ public:
+  MOCK_METHOD(void, Cancel, (), (override));
+  MOCK_METHOD(future<bool>, Start, (), (override));
+  MOCK_METHOD(future<absl::optional<FakeResponse>>, Read, (), (override));
+  MOCK_METHOD(future<bool>, Write, (FakeRequest const&, grpc::WriteOptions),
+              (override));
+  MOCK_METHOD(future<bool>, WritesDone, (), (override));
+  MOCK_METHOD(future<Status>, Finish, (), (override));
+};
+
+TEST(AsyncStreamReadWriteAuth, Start) {
+  auto factory = [](std::unique_ptr<grpc::ClientContext>) {
+    auto mock = absl::make_unique<MockStream>();
+    EXPECT_CALL(*mock, Start).WillOnce([] { return make_ready_future(true); });
+    EXPECT_CALL(*mock, Write)
+        .WillOnce([](FakeRequest const&, grpc::WriteOptions) {
+          return make_ready_future(true);
+        });
+    EXPECT_CALL(*mock, Read).WillOnce([] {
+      return make_ready_future(absl::make_optional(FakeResponse{"k0", "v0"}));
+    });
+    EXPECT_CALL(*mock, WritesDone).WillOnce([] {
+      return make_ready_future(true);
+    });
+    EXPECT_CALL(*mock, Finish).WillOnce([] {
+      return make_ready_future(Status{});
+    });
+    return std::unique_ptr<BaseStream>(std::move(mock));
+  };
+  auto strategy = std::make_shared<MockAuthenticationStrategy>();
+  EXPECT_CALL(*strategy, AsyncConfigureContext)
+      .WillOnce([](std::unique_ptr<grpc::ClientContext> context) {
+        return make_ready_future(make_status_or(std::move(context)));
+      });
+  auto uut = std::make_shared<AuthStream>(
+      absl::make_unique<grpc::ClientContext>(), strategy, factory);
+  EXPECT_TRUE(uut->Start().get());
+  EXPECT_TRUE(uut->Write(FakeRequest{"k"}, grpc::WriteOptions()).get());
+  auto response = uut->Read().get();
+  ASSERT_TRUE(response.has_value());
+  EXPECT_EQ(response->key, "k0");
+  EXPECT_EQ(response->value, "v0");
+  EXPECT_TRUE(uut->WritesDone().get());
+  EXPECT_THAT(uut->Finish().get(), IsOk());
+}
+
+}  // namespace
+}  // namespace internal
+}  // namespace GOOGLE_CLOUD_CPP_NS
+}  // namespace cloud
+}  // namespace google

google/cloud/pubsub/internal/streaming_subscription_batch_source_test.cc

@@ -325,7 +325,7 @@ TEST(StreamingSubscriptionBatchSourceTest, StartUnexpected) {
       .WillRepeatedly([](google::cloud::CompletionQueue&,
                          std::unique_ptr<grpc::ClientContext>,
                          google::pubsub::v1::StreamingPullRequest const&) {
-        return std::unique_ptr<SubscriberStub::AsyncPullStream>{};
+        return std::shared_ptr<SubscriberStub::AsyncPullStream>{};
       });
 
   auto shutdown = std::make_shared<SessionShutdownManager>();

google/cloud/pubsub/internal/subscriber_auth.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "google/cloud/pubsub/internal/subscriber_auth.h"
+#include "google/cloud/internal/async_read_write_stream_auth.h"
 
 namespace google {
 namespace cloud {
@@ -68,19 +69,23 @@ Status SubscriberAuth::ModifyPushConfig(
   return child_->ModifyPushConfig(context, request);
 }
 
-std::unique_ptr<SubscriberStub::AsyncPullStream>
+std::shared_ptr<SubscriberStub::AsyncPullStream>
 SubscriberAuth::AsyncStreamingPull(
     google::cloud::CompletionQueue& cq,
     std::unique_ptr<grpc::ClientContext> context,
     google::pubsub::v1::StreamingPullRequest const& request) {
-  using ErrorStream =
-      ::google::cloud::internal::AsyncStreamingReadWriteRpcError<
-          google::pubsub::v1::StreamingPullRequest,
-          google::pubsub::v1::StreamingPullResponse>;
-
-  auto status = auth_->ConfigureContext(*context);
-  if (!status.ok()) return absl::make_unique<ErrorStream>(std::move(status));
-  return child_->AsyncStreamingPull(cq, std::move(context), request);
+  using StreamAuth = google::cloud::internal::AsyncStreamingReadWriteRpcAuth<
+      google::pubsub::v1::StreamingPullRequest,
+      google::pubsub::v1::StreamingPullResponse>;
+
+  auto child = child_;
+  auto call = [child, cq,
+               request](std::unique_ptr<grpc::ClientContext> ctx) mutable {
+    return child->AsyncStreamingPull(cq, std::move(ctx), request);
+  };
+  auto factory = StreamAuth::StreamFactory(std::move(call));
+  return std::make_shared<StreamAuth>(std::move(context), auth_,
+                                      std::move(factory));
 }
 
 future<Status> SubscriberAuth::AsyncAcknowledge(

google/cloud/pubsub/internal/subscriber_auth.h

@@ -56,7 +56,7 @@ class SubscriberAuth : public SubscriberStub {
       grpc::ClientContext& context,
       google::pubsub::v1::ModifyPushConfigRequest const& request) override;
 
-  std::unique_ptr<AsyncPullStream> AsyncStreamingPull(
+  std::shared_ptr<AsyncPullStream> AsyncStreamingPull(
       google::cloud::CompletionQueue& cq,
       std::unique_ptr<grpc::ClientContext> context,
       google::pubsub::v1::StreamingPullRequest const& request) override;

google/cloud/pubsub/internal/subscriber_auth_test.cc

@@ -140,16 +140,18 @@ TEST(SubscriberAuthTest, AsyncStreamingPull) {
         return absl::make_unique<ErrorStream>(
             Status(StatusCode::kPermissionDenied, "uh-oh"));
       });
-  auto under_test = SubscriberAuth(MakeTypicalMockAuth(), mock);
+  auto under_test = SubscriberAuth(MakeTypicalAsyncMockAuth(), mock);
   google::cloud::CompletionQueue cq;
   google::pubsub::v1::StreamingPullRequest request;
   auto auth_failure = under_test.AsyncStreamingPull(
       cq, absl::make_unique<grpc::ClientContext>(), request);
+  ASSERT_FALSE(auth_failure->Start().get());
   EXPECT_THAT(auth_failure->Finish().get(),
               StatusIs(StatusCode::kInvalidArgument));
 
   auto auth_success = under_test.AsyncStreamingPull(
       cq, absl::make_unique<grpc::ClientContext>(), request);
+  ASSERT_FALSE(auth_success->Start().get());
   EXPECT_THAT(auth_success->Finish().get(),
               StatusIs(StatusCode::kPermissionDenied));
 }

google/cloud/pubsub/internal/subscriber_logging.cc

@@ -92,7 +92,7 @@ Status SubscriberLogging::ModifyPushConfig(
       context, request, __func__, tracing_options_);
 }
 
-std::unique_ptr<SubscriberStub::AsyncPullStream>
+std::shared_ptr<SubscriberStub::AsyncPullStream>
 SubscriberLogging::AsyncStreamingPull(
     google::cloud::CompletionQueue& cq,
     std::unique_ptr<grpc::ClientContext> context,
@@ -103,8 +103,8 @@ SubscriberLogging::AsyncStreamingPull(
                  << " << request=" << DebugString(request, tracing_options_);
   auto stream = child_->AsyncStreamingPull(cq, std::move(context), request);
   if (!trace_streams_) return stream;
-  return absl::make_unique<LoggingAsyncPullStream>(
-      std::move(stream), tracing_options_, request_id);
+  return std::make_shared<LoggingAsyncPullStream>(std::move(stream),
+                                                  tracing_options_, request_id);
 }
 
 future<Status> SubscriberLogging::AsyncAcknowledge(
@@ -201,7 +201,7 @@ StatusOr<google::pubsub::v1::SeekResponse> SubscriberLogging::Seek(
 }
 
 LoggingAsyncPullStream::LoggingAsyncPullStream(
-    std::unique_ptr<SubscriberStub::AsyncPullStream> child,
+    std::shared_ptr<SubscriberStub::AsyncPullStream> child,
     TracingOptions tracing_options, std::string request_id)
     : child_(std::move(child)),
       tracing_options_(std::move(tracing_options)),

google/cloud/pubsub/internal/subscriber_logging.h

@@ -58,7 +58,7 @@ class SubscriberLogging : public SubscriberStub {
       grpc::ClientContext& context,
       google::pubsub::v1::ModifyPushConfigRequest const& request) override;
 
-  std::unique_ptr<AsyncPullStream> AsyncStreamingPull(
+  std::shared_ptr<AsyncPullStream> AsyncStreamingPull(
       google::cloud::CompletionQueue& cq,
       std::unique_ptr<grpc::ClientContext> context,
       google::pubsub::v1::StreamingPullRequest const& request) override;
@@ -105,7 +105,7 @@ class SubscriberLogging : public SubscriberStub {
 
 class LoggingAsyncPullStream : public SubscriberStub::AsyncPullStream {
  public:
-  LoggingAsyncPullStream(std::unique_ptr<SubscriberStub::AsyncPullStream> child,
+  LoggingAsyncPullStream(std::shared_ptr<SubscriberStub::AsyncPullStream> child,
                          TracingOptions tracing_options,
                          std::string request_id);
 
@@ -119,7 +119,7 @@ class LoggingAsyncPullStream : public SubscriberStub::AsyncPullStream {
   future<Status> Finish() override;
 
  private:
-  std::unique_ptr<SubscriberStub::AsyncPullStream> child_;
+  std::shared_ptr<SubscriberStub::AsyncPullStream> child_;
   TracingOptions tracing_options_;
   std::string request_id_;
 };

google/cloud/pubsub/internal/subscriber_metadata.cc

@@ -70,7 +70,7 @@ Status SubscriberMetadata::ModifyPushConfig(
   return child_->ModifyPushConfig(context, request);
 }
 
-std::unique_ptr<SubscriberStub::AsyncPullStream>
+std::shared_ptr<SubscriberStub::AsyncPullStream>
 SubscriberMetadata::AsyncStreamingPull(
     google::cloud::CompletionQueue& cq,
     std::unique_ptr<grpc::ClientContext> context,

google/cloud/pubsub/internal/subscriber_metadata.h

@@ -54,7 +54,7 @@ class SubscriberMetadata : public SubscriberStub {
       grpc::ClientContext& context,
       google::pubsub::v1::ModifyPushConfigRequest const& request) override;
 
-  std::unique_ptr<AsyncPullStream> AsyncStreamingPull(
+  std::shared_ptr<AsyncPullStream> AsyncStreamingPull(
       google::cloud::CompletionQueue& cq,
       std::unique_ptr<grpc::ClientContext> context,
       google::pubsub::v1::StreamingPullRequest const& request) override;