ARROW-8732: [C++] Add basic cancellation API

cpp/src/arrow/CMakeLists.txt

@@ -181,6 +181,7 @@ set(ARROW_SRCS
     util/bitmap_builders.cc
     util/bitmap_ops.cc
     util/bpacking.cc
+    util/cancel.cc
     util/compression.cc
     util/cpu_info.cc
     util/decimal.cc

cpp/src/arrow/status.cc

@@ -68,6 +68,9 @@ std::string Status::CodeAsString(StatusCode code) {
     case StatusCode::Invalid:
       type = "Invalid";
       break;
+    case StatusCode::Cancelled:
+      type = "Cancelled";
+      break;
     case StatusCode::IOError:
       type = "IOError";
       break;

cpp/src/arrow/status.h

@@ -83,6 +83,7 @@ enum class StatusCode : char {
   IOError = 5,
   CapacityError = 6,
   IndexError = 7,
+  Cancelled = 8,
   UnknownError = 9,
   NotImplemented = 10,
   SerializationError = 11,
@@ -204,6 +205,12 @@ class ARROW_MUST_USE_TYPE ARROW_EXPORT Status : public util::EqualityComparable<
     return Status::FromArgs(StatusCode::Invalid, std::forward<Args>(args)...);
   }
 
+  /// Return an error status for cancelled operation
+  template <typename... Args>
+  static Status Cancelled(Args&&... args) {
+    return Status::FromArgs(StatusCode::Cancelled, std::forward<Args>(args)...);
+  }
+
   /// Return an error status when an index is out of bounds
   template <typename... Args>
   static Status IndexError(Args&&... args) {
@@ -263,6 +270,8 @@ class ARROW_MUST_USE_TYPE ARROW_EXPORT Status : public util::EqualityComparable<
   bool IsKeyError() const { return code() == StatusCode::KeyError; }
   /// Return true iff the status indicates invalid data.
   bool IsInvalid() const { return code() == StatusCode::Invalid; }
+  /// Return true iff the status indicates a cancelled operation.
+  bool IsCancelled() const { return code() == StatusCode::Cancelled; }
   /// Return true iff the status indicates an IO-related failure.
   bool IsIOError() const { return code() == StatusCode::IOError; }
   /// Return true iff the status indicates a container reaching capacity limits.

cpp/src/arrow/util/CMakeLists.txt

@@ -67,9 +67,10 @@ add_arrow_test(utility-test
 
 add_arrow_test(threading-utility-test
                SOURCES
-               future_test
-               task_group_test
-               thread_pool_test)
+               cancel_test.cc
+               future_test.cc
+               task_group_test.cc
+               thread_pool_test.cc)
 
 add_arrow_benchmark(bit_block_counter_benchmark)
 add_arrow_benchmark(bit_util_benchmark)

cpp/src/arrow/util/cancel.cc

@@ -0,0 +1,124 @@
+// 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 <atomic>
+#include <mutex>
+#include <utility>
+
+#include "arrow/util/cancel.h"
+#include "arrow/util/logging.h"
+#include "arrow/util/visibility.h"
+
+namespace arrow {
+
+StopCallback::StopCallback(StopToken* token, Callable cb)
+    : token_(token), cb_(std::move(cb)) {
+  if (token_ != nullptr) {
+    DCHECK(cb_);
+    // May call *this
+    token_->SetCallback(this);
+  }
+}
+
+StopCallback::~StopCallback() {
+  if (token_ != nullptr) {
+    token_->RemoveCallback(this);
+  }
+}
+
+StopCallback::StopCallback(StopCallback&& other) { *this = std::move(other); }
+
+StopCallback& StopCallback::operator=(StopCallback&& other) {
+  token_ = other.token_;
+  if (token_ != nullptr) {
+    other.token_ = nullptr;
+    token_->RemoveCallback(&other);
+  }
+  cb_ = std::move(other.cb_);
+  if (token_ != nullptr) {
+    // May call *this
+    token_->SetCallback(this);
+  }
+  return *this;
+}
+
+void StopCallback::Call(const Status& st) {
+  if (cb_) {
+    // Forget callable after calling it
+    Callable local_cb;
+    cb_.swap(local_cb);
+    local_cb(st);
+  }
+}
+
+// NOTE: We care mainly about the making the common case (not cancelled) fast.
+
+struct StopToken::Impl {
+  std::atomic<bool> requested_{false};
+  std::mutex mutex_;
+  StopCallback* cb_{nullptr};
+  Status cancel_error_;
+};
+
+StopToken::StopToken() : impl_(new Impl()) {}
+
+StopToken::~StopToken() {}
+
+Status StopToken::Poll() {
+  if (impl_->requested_) {
+    std::lock_guard<std::mutex> lock(impl_->mutex_);
+    return impl_->cancel_error_;
+  }
+  return Status::OK();
+}
+
+bool StopToken::IsStopRequested() { return impl_->requested_; }
+
+void StopToken::RequestStop() { RequestStop(Status::Cancelled("Operation cancelled")); }
+
+void StopToken::RequestStop(Status st) {
+  std::lock_guard<std::mutex> lock(impl_->mutex_);
+  DCHECK(!st.ok());
+  if (!impl_->requested_) {
+    impl_->requested_ = true;
+    impl_->cancel_error_ = std::move(st);
+    if (impl_->cb_) {
+      impl_->cb_->Call(impl_->cancel_error_);
+    }
+  }
+}
+
+StopCallback StopToken::SetCallback(StopCallback::Callable cb) {
+  return StopCallback(this, std::move(cb));
+}
+
+void StopToken::SetCallback(StopCallback* cb) {
+  std::lock_guard<std::mutex> lock(impl_->mutex_);
+  DCHECK_EQ(impl_->cb_, nullptr);
+  impl_->cb_ = cb;
+  if (impl_->requested_) {
+    impl_->cb_->Call(impl_->cancel_error_);
+  }
+}
+
+void StopToken::RemoveCallback(StopCallback* cb) {
+  std::lock_guard<std::mutex> lock(impl_->mutex_);
+  DCHECK_EQ(impl_->cb_, cb);
+  impl_->cb_ = nullptr;
+}
+
+}  // namespace arrow

cpp/src/arrow/util/cancel.h

@@ -0,0 +1,86 @@
+// 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 <functional>
+#include <memory>
+#include <string>
+
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/macros.h"
+#include "arrow/util/visibility.h"
+
+namespace arrow {
+
+class StopToken;
+
+// A RAII wrapper that automatically registers and unregisters
+// a callback to a StopToken.
+class ARROW_MUST_USE_TYPE ARROW_EXPORT StopCallback {
+ public:
+  using Callable = std::function<void(const Status&)>;
+  StopCallback(StopToken* token, Callable cb);
+  ~StopCallback();
+
+  ARROW_DISALLOW_COPY_AND_ASSIGN(StopCallback);
+  StopCallback(StopCallback&&);
+  StopCallback& operator=(StopCallback&&);
+
+  void Call(const Status&);
+
+ protected:
+  StopToken* token_;
+  Callable cb_;
+};
+
+class ARROW_EXPORT StopToken {
+ public:
+  StopToken();
+  ~StopToken();
+  ARROW_DISALLOW_COPY_AND_ASSIGN(StopToken);
+
+  // NOTE: all APIs here are non-blocking.  For consumers, waiting is done
+  // at a higher level using e.g. Future.  Producers don't have to wait
+  // on a StopToken.
+
+  // Consumer API (the side that stops)
+  void RequestStop();
+  void RequestStop(Status error);
+
+  // Producer API (the side that gets asked to stopped)
+  Status Poll();
+  bool IsStopRequested();
+
+  // Register a callback that will be called whenever cancellation happens.
+  // Note the callback may be called immediately, if cancellation was already
+  // requested.  The callback will be unregistered when the returned object
+  // is destroyed.
+  StopCallback SetCallback(StopCallback::Callable cb);
+
+ protected:
+  struct Impl;
+  std::unique_ptr<Impl> impl_;
+
+  void SetCallback(StopCallback* cb);
+  void RemoveCallback(StopCallback* cb);
+
+  friend class StopCallback;
+};
+
+}  // namespace arrow