ARROW-11841: [R][C++] Allow cancelling long-running commands

r/R/arrow-package.R

@@ -102,7 +102,10 @@ supported_dplyr_methods <- list(
     configure_tzdb()
   }
 
-  # register extension types that we use internally
+  # Set interrupt handlers
+  SetEnableSignalStopSource(TRUE)
+
+  # Register extension types that we use internally
   reregister_extension_type(vctrs_extension_type(vctrs::unspecified()))
 
   invisible()
@@ -142,6 +145,19 @@ configure_tzdb <- function() {
   })
 }
 
+# Clean up the StopSource that was registered in .onLoad() so that if the
+# package is reloaded we don't get an error from C++ informing us that
+# a StopSource has already been set up.
+.onUnload <- function(...) {
+  DeinitializeMainRThread()
+}
+
+# While .onUnload should be sufficient, devtools::load_all() does not call it
+# (but it does call .onDetach()). It is safe to call DeinitializeMainRThread()
+# more than once.
+.onDetach <- function(...) {
+  DeinitializeMainRThread()
+}
 
 # True when the OS is linux + and the R version is development
 # helpful for skipping on Valgrind, and the sanitizer checks (clang + gcc) on cran

r/src/arrow_types.h

@@ -94,8 +94,8 @@ class UnwindProtectDetail : public StatusDetail {
   virtual std::string ToString() const { return "R code execution error"; }
 };
 
-static inline Status StatusUnwindProtect(SEXP token) {
-  return Status::Invalid("R code execution error")
+static inline Status StatusUnwindProtect(SEXP token, std::string reason = "") {
+  return Status::Invalid("R code execution error (", reason, ")")
       .WithDetail(std::make_shared<UnwindProtectDetail>(token));
 }
 

r/src/compute-exec.cpp

@@ -73,10 +73,14 @@ class ExecPlanReader : public arrow::RecordBatchReader {
   ExecPlanReader(const std::shared_ptr<arrow::compute::ExecPlan>& plan,
                  const std::shared_ptr<arrow::Schema>& schema,
                  arrow::AsyncGenerator<std::optional<compute::ExecBatch>> sink_gen)
-      : schema_(schema), plan_(plan), sink_gen_(sink_gen), status_(PLAN_NOT_STARTED) {}
+      : schema_(schema),
+        plan_(plan),
+        sink_gen_(sink_gen),
+        plan_status_(PLAN_NOT_STARTED),
+        stop_token_(MainRThread::GetInstance().GetStopToken()) {}
 
   std::string PlanStatus() const {
-    switch (status_) {
+    switch (plan_status_) {
       case PLAN_NOT_STARTED:
         return "PLAN_NOT_STARTED";
       case PLAN_RUNNING:
@@ -91,21 +95,27 @@ class ExecPlanReader : public arrow::RecordBatchReader {
   std::shared_ptr<arrow::Schema> schema() const override { return schema_; }
 
   arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch>* batch_out) override {
-    // TODO(ARROW-11841) check a StopToken to potentially cancel this plan
-
     // If this is the first batch getting pulled, tell the exec plan to
     // start producing
-    if (status_ == PLAN_NOT_STARTED) {
+    if (plan_status_ == PLAN_NOT_STARTED) {
       ARROW_RETURN_NOT_OK(StartProducing());
     }
 
     // If we've closed the reader, keep sending nullptr
     // (consistent with what most RecordBatchReader subclasses do)
-    if (status_ == PLAN_FINISHED) {
+    if (plan_status_ == PLAN_FINISHED) {
       batch_out->reset();
       return arrow::Status::OK();
     }
 
+    // Check for cancellation and stop the plan if we have a request. When
+    // the ExecPlan supports passing a StopToken and handling this itself,
+    // this will be redundant.
+    if (stop_token_.IsStopRequested()) {
+      StopProducing();
+      return stop_token_.Poll();
+    }
+
     auto out = sink_gen_().result();
     if (!out.ok()) {
       StopProducing();
@@ -141,16 +151,17 @@ class ExecPlanReader : public arrow::RecordBatchReader {
   std::shared_ptr<arrow::Schema> schema_;
   std::shared_ptr<arrow::compute::ExecPlan> plan_;
   arrow::AsyncGenerator<std::optional<compute::ExecBatch>> sink_gen_;
-  int status_;
+  ExecPlanReaderStatus plan_status_;
+  arrow::StopToken stop_token_;
 
   arrow::Status StartProducing() {
     ARROW_RETURN_NOT_OK(plan_->StartProducing());
-    status_ = PLAN_RUNNING;
+    plan_status_ = PLAN_RUNNING;
     return arrow::Status::OK();
   }
 
   void StopProducing() {
-    if (status_ == PLAN_RUNNING) {
+    if (plan_status_ == PLAN_RUNNING) {
       // We're done with the plan, but it may still need some time
       // to finish and clean up after itself. To do this, we give a
       // callable with its own copy of the shared_ptr<ExecPlan> so
@@ -164,7 +175,7 @@ class ExecPlanReader : public arrow::RecordBatchReader {
       }
     }
 
-    status_ = PLAN_FINISHED;
+    plan_status_ = PLAN_FINISHED;
     plan_.reset();
     sink_gen_ = arrow::MakeEmptyGenerator<std::optional<compute::ExecBatch>>();
   }

r/src/csv.cpp

@@ -31,7 +31,7 @@ std::shared_ptr<arrow::csv::WriteOptions> csv___WriteOptions__initialize(
       std::make_shared<arrow::csv::WriteOptions>(arrow::csv::WriteOptions::Defaults());
   res->include_header = cpp11::as_cpp<bool>(options["include_header"]);
   res->batch_size = cpp11::as_cpp<int>(options["batch_size"]);
-  res->io_context = arrow::io::IOContext(gc_memory_pool());
+  res->io_context = MainRThread::GetInstance().CancellableIOContext();
   return res;
 }
 
@@ -154,9 +154,9 @@ std::shared_ptr<arrow::csv::TableReader> csv___TableReader__Make(
     const std::shared_ptr<arrow::csv::ReadOptions>& read_options,
     const std::shared_ptr<arrow::csv::ParseOptions>& parse_options,
     const std::shared_ptr<arrow::csv::ConvertOptions>& convert_options) {
-  return ValueOrStop(arrow::csv::TableReader::Make(arrow::io::IOContext(gc_memory_pool()),
-                                                   input, *read_options, *parse_options,
-                                                   *convert_options));
+  return ValueOrStop(arrow::csv::TableReader::Make(
+      MainRThread::GetInstance().CancellableIOContext(), input, *read_options,
+      *parse_options, *convert_options));
 }
 
 // [[arrow::export]]

r/src/extension-impl.cpp

@@ -80,7 +80,7 @@ arrow::Result<std::shared_ptr<arrow::DataType>> RExtensionType::Deserialize(
   // an event loop from wherever this *might* be called is high and hard to
   // predict. As a compromise, just create the instance when it is safe to
   // do so.
-  if (GetMainRThread().IsMainThread()) {
+  if (MainRThread::GetInstance().IsMainThread()) {
     r6_instance();
   }
 

r/src/filesystem.cpp

@@ -16,6 +16,7 @@
 // under the License.
 
 #include "./arrow_types.h"
+#include "./safe-call-into-r.h"
 
 #include <arrow/filesystem/filesystem.h>
 #include <arrow/filesystem/localfs.h>
@@ -239,7 +240,7 @@ std::string fs___FileSystem__type_name(
 // [[arrow::export]]
 std::shared_ptr<fs::LocalFileSystem> fs___LocalFileSystem__create() {
   // Affects OpenInputFile/OpenInputStream
-  auto io_context = arrow::io::IOContext(gc_memory_pool());
+  auto io_context = MainRThread::GetInstance().CancellableIOContext();
   return std::make_shared<fs::LocalFileSystem>(io_context);
 }
 
@@ -334,7 +335,7 @@ std::shared_ptr<fs::S3FileSystem> fs___S3FileSystem__create(
   s3_opts.allow_bucket_creation = allow_bucket_creation;
   s3_opts.allow_bucket_deletion = allow_bucket_deletion;
 
-  auto io_context = arrow::io::IOContext(gc_memory_pool());
+  auto io_context = MainRThread::GetInstance().CancellableIOContext();
   return ValueOrStop(fs::S3FileSystem::Make(s3_opts, io_context));
 }
 
@@ -412,7 +413,7 @@ std::shared_ptr<fs::GcsFileSystem> fs___GcsFileSystem__Make(bool anonymous,
     gcs_opts.default_metadata = strings_to_kvm(options["default_metadata"]);
   }
 
-  auto io_context = arrow::io::IOContext(gc_memory_pool());
+  auto io_context = MainRThread::GetInstance().CancellableIOContext();
   // TODO(ARROW-16884): update when this returns Result
   return fs::GcsFileSystem::Make(gcs_opts, io_context);
 }

r/src/r_to_arrow.cpp

@@ -296,7 +296,7 @@ class AsArrowArrayConverter : public RConverter {
       arrays_.push_back(std::move(array));
       return Status::OK();
     } catch (cpp11::unwind_exception& e) {
-      return StatusUnwindProtect(e.token);
+      return StatusUnwindProtect(e.token, "calling as_arrow_array()");
     }
   }
 

r/src/safe-call-into-r-impl.cpp

@@ -21,18 +21,33 @@
 #include <functional>
 #include <thread>
 
-MainRThread& GetMainRThread() {
+MainRThread& MainRThread::GetInstance() {
   static MainRThread main_r_thread;
   return main_r_thread;
 }
 
 // [[arrow::export]]
-void InitializeMainRThread() { GetMainRThread().Initialize(); }
+void InitializeMainRThread() { MainRThread::GetInstance().Initialize(); }
+
+// [[arrow::export]]
+void DeinitializeMainRThread() { MainRThread::GetInstance().Deinitialize(); }
+
+// [[arrow::export]]
+bool SetEnableSignalStopSource(bool enabled) {
+  bool was_enabled = MainRThread::GetInstance().SignalStopSourceEnabled();
+  if (was_enabled && !enabled) {
+    MainRThread::GetInstance().DisableSignalStopSource();
+  } else if (!was_enabled && enabled) {
+    MainRThread::GetInstance().EnableSignalStopSource();
+  }
+
+  return was_enabled;
+}
 
 // [[arrow::export]]
 bool CanRunWithCapturedR() {
 #if defined(HAS_UNWIND_PROTECT)
-  return GetMainRThread().Executor() == nullptr;
+  return MainRThread::GetInstance().Executor() == nullptr;
 #else
   return false;
 #endif
@@ -42,31 +57,28 @@ bool CanRunWithCapturedR() {
 std::string TestSafeCallIntoR(cpp11::function r_fun_that_returns_a_string,
                               std::string opt) {
   if (opt == "async_with_executor") {
-    std::thread* thread_ptr;
+    std::thread thread;
 
-    auto result =
-        RunWithCapturedR<std::string>([&thread_ptr, r_fun_that_returns_a_string]() {
-          auto fut = arrow::Future<std::string>::Make();
-          thread_ptr = new std::thread([fut, r_fun_that_returns_a_string]() mutable {
-            auto result = SafeCallIntoR<std::string>([&] {
-              return cpp11::as_cpp<std::string>(r_fun_that_returns_a_string());
-            });
+    auto result = RunWithCapturedR<std::string>([&thread, r_fun_that_returns_a_string]() {
+      auto fut = arrow::Future<std::string>::Make();
+      thread = std::thread([&fut, r_fun_that_returns_a_string]() {
+        auto result = SafeCallIntoR<std::string>(
+            [&] { return cpp11::as_cpp<std::string>(r_fun_that_returns_a_string()); });
 
-            fut.MarkFinished(result);
-          });
+        fut.MarkFinished(result);
+      });
 
-          return fut;
-        });
+      return fut;
+    });
 
-    thread_ptr->join();
-    delete thread_ptr;
+    if (thread.joinable()) {
+      thread.join();
+    }
 
     return arrow::ValueOrStop(result);
   } else if (opt == "async_without_executor") {
-    std::thread* thread_ptr;
-
     auto fut = arrow::Future<std::string>::Make();
-    thread_ptr = new std::thread([fut, r_fun_that_returns_a_string]() mutable {
+    std::thread thread([&fut, r_fun_that_returns_a_string]() {
       auto result = SafeCallIntoR<std::string>(
           [&] { return cpp11::as_cpp<std::string>(r_fun_that_returns_a_string()); });
 
@@ -77,8 +89,7 @@ std::string TestSafeCallIntoR(cpp11::function r_fun_that_returns_a_string,
       }
     });
 
-    thread_ptr->join();
-    delete thread_ptr;
+    thread.join();
 
     // We should be able to get this far, but fut will contain an error
     // because it tried to evaluate R code from another thread