fix: CompletionQueue::RunAsync is always async

google/cloud/completion_queue.cc

@@ -76,6 +76,38 @@ class AsyncTimerFuture : public internal::AsyncGrpcOperation {
   std::unique_ptr<grpc::Alarm> alarm_;
 };
 
+class AsyncFunction : public internal::AsyncGrpcOperation {
+ public:
+  AsyncFunction(std::unique_ptr<internal::RunAsyncBase> fun,
+                std::unique_ptr<grpc::Alarm> alarm)
+      : fun_(std::move(fun)), alarm_(std::move(alarm)) {}
+
+  void Set(grpc::CompletionQueue& cq,
+           std::chrono::system_clock::time_point deadline, void* tag) {
+    if (alarm_) {
+      alarm_->Set(&cq, deadline, tag);
+    }
+  }
+
+  void Cancel() override {
+    if (alarm_) {
+      alarm_->Cancel();
+    }
+  }
+
+ private:
+  bool Notify(bool) override {
+    fun_->exec();
+    fun_.reset();
+    return true;
+  }
+
+  std::shared_ptr<internal::CompletionQueueImpl> cq_;
+  std::unique_ptr<internal::RunAsyncBase> fun_;
+  // Holds the underlying handle, it might be a nullptr in tests.
+  std::unique_ptr<grpc::Alarm> alarm_;
+};
+
 }  // namespace
 
 CompletionQueue::CompletionQueue() : impl_(new internal::CompletionQueueImpl) {}
@@ -95,6 +127,13 @@ CompletionQueue::MakeDeadlineTimer(
   return op->GetFuture();
 }
 
+void CompletionQueue::RunAsyncImpl(std::unique_ptr<internal::RunAsyncBase> f) {
+  auto deadline = std::chrono::system_clock::now();
+  auto op = std::make_shared<AsyncFunction>(std::move(f), impl_->CreateAlarm());
+  impl_->StartOperation(
+      op, [&](void* tag) { op->Set(impl_->cq(), deadline, tag); });
+}
+
 }  // namespace GOOGLE_CLOUD_CPP_NS
 }  // namespace cloud
 }  // namespace google

google/cloud/completion_queue.h

@@ -20,10 +20,20 @@
 #include "google/cloud/internal/completion_queue_impl.h"
 #include "google/cloud/status_or.h"
 #include "google/cloud/version.h"
+#include "absl/memory/memory.h"
+#include "absl/meta/type_traits.h"
 
 namespace google {
 namespace cloud {
 inline namespace GOOGLE_CLOUD_CPP_NS {
+namespace internal {
+// Type erase the callables in RunAsync()
+struct RunAsyncBase {
+  virtual ~RunAsyncBase() = default;
+  virtual void exec() = 0;
+};
+}  // namespace internal
+
 /**
  * Call the functor associated with asynchronous operations when they complete.
  */
@@ -178,20 +188,29 @@ class CompletionQueue {
             typename std::enable_if<
                 internal::CheckRunAsyncCallback<Functor>::value, int>::type = 0>
   void RunAsync(Functor&& functor) {
-    auto impl = impl_;
-    MakeRelativeTimer(std::chrono::seconds(0))
-        .then(
-            [impl, functor](
-                future<
-                    StatusOr<std::chrono::system_clock::time_point>>) mutable {
-              // We intentionally ignore the status here; the functor is always
-              // called, even after a call to `CancelAll`.
-              CompletionQueue cq(impl);
-              functor(cq);
-            });
+    class Wrapper : public internal::RunAsyncBase {
+     public:
+      Wrapper(std::weak_ptr<internal::CompletionQueueImpl> impl, Functor&& f)
+          : impl_(std::move(impl)), fun_(std::forward<Functor>(f)) {}
+      ~Wrapper() override = default;
+      void exec() override {
+        auto impl = impl_.lock();
+        if (!impl) return;
+        CompletionQueue cq(std::move(impl));
+        fun_(cq);
+      }
+
+     private:
+      std::weak_ptr<internal::CompletionQueueImpl> impl_;
+      absl::decay_t<Functor> fun_;
+    };
+    RunAsyncImpl(
+        absl::make_unique<Wrapper>(impl_, std::forward<Functor>(functor)));
   }
 
  private:
+  void RunAsyncImpl(std::unique_ptr<internal::RunAsyncBase>);
+
   std::shared_ptr<internal::CompletionQueueImpl> impl_;
 };
 

google/cloud/completion_queue_test.cc

@@ -32,6 +32,7 @@ namespace btadmin = ::google::bigtable::admin::v2;
 namespace btproto = ::google::bigtable::v2;
 using ::google::cloud::testing_util::MockCompletionQueue;
 using ::testing::_;
+using ::testing::Contains;
 using ::testing::StrictMock;
 
 class MockClient {
@@ -341,6 +342,61 @@ TEST(CompletionQueueTest, RunAsyncCompletionQueueDestroyed) {
   done_promise.get_future().get();
 }
 
+TEST(CompletionQueueTest, RunAsyncMoveOnly) {
+  struct MoveOnly {
+    promise<void> p;
+    void operator()(CompletionQueue&) { p.set_value(); }
+  };
+  static_assert(!std::is_copy_assignable<MoveOnly>::value,
+                "MoveOnly test type should not copy-assignable");
+
+  promise<void> p;
+  auto done = p.get_future();
+  CompletionQueue cq;
+  std::thread t{[&cq] { cq.Run(); }};
+  cq.RunAsync(MoveOnly{std::move(p)});
+  done.get();
+  cq.Shutdown();
+  t.join();
+}
+
+TEST(CompletionQueueTest, RunAsyncThread) {
+  CompletionQueue cq;
+
+  std::set<std::thread::id> runner_ids;
+  auto constexpr kRunners = 8;
+  std::vector<std::thread> runners(kRunners);
+  for (auto& t : runners) {
+    promise<std::thread::id> started;
+    auto f = started.get_future();
+    t = std::thread(
+        [&cq](promise<std::thread::id> p) {
+          p.set_value(std::this_thread::get_id());
+          cq.Run();
+        },
+        std::move(started));
+    runner_ids.insert(f.get());
+  }
+
+  auto constexpr kIterations = 10000;
+  std::vector<promise<std::thread::id>> pending(kIterations);
+  std::vector<future<std::thread::id>> actual;
+  for (int i = 0; i != kIterations; ++i) {
+    auto& p = pending[i];
+    actual.push_back(p.get_future());
+    cq.RunAsync(
+        [&p](CompletionQueue&) { p.set_value(std::this_thread::get_id()); });
+  }
+
+  for (auto& done : actual) {
+    auto id = done.get();
+    EXPECT_THAT(runner_ids, Contains(id));
+  }
+
+  cq.Shutdown();
+  for (auto& t : runners) t.join();
+}
+
 // Sets up a timer that reschedules itself and verifies we can shut down
 // cleanly whether we call `CancelAll()` on the queue first or not.
 namespace {