fix(pubsub): Change AsyncReadWriteStreamAuth to be usable with unique_ptr

google/cloud/internal/async_read_write_stream_auth.h

@@ -20,6 +20,7 @@
 #include "google/cloud/version.h"
 #include <functional>
 #include <memory>
+#include <mutex>
 
 namespace google {
 namespace cloud {
@@ -28,79 +29,103 @@ namespace internal {
 
 template <typename Request, typename Response>
 class AsyncStreamingReadWriteRpcAuth
-    : public AsyncStreamingReadWriteRpc<Request, Response>,
-      public std::enable_shared_from_this<
-          AsyncStreamingReadWriteRpcAuth<Request, Response>> {
+    : public AsyncStreamingReadWriteRpc<Request, Response> {
  public:
   using StreamFactory = std::function<
-      std::shared_ptr<AsyncStreamingReadWriteRpc<Request, Response>>(
+      std::unique_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)) {}
+      : auth_(std::move(auth)),
+        state_(std::make_shared<SharedState>(std::move(factory),
+                                             std::move(context))) {}
 
-  void Cancel() override {
-    if (context_) return context_->TryCancel();
-    if (stream_) return stream_->Cancel();
-  }
+  void Cancel() override { state_->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_))
+    auto weak = std::weak_ptr<SharedState>(state_);
+    return auth_->AsyncConfigureContext(state_->ReleaseInitialContext())
         .then([weak](future<Result> f) mutable {
-          if (auto self = weak.lock()) return self->OnStart(f.get());
+          if (auto state = weak.lock()) return state->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();
+    std::lock_guard<std::mutex> g{state_->mu};
+    return state_->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));
+    std::lock_guard<std::mutex> g{state_->mu};
+    return state_->stream->Write(request, std::move(options));
   }
 
   future<bool> WritesDone() override {
-    if (!stream_) return make_ready_future(false);
-    return stream_->WritesDone();
+    std::lock_guard<std::mutex> g{state_->mu};
+    return state_->stream->WritesDone();
   }
 
-  future<Status> Finish() override {
-    if (!stream_) {
-      return make_ready_future(
-          Status(StatusCode::kInvalidArgument,
-                 "uninitialized GrpcReadWriteStreamAuth<>"));
-    }
-    return stream_->Finish();
-  }
+  future<Status> Finish() override { return state_->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);
+  struct SharedState {
+    SharedState(StreamFactory factory,
+                std::unique_ptr<grpc::ClientContext> initial_context)
+        : factory(std::move(factory)),
+          initial_context(std::move(initial_context)),
+          stream(absl::make_unique<
+                 AsyncStreamingReadWriteRpcError<Request, Response>>(
+              Status(StatusCode::kInternal, "Stream is not yet started."))) {}
+
+    std::unique_ptr<grpc::ClientContext> ReleaseInitialContext() {
+      std::lock_guard<std::mutex> g{mu};
+      return std::move(initial_context);
+    }
+
+    future<bool> OnStart(
+        StatusOr<std::unique_ptr<grpc::ClientContext>> context) {
+      std::lock_guard<std::mutex> g{mu};
+      if (cancelled) return make_ready_future(false);
+      if (context) {
+        stream = factory(*std::move(context));
+      } else {
+        stream = absl::make_unique<
+            AsyncStreamingReadWriteRpcError<Request, Response>>(
+            std::move(context).status());
+      }
+      return stream->Start();
+    }
+
+    future<Status> Finish() {
+      std::lock_guard<std::mutex> g{mu};
+      cancelled = true;  // ensure stream is not recreated after Finish
+      return stream->Finish();
+    }
+
+    void Cancel() {
+      std::lock_guard<std::mutex> g{mu};
+      if (cancelled) return;
+      cancelled = true;
+      if (initial_context) initial_context->TryCancel();
+      stream->Cancel();
     }
-    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_;
+    StreamFactory const factory;
+    std::mutex mu;
+    std::unique_ptr<grpc::ClientContext>
+        initial_context;  // ABSL_GUARDED_BY(mu)
+    std::unique_ptr<AsyncStreamingReadWriteRpc<Request, Response>>
+        stream;              // ABSL_GUARDED_BY(mu)
+    bool cancelled = false;  // ABSL_GUARDED_BY(mu)
+  };
+
+  std::shared_ptr<GrpcAuthenticationStrategy> const auth_;
+  std::shared_ptr<SharedState> const state_;
 };
 
 }  // namespace internal

google/cloud/internal/async_read_write_stream_auth_test.cc

@@ -29,6 +29,7 @@ namespace {
 
 using ::google::cloud::testing_util::IsOk;
 using ::google::cloud::testing_util::MockAuthenticationStrategy;
+using ::testing::StrictMock;
 
 struct FakeRequest {
   std::string key;
@@ -54,8 +55,9 @@ class MockStream : public BaseStream {
 };
 
 TEST(AsyncStreamReadWriteAuth, Start) {
-  auto factory = [](std::unique_ptr<grpc::ClientContext>) {
-    auto mock = absl::make_unique<MockStream>();
+  auto factory = AuthStream::StreamFactory([](std::unique_ptr<
+                                               grpc::ClientContext>) {
+    auto mock = absl::make_unique<StrictMock<MockStream>>();
     EXPECT_CALL(*mock, Start).WillOnce([] { return make_ready_future(true); });
     EXPECT_CALL(*mock, Write)
         .WillOnce([](FakeRequest const&, grpc::WriteOptions) {
@@ -71,13 +73,13 @@ TEST(AsyncStreamReadWriteAuth, Start) {
       return make_ready_future(Status{});
     });
     return std::unique_ptr<BaseStream>(std::move(mock));
-  };
-  auto strategy = std::make_shared<MockAuthenticationStrategy>();
+  });
+  auto strategy = std::make_shared<StrictMock<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>(
+  auto uut = absl::make_unique<AuthStream>(
       absl::make_unique<grpc::ClientContext>(), strategy, factory);
   EXPECT_TRUE(uut->Start().get());
   EXPECT_TRUE(uut->Write(FakeRequest{"k"}, grpc::WriteOptions()).get());

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::shared_ptr<SubscriberStub::AsyncPullStream>{};
+        return std::unique_ptr<SubscriberStub::AsyncPullStream>{};
       });
 
   auto shutdown = std::make_shared<SessionShutdownManager>();

google/cloud/pubsub/internal/subscriber_auth.cc

@@ -69,7 +69,7 @@ Status SubscriberAuth::ModifyPushConfig(
   return child_->ModifyPushConfig(context, request);
 }
 
-std::shared_ptr<SubscriberStub::AsyncPullStream>
+std::unique_ptr<SubscriberStub::AsyncPullStream>
 SubscriberAuth::AsyncStreamingPull(
     google::cloud::CompletionQueue& cq,
     std::unique_ptr<grpc::ClientContext> context,
@@ -83,9 +83,8 @@ SubscriberAuth::AsyncStreamingPull(
                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));
+  return absl::make_unique<StreamAuth>(
+      std::move(context), auth_, StreamAuth::StreamFactory(std::move(call)));
 }
 
 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::shared_ptr<AsyncPullStream> AsyncStreamingPull(
+  std::unique_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

@@ -31,6 +31,7 @@ using ::google::cloud::testing_util::StatusIs;
 using ::testing::IsNull;
 using ::testing::Not;
 using ::testing::Return;
+using ::testing::StrictMock;
 
 TEST(SubscriberAuthTest, CreateSubscription) {
   auto mock = std::make_shared<pubsub_testing::MockSubscriberStub>();
@@ -128,31 +129,53 @@ TEST(SubscriberAuthTest, ModifyPushConfig) {
   EXPECT_THAT(auth_success, StatusIs(StatusCode::kPermissionDenied));
 }
 
-TEST(SubscriberAuthTest, AsyncStreamingPull) {
+TEST(SubscriberAuthTest, AsyncStreamingPullFailedAuth) {
+  auto mock =
+      std::make_shared<StrictMock<pubsub_testing::MockSubscriberStub>>();
+  auto auth = std::make_shared<testing_util::MockAuthenticationStrategy>();
+  EXPECT_CALL(*auth, AsyncConfigureContext)
+      .WillOnce([](std::unique_ptr<grpc::ClientContext>)
+                    -> future<StatusOr<std::unique_ptr<grpc::ClientContext>>> {
+        return make_ready_future(StatusOr<std::unique_ptr<grpc::ClientContext>>(
+            Status(StatusCode::kInvalidArgument, "cannot-set-credentials")));
+      });
+  auto under_test = SubscriberAuth(auth, 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));
+}
+
+TEST(SubscriberAuthTest, AsyncStreamingPullAuthSuccess) {
   using ErrorStream =
       ::google::cloud::internal::AsyncStreamingReadWriteRpcError<
           google::pubsub::v1::StreamingPullRequest,
           google::pubsub::v1::StreamingPullResponse>;
 
-  auto mock = std::make_shared<pubsub_testing::MockSubscriberStub>();
+  auto mock =
+      std::make_shared<StrictMock<pubsub_testing::MockSubscriberStub>>();
   EXPECT_CALL(*mock, AsyncStreamingPull)
       .WillOnce([](::testing::Unused, ::testing::Unused, ::testing::Unused) {
         return absl::make_unique<ErrorStream>(
             Status(StatusCode::kPermissionDenied, "uh-oh"));
       });
-  auto under_test = SubscriberAuth(MakeTypicalAsyncMockAuth(), mock);
+  auto auth = std::make_shared<testing_util::MockAuthenticationStrategy>();
+  EXPECT_CALL(*auth, AsyncConfigureContext)
+      .WillOnce([](std::unique_ptr<grpc::ClientContext> context) {
+        context->set_credentials(
+            grpc::AccessTokenCredentials("test-only-invalid"));
+        return make_ready_future(make_status_or(std::move(context)));
+      });
+  auto under_test = SubscriberAuth(auth, 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::shared_ptr<SubscriberStub::AsyncPullStream>
+std::unique_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 std::make_shared<LoggingAsyncPullStream>(std::move(stream),
-                                                  tracing_options_, request_id);
+  return absl::make_unique<LoggingAsyncPullStream>(
+      std::move(stream), tracing_options_, request_id);
 }
 
 future<Status> SubscriberLogging::AsyncAcknowledge(

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::shared_ptr<AsyncPullStream> AsyncStreamingPull(
+  std::unique_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_metadata.cc

@@ -70,7 +70,7 @@ Status SubscriberMetadata::ModifyPushConfig(
   return child_->ModifyPushConfig(context, request);
 }
 
-std::shared_ptr<SubscriberStub::AsyncPullStream>
+std::unique_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::shared_ptr<AsyncPullStream> AsyncStreamingPull(
+  std::unique_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_round_robin.cc

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

google/cloud/pubsub/internal/subscriber_round_robin.h

@@ -56,7 +56,7 @@ class SubscriberRoundRobin : public SubscriberStub {
       grpc::ClientContext& context,
       google::pubsub::v1::ModifyPushConfigRequest const& request) override;
 
-  std::shared_ptr<AsyncPullStream> AsyncStreamingPull(
+  std::unique_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_stub.cc

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

google/cloud/pubsub/internal/subscriber_stub.h

@@ -76,7 +76,7 @@ class SubscriberStub {
       google::pubsub::v1::StreamingPullResponse>;
 
   /// Start a bi-directional stream to read messages and send ack/nacks.
-  virtual std::shared_ptr<AsyncPullStream> AsyncStreamingPull(
+  virtual std::unique_ptr<AsyncPullStream> AsyncStreamingPull(
       google::cloud::CompletionQueue&, std::unique_ptr<grpc::ClientContext>,
       google::pubsub::v1::StreamingPullRequest const& request) = 0;
 

google/cloud/pubsub/testing/mock_subscriber_stub.h

@@ -61,7 +61,7 @@ class MockSubscriberStub : public pubsub_internal::SubscriberStub {
                google::pubsub::v1::ModifyPushConfigRequest const& request),
               (override));
 
-  MOCK_METHOD(std::shared_ptr<pubsub_internal::SubscriberStub::AsyncPullStream>,
+  MOCK_METHOD(std::unique_ptr<pubsub_internal::SubscriberStub::AsyncPullStream>,
               AsyncStreamingPull,
               (google::cloud::CompletionQueue&,
                std::unique_ptr<grpc::ClientContext>,