with test

Author: paleolimbot

r/adbcdrivermanager/R/async.R

@@ -27,6 +27,25 @@ adbc_async_task_status <- function(task) {
   .Call(RAdbcAsyncTaskWaitFor, task, 0)
 }
 
+adbc_async_task_set_callback <- function(task, callback, loop = later::current_loop()) {
+  # If the task is completed, run the callback (or else the callback
+  # will not run)
+  if (adbc_async_task_status(task) == "ready") {
+    result <- adbc_async_task_result(task)
+    callback(result)
+  } else {
+    .Call(RAdbcAsyncTaskSetCallback, task, callback, loop$id)
+  }
+
+  invisible(task)
+}
+
+adbc_async_task_run_callback <- function(task) {
+  callback <- task$callback
+  result <- adbc_async_task_result(task)
+  callback(result)
+}
+
 adbc_async_task_wait_non_cancellable <- function(task, resolution = 0.05) {
   .Call(RAdbcAsyncTaskWaitFor, task, round(resolution * 1000))
 }

r/adbcdrivermanager/src/async.cc

@@ -45,30 +45,23 @@ static void later_task_callback_wrapper(void* data);
 enum class RAdbcAsyncTaskStatus { NOT_STARTED, STARTED, READY };
 
 struct RAdbcAsyncTask {
-  RAdbcAsyncTask() : callback_sexp(R_NilValue), callback_data_sexp(R_NilValue) {}
+  RAdbcAsyncTask() : callback_data_sexp(R_NilValue) {}
 
-  void SetCallback(SEXP callback, SEXP data, int loop_id) {
-    if (callback_sexp != R_NilValue) {
-      return;
-    }
-
-    callback_sexp = callback;
+  void SetCallback(SEXP data, int loop_id) {
     callback_data_sexp = data;
     later_loop_id = loop_id;
     later_ensure_initialized();
   }
 
   void ScheduleCallbackIfSet() {
-    if (callback_sexp != R_NilValue) {
+    if (callback_data_sexp != R_NilValue) {
       later_execLaterNative2(&later_task_callback_wrapper, this, 0, later_loop_id);
-      callback_sexp = R_NilValue;
     }
   }
 
   AdbcError* return_error{nullptr};
   int* return_code{nullptr};
 
-  SEXP callback_sexp;
   SEXP callback_data_sexp;
   int later_loop_id{-1};
 
@@ -79,9 +72,8 @@ struct RAdbcAsyncTask {
 static void later_task_callback_wrapper(void* data) {
   auto task = reinterpret_cast<RAdbcAsyncTask*>(data);
 
-  SEXP func_sym = PROTECT(Rf_install("adbc_async_run_callback"));
-  SEXP func_call =
-      PROTECT(Rf_lang3(func_sym, task->callback_sexp, task->callback_data_sexp));
+  SEXP func_sym = PROTECT(Rf_install("adbc_async_task_run_callback"));
+  SEXP func_call = PROTECT(Rf_lang2(func_sym, task->callback_data_sexp));
   SEXP pkg_chr = PROTECT(Rf_mkString("adbcdrivermanager"));
   SEXP pkg_ns = PROTECT(R_FindNamespace(pkg_chr));
   Rf_eval(func_call, pkg_ns);
@@ -140,7 +132,7 @@ extern "C" SEXP RAdbcAsyncTaskSetCallback(SEXP task_xptr, SEXP callback_sexp,
   int loop_id = adbc_as_int(loop_id_sexp);
 
   SET_VECTOR_ELT(task_prot, 3, callback_sexp);
-  task->SetCallback(callback_sexp, task_xptr, loop_id);
+  task->SetCallback(task_xptr, loop_id);
   return R_NilValue;
 }
 

r/adbcdrivermanager/src/init.c

@@ -21,6 +21,7 @@
 
 /* generated by tools/make-callentries.R */
 SEXP RAdbcAsyncTaskNew(SEXP error_xptr);
+SEXP RAdbcAsyncTaskSetCallback(SEXP task_xptr, SEXP callback_sexp, SEXP loop_id_sexp);
 SEXP RAdbcAsyncTaskData(SEXP task_xptr);
 SEXP RAdbcAsyncTaskWaitFor(SEXP task_xptr, SEXP duration_ms_sexp);
 SEXP RAdbcAsyncTaskLaunchSleep(SEXP task_xptr, SEXP duration_ms_sexp);
@@ -110,6 +111,7 @@ SEXP RAdbcXptrSetProtected(SEXP xptr, SEXP prot);
 
 static const R_CallMethodDef CallEntries[] = {
     {"RAdbcAsyncTaskNew", (DL_FUNC)&RAdbcAsyncTaskNew, 1},
+    {"RAdbcAsyncTaskSetCallback", (DL_FUNC)&RAdbcAsyncTaskSetCallback, 3},
     {"RAdbcAsyncTaskData", (DL_FUNC)&RAdbcAsyncTaskData, 1},
     {"RAdbcAsyncTaskWaitFor", (DL_FUNC)&RAdbcAsyncTaskWaitFor, 2},
     {"RAdbcAsyncTaskLaunchSleep", (DL_FUNC)&RAdbcAsyncTaskLaunchSleep, 2},

r/adbcdrivermanager/tests/testthat/test-async.R

@@ -74,6 +74,24 @@ test_that("async task waiter works", {
   )
 })
 
+test_that("async tasks can set an R callback", {
+  skip_if_not_installed("later")
+
+  async_called <- FALSE
+  sleep_task <- adbc_async_sleep(200)
+  adbc_async_task_set_callback(sleep_task, function(x) { async_called <<- TRUE })
+  Sys.sleep(0.4)
+  later::run_now()
+  expect_true(async_called)
+
+  # Ensure the callback runs even if the task is already finished
+  async_called <- FALSE
+  sleep_task <- adbc_async_sleep(0)
+  adbc_async_task_set_callback(sleep_task, function(x) { async_called <<- TRUE })
+  Sys.sleep(0.1)
+  expect_true(async_called)
+})
+
 test_that("async task can be converted to a promise", {
   skip_if_not_installed("promises")