ARROW-13680: [C++] Create an asynchronous nursery to simplify capture logic

cpp/src/arrow/CMakeLists.txt

@@ -190,6 +190,7 @@ set(ARROW_SRCS
     io/slow.cc
     io/stdio.cc
     io/transform.cc
+    util/async_nursery.cc
     util/basic_decimal.cc
     util/bit_block_counter.cc
     util/bit_run_reader.cc

cpp/src/arrow/util/CMakeLists.txt

@@ -42,6 +42,7 @@ add_arrow_test(utility-test
                SOURCES
                align_util_test.cc
                async_generator_test.cc
+               async_nursery_test.cc
                bit_block_counter_test.cc
                bit_util_test.cc
                cache_test.cc

cpp/src/arrow/util/async_nursery.cc

@@ -0,0 +1,123 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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 "arrow/util/async_nursery.h"
+
+#include "arrow/util/logging.h"
+
+namespace arrow {
+namespace util {
+
+AsyncCloseable::AsyncCloseable() : on_closed_(Future<>::Make()) {}
+AsyncCloseable::AsyncCloseable(AsyncCloseable* parent) : on_closed_(Future<>::Make()) {
+  parent->AddDependentTask(OnClosed());
+}
+
+AsyncCloseable::~AsyncCloseable() {
+  DCHECK_NE(nursery_, nullptr) << "An AsyncCloseable must be created with a nursery "
+                                  "using MakeSharedCloseable or MakeUniqueCloseable";
+}
+
+const Future<>& AsyncCloseable::OnClosed() { return on_closed_; }
+
+void AsyncCloseable::AddDependentTask(const Future<>& task) {
+  DCHECK(!closed_);
+  if (num_tasks_outstanding_.fetch_add(1) == 1) {
+    tasks_finished_ = Future<>::Make();
+  }
+  task.AddCallback([this](const Status& st) {
+    if (num_tasks_outstanding_.fetch_sub(1) == 1 && closed_.load()) {
+      tasks_finished_.MarkFinished(st);
+    }
+  });
+}
+
+void AsyncCloseable::SetNursery(Nursery* nursery) { nursery_ = nursery; }
+
+void AsyncCloseable::Destroy() {
+  DCHECK_NE(nursery_, nullptr);
+  closed_ = true;
+  nursery_->num_closeables_destroyed_.fetch_add(1);
+  Future<> finish_fut;
+  if (tasks_finished_.is_valid()) {
+    if (num_tasks_outstanding_.fetch_sub(1) > 1) {
+      finish_fut = AllComplete({DoClose(), tasks_finished_});
+    } else {
+      // Any added tasks have already finished so there is nothing to wait for
+      finish_fut = DoClose();
+    }
+  } else {
+    // No dependent tasks were added
+    finish_fut = DoClose();
+  }
+  finish_fut.AddCallback([this](const Status& st) {
+    if (on_closed_.is_valid()) {
+      on_closed_.MarkFinished(st);
+    }
+    nursery_->OnTaskFinished(st);
+    delete this;
+  });
+}
+
+Status AsyncCloseable::CheckClosed() const {
+  if (closed_.load()) {
+    return Status::Invalid("Invalid operation after Close");
+  }
+  return Status::OK();
+}
+
+void AsyncCloseablePimpl::Init(AsyncCloseable* impl) { impl_ = impl; }
+void AsyncCloseablePimpl::Destroy() { impl_->Destroy(); }
+void AsyncCloseablePimpl::SetNursery(Nursery* nursery) { impl_->SetNursery(nursery); }
+
+Nursery::Nursery() : finished_(Future<>::Make()) {}
+
+Status Nursery::WaitForFinish() {
+  if (num_closeables_destroyed_.load() != num_closeables_created_.load()) {
+    return Status::UnknownError(
+        "Not all closeables that were created during the nursery were destroyed.  "
+        "Something must be holding onto a shared_ptr/unique_ptr reference.");
+  }
+  if (num_tasks_outstanding_.fetch_sub(1) == 1) {
+    // All tasks done, nothing to wait for
+    return Status::OK();
+  }
+  return finished_.status();
+}
+
+void Nursery::OnTaskFinished(Status st) {
+  if (num_tasks_outstanding_.fetch_sub(1) == 1) {
+    finished_.MarkFinished(std::move(st));
+  }
+}
+
+Status Nursery::RunInNursery(std::function<void(Nursery*)> task) {
+  Nursery nursery;
+  task(&nursery);
+  return nursery.WaitForFinish();
+}
+
+Status Nursery::RunInNursery(std::function<Status(Nursery*)> task) {
+  Nursery nursery;
+  Status task_st = task(&nursery);
+  // Need to wait for everything to finish, even if invalid status
+  Status close_st = nursery.WaitForFinish();
+  return task_st & close_st;
+}
+
+}  // namespace util
+}  // namespace arrow

cpp/src/arrow/util/async_nursery.h

@@ -0,0 +1,154 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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.
+
+#pragma once
+
+#include <list>
+
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/future.h"
+#include "arrow/util/mutex.h"
+
+namespace arrow {
+namespace util {
+
+class Nursery;
+class AsyncCloseablePimpl;
+
+template <typename T>
+struct DestroyingDeleter {
+  void operator()(T* p) { p->Destroy(); }
+};
+
+/// An object which should be asynchronously closed before it is destroyed
+///
+/// Any AsyncCloseable must be kept alive until its parent is destroyed (this is a given
+/// if the parent is a nursery).  For shorter lived tasks/objects consider
+/// OwnedAsyncCloseable adding a dependent task.
+class ARROW_EXPORT AsyncCloseable : public std::enable_shared_from_this<AsyncCloseable> {
+ public:
+  AsyncCloseable();
+  explicit AsyncCloseable(AsyncCloseable* parent);
+  virtual ~AsyncCloseable();
+
+  /// Returns a future that is completed when this object is finished closing
+  const Future<>& OnClosed();
+
+ protected:
+  /// Subclasses should override this and perform any cleanup.  Once the future returned
+  /// by this method finishes then this object is eligible for destruction and any
+  /// reference to `this` will be invalid
+  virtual Future<> DoClose() = 0;
+
+  /// This method is called by subclasses to add tasks which must complete before the
+  /// object can be safely deleted
+  void AddDependentTask(const Future<>& task);
+  /// This method can be called by subclasses for error checking purposes.  It will
+  /// return an invalid status if this object has started closing
+  Status CheckClosed() const;
+
+  Nursery* nursery_;
+
+ private:
+  void SetNursery(Nursery* nursery);
+  void Destroy();
+
+  Future<> on_closed_;
+  Future<> tasks_finished_;
+  std::atomic<bool> closed_{false};
+  std::atomic<uint32_t> num_tasks_outstanding_{1};
+
+  friend Nursery;
+  template <typename T>
+  friend struct DestroyingDeleter;
+  friend AsyncCloseablePimpl;
+};
+
+class ARROW_EXPORT AsyncCloseablePimpl {
+ protected:
+  void Init(AsyncCloseable* impl);
+
+ private:
+  void SetNursery(Nursery* nursery);
+  void Destroy();
+
+  AsyncCloseable* impl_;
+
+  friend Nursery;
+  template <typename T>
+  friend struct DestroyingDeleter;
+};
+
+class ARROW_EXPORT Nursery {
+ public:
+  template <typename T, typename... Args>
+  typename std::enable_if<!std::is_array<T>::value, std::shared_ptr<T>>::type
+  MakeSharedCloseable(Args&&... args) {
+    static_assert(std::is_base_of<AsyncCloseable, T>::value,
+                  "Nursery::MakeSharedCloseable only works with AsyncCloseable types");
+    num_closeables_created_.fetch_add(1);
+    num_tasks_outstanding_.fetch_add(1);
+    std::shared_ptr<T> shared_closeable(new T(std::forward<Args&&>(args)...),
+                                        DestroyingDeleter<T>());
+    shared_closeable->SetNursery(this);
+    return shared_closeable;
+  }
+
+  template <typename T, typename... Args>
+  typename std::enable_if<!std::is_array<T>::value,
+                          std::unique_ptr<T, DestroyingDeleter<T>>>::type
+  MakeUniqueCloseable(Args&&... args) {
+    static_assert(std::is_base_of<AsyncCloseable, T>::value,
+                  "Nursery::MakeUniqueCloseable only works with AsyncCloseable types");
+    num_closeables_created_.fetch_add(1);
+    num_tasks_outstanding_.fetch_add(1);
+    auto unique_closeable = std::unique_ptr<T, DestroyingDeleter<T>>(
+        new T(std::forward<Args>(args)...), DestroyingDeleter<T>());
+    unique_closeable->SetNursery(this);
+    return unique_closeable;
+  }
+
+  template <typename T>
+  void AddDependentTask(const Future<T>& task) {
+    num_tasks_outstanding_.fetch_add(1);
+    task.AddCallback([this](const Result<T>& res) { OnTaskFinished(res.status()); });
+  }
+
+  /// Runs `task` within a nursery.  This method will not return until
+  /// all roots added by the task have been closed and destroyed.
+  static Status RunInNursery(std::function<void(Nursery*)> task);
+  static Status RunInNursery(std::function<Status(Nursery*)> task);
+
+ protected:
+  Status WaitForFinish();
+  void OnTaskFinished(Status st);
+
+  Nursery();
+
+  // Rather than keep a separate closing flag (and requiring mutex) we treat
+  // "closing" as a default task
+  std::atomic<uint32_t> num_tasks_outstanding_{1};
+  std::atomic<uint32_t> num_closeables_created_{0};
+  std::atomic<uint32_t> num_closeables_destroyed_{0};
+  Future<> finished_;
+
+  friend AsyncCloseable;
+};
+
+}  // namespace util
+}  // namespace arrow