gRPC: fix cases where gRPC call could be finished twice

Firestore/core/src/firebase/firestore/remote/grpc_stream.cc

@@ -21,6 +21,7 @@
 
 #include "Firestore/core/src/firebase/firestore/remote/grpc_connection.h"
 #include "Firestore/core/src/firebase/firestore/remote/grpc_util.h"
+#include "Firestore/core/src/firebase/firestore/util/log.h"
 
 namespace firebase {
 namespace firestore {
@@ -96,6 +97,7 @@ GrpcStream::GrpcStream(
 }
 
 GrpcStream::~GrpcStream() {
+  LOG_DEBUG("GrpcStream('%s'): destroying stream", this);
   HARD_ASSERT(completions_.empty(),
               "GrpcStream is being destroyed without proper shutdown");
   MaybeUnregister();
@@ -156,11 +158,15 @@ void GrpcStream::MaybeWrite(absl::optional<BufferedWrite> maybe_write) {
 }
 
 void GrpcStream::FinishImmediately() {
+  LOG_DEBUG("GrpcStream('%s'): finishing without notifying observers", this);
+
   Shutdown();
   UnsetObserver();
 }
 
 void GrpcStream::FinishAndNotify(const Status& status) {
+  LOG_DEBUG("GrpcStream('%s'): finishing and notifying observers", this);
+
   Shutdown();
 
   if (observer_) {
@@ -173,22 +179,27 @@ void GrpcStream::FinishAndNotify(const Status& status) {
 }
 
 void GrpcStream::Shutdown() {
+  LOG_DEBUG("GrpcStream('%s'): shutting down; completions: %s, is finished: %s",
+            this, completions_.size(), is_grpc_call_finished_);
+
   MaybeUnregister();
-  if (completions_.empty()) {
-    // Nothing to cancel -- either the call was already finished, or it has
-    // never been started.
-    return;
+
+  if (!completions_.empty() && !is_grpc_call_finished_) {
+    // Important: during normal operation, the stream always has a pending read
+    // operation, so `Shutdown` would hang indefinitely if we didn't cancel the
+    // `context_`. However, if the stream has already failed, avoid canceling
+    // the context to avoid overwriting the status captured during the
+    // `OnOperationFailed`.
+
+    context_->TryCancel();
+    FinishGrpcCall({});
   }
 
-  // Important: since the stream always has a pending read operation,
-  // cancellation has to be called, or else the read would hang forever, and
-  // finish operation will never get completed.
-  // (on the other hand, when an operation fails, cancellation should not be
-  // called, otherwise the real failure cause will be overwritten by status
-  // "canceled".)
-  context_->TryCancel();
-  FinishCall({});
-  // Wait until "finish" is off the queue.
+  // Drain the completions -- `Shutdown` guarantees to bring the stream into
+  // destructible state. Two possibilities here:
+  // - completions are empty -- nothing to block on;
+  // - the only completion is "finish" (whether enqueued by this function or
+  //   previously) -- "finish" is a very fast operation.
   FastFinishCompletionsBlocking();
 }
 
@@ -199,7 +210,12 @@ void GrpcStream::MaybeUnregister() {
   }
 }
 
-void GrpcStream::FinishCall(const OnSuccess& callback) {
+void GrpcStream::FinishGrpcCall(const OnSuccess& callback) {
+  LOG_DEBUG("GrpcStream('%s'): finishing the underlying call", this);
+
+  HARD_ASSERT(!is_grpc_call_finished_, "FinishGrpcCall called twice");
+  is_grpc_call_finished_ = true;
+
   // All completions issued by this call must be taken off the queue before
   // finish operation can be enqueued.
   FastFinishCompletionsBlocking();
@@ -208,6 +224,9 @@ void GrpcStream::FinishCall(const OnSuccess& callback) {
 }
 
 void GrpcStream::FastFinishCompletionsBlocking() {
+  LOG_DEBUG("GrpcStream('%s'): fast finishing %s completion(s)", this,
+            completions_.size());
+
   // TODO(varconst): reset buffered_writer_? Should not be necessary, because it
   // should never be called again after a call to Finish.
 
@@ -226,30 +245,35 @@ void GrpcStream::FastFinishCompletionsBlocking() {
 
 bool GrpcStream::WriteAndFinish(grpc::ByteBuffer&& message) {
   bool did_last_write = false;
+  if (!is_grpc_call_finished_) {
+    did_last_write = TryLastWrite(std::move(message));
+  }
+  FinishImmediately();
+  return did_last_write;
+}
 
+bool GrpcStream::TryLastWrite(grpc::ByteBuffer&& message) {
   absl::optional<BufferedWrite> maybe_write =
       buffered_writer_.EnqueueWrite(std::move(message));
   // Only bother with the last write if there is no active write at the moment.
-  if (maybe_write) {
-    BufferedWrite last_write = std::move(maybe_write).value();
-    GrpcCompletion* completion = NewCompletion(Type::Write, {});
-    *completion->message() = last_write.message;
-    call_->WriteLast(*completion->message(), grpc::WriteOptions{}, completion);
-
-    // Empirically, the write normally takes less than a millisecond to finish
-    // (both with and without network connection), and never more than several
-    // dozen milliseconds. Nevertheless, ensure `WriteAndFinish` doesn't hang if
-    // there happen to be circumstances under which the write may block
-    // indefinitely (in that case, rely on the fact that canceling a gRPC call
-    // makes all pending operations come back from the queue quickly).
-    auto status = completion->WaitUntilOffQueue(std::chrono::milliseconds(500));
-    if (status == std::future_status::ready) {
-      did_last_write = true;
-    }
+  if (!maybe_write) {
+    return false;
   }
 
-  FinishImmediately();
-  return did_last_write;
+  BufferedWrite last_write = std::move(maybe_write).value();
+  GrpcCompletion* completion = NewCompletion(Type::Write, {});
+  *completion->message() = last_write.message;
+  call_->WriteLast(*completion->message(), grpc::WriteOptions{}, completion);
+
+  // Empirically, the write normally takes less than a millisecond to finish
+  // (both with and without network connection), and never more than several
+  // dozen milliseconds. Nevertheless, ensure `WriteAndFinish` doesn't hang if
+  // there happen to be circumstances under which the write may block
+  // indefinitely (in that case, rely on the fact that canceling a gRPC call
+  // makes all pending operations come back from the queue quickly).
+
+  auto status = completion->WaitUntilOffQueue(std::chrono::milliseconds(500));
+  return status == std::future_status::ready;
 }
 
 GrpcStream::Metadata GrpcStream::GetResponseHeaders() const {
@@ -277,7 +301,13 @@ void GrpcStream::OnWrite() {
 }
 
 void GrpcStream::OnOperationFailed() {
-  FinishCall([this](const GrpcCompletion* completion) {
+  if (is_grpc_call_finished_) {
+    // If a finish operation has been enqueued already (possibly by a previous
+    // failed operation), there's nothing to do.
+    return;
+  }
+
+  FinishGrpcCall([this](const GrpcCompletion* completion) {
     Status status = ConvertStatus(*completion->status());
     FinishAndNotify(status);
   });
@@ -303,6 +333,8 @@ GrpcCompletion* GrpcStream::NewCompletion(Type tag,
         } else {
           // Use the same error-handling for all operations; all errors are
           // unrecoverable.
+          LOG_DEBUG("GrpcStream('%s'): operation of type %s failed", this,
+                    completion->type());
           OnOperationFailed();
         }
       };

Firestore/core/src/firebase/firestore/remote/grpc_stream.h

@@ -180,6 +180,7 @@ class GrpcStream : public GrpcCall {
  private:
   void Read();
   void MaybeWrite(absl::optional<internal::BufferedWrite> maybe_write);
+  bool TryLastWrite(grpc::ByteBuffer&& message);
 
   void Shutdown();
   void UnsetObserver() {
@@ -200,7 +201,7 @@ class GrpcStream : public GrpcCall {
   // was started. Presumes that any pending completions will quickly come off
   // the queue and will block until they do, so this must only be invoked when
   // the current call either failed (`OnOperationFailed`) or canceled.
-  void FinishCall(const OnSuccess& callback);
+  void FinishGrpcCall(const OnSuccess& callback);
 
   // Blocks until all the completions issued by this stream come out from the
   // gRPC completion queue. Once they do, it is safe to delete this `GrpcStream`
@@ -231,6 +232,9 @@ class GrpcStream : public GrpcCall {
   internal::BufferedWriter buffered_writer_;
 
   std::vector<GrpcCompletion*> completions_;
+
+  // gRPC asserts that a call is finished exactly once.
+  bool is_grpc_call_finished_ = false;
 };
 
 }  // namespace remote

Firestore/core/test/firebase/firestore/remote/grpc_stream_test.cc

@@ -448,6 +448,54 @@ TEST_F(GrpcStreamTest, ObserverCanImmediatelyDestroyStreamOnFinishAndNotify) {
   });
 }
 
+// Double finish
+
+TEST_F(GrpcStreamTest, DoubleFinish_FailThenFinishImmediately) {
+  worker_queue.EnqueueBlocking([&] { stream->Start(); });
+
+  ForceFinish({{Type::Read, Error}});
+  KeepPollingGrpcQueue();
+  worker_queue.EnqueueBlocking(
+      [&] { EXPECT_NO_THROW(stream->FinishImmediately()); });
+}
+
+TEST_F(GrpcStreamTest, DoubleFinish_FailThenWriteAndFinish) {
+  worker_queue.EnqueueBlocking([&] { stream->Start(); });
+
+  ForceFinish({{Type::Read, Error}});
+  KeepPollingGrpcQueue();
+  worker_queue.EnqueueBlocking(
+      [&] { EXPECT_NO_THROW(stream->WriteAndFinish({})); });
+}
+
+TEST_F(GrpcStreamTest, DoubleFinish_FailThenFailAgain) {
+  worker_queue.EnqueueBlocking([&] {
+    stream->Start();
+    stream->Write({});
+  });
+
+  int failures_count = 0;
+  auto future = tester.ForceFinishAsync([&](GrpcCompletion* completion) {
+    switch (completion->type()) {
+      case Type::Read:
+      case Type::Write:
+        ++failures_count;
+        completion->Complete(false);
+        return failures_count == 2;
+      default:
+        UnexpectedType(completion);
+        return true;
+    }
+  });
+  future.wait();
+  worker_queue.EnqueueBlocking([] {});
+
+  // Normally, "Finish" never fails, but for the test it's easier to abuse the
+  // finish operation that has already been enqueued by `OnOperationFailed`
+  // rather than adding a new operation.
+  ForceFinish({{Type::Finish, Error}});
+}
+
 }  // namespace remote
 }  // namespace firestore
 }  // namespace firebase