From 9cc64fa25d53d184d048f172952117fc5f87d93a Mon Sep 17 00:00:00 2001
From: Henrik Bengtsson <henrik.bengtsson@gmail.com>
Date: Wed, 29 May 2024 14:28:25 -0700
Subject: [PATCH] Added support for backend maintainers to specify "cleanup"
 hook functions [#719]

---
 DESCRIPTION           |  2 +-
 NEWS.md               |  7 +++++++
 R/cluster.R           |  3 +++
 R/multisession.R      |  3 +++
 R/zzz.plan.R          | 18 +++++++++++++-----
 man/future.options.Rd |  4 ++--
 man/multisession.Rd   |  2 +-
 7 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index cd67f309..ff41d7d7 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,5 +1,5 @@
 Package: future
-Version: 1.33.2-9001
+Version: 1.33.2-9002
 Title: Unified Parallel and Distributed Processing in R for Everyone
 Imports:
     digest,
diff --git a/NEWS.md b/NEWS.md
index 7d6a17e4..9010fa95 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,5 +1,12 @@
 # Version (development version)
 
+## New Features
+
+ * Added support for backend maintainers to specify "cleanup" hook
+   functions on future strategies, which are called when switching
+   future plan. These hook functions are specified via the optional
+   `cleanup` attribute, cf. `attr(cluster, "cleanup")`.
+
 ## Bug Fixes
 
  * `resolved()` for `ClusterFuture`:s would produce `Error:
diff --git a/R/cluster.R b/R/cluster.R
index f9f1cb7e..01e28cc6 100644
--- a/R/cluster.R
+++ b/R/cluster.R
@@ -46,4 +46,7 @@ cluster <- function(..., persistent = FALSE, workers = availableWorkers(), envir
 }
 class(cluster) <- c("cluster", "multiprocess", "future", "function")
 attr(cluster, "init") <- TRUE
+attr(cluster, "cleanup") <- function() {
+  ClusterRegistry(action = "stop")
+}
 attr(cluster, "tweakable") <- quote(c(makeClusterPSOCK_args(), "persistent"))
diff --git a/R/multisession.R b/R/multisession.R
index 733ac355..1e6ce5f3 100644
--- a/R/multisession.R
+++ b/R/multisession.R
@@ -81,4 +81,7 @@ multisession <- function(..., workers = availableCores(), lazy = FALSE, rscript_
 }
 class(multisession) <- c("multisession", "cluster", "multiprocess", "future", "function")
 attr(multisession, "init") <- TRUE
+attr(multisession, "cleanup") <- function() {
+  ClusterRegistry(action = "stop")
+}
 attr(multisession, "untweakable") <- c("persistent")
diff --git a/R/zzz.plan.R b/R/zzz.plan.R
index c9795f11..dc2d18a1 100644
--- a/R/zzz.plan.R
+++ b/R/zzz.plan.R
@@ -177,7 +177,16 @@ plan <- local({
   })
 
   plan_cleanup <- function() {
-    ClusterRegistry(action = "stop")
+    evaluator <- stack[[1L]]
+    
+    cleanup <- attr(evaluator, "cleanup", exact = TRUE)
+    if (!is.null(cleanup)) {
+      if (is.function(cleanup)) {
+        cleanup()
+      } else {
+        stop(FutureError(sprintf("Unknown type of 'cleanup' attribute on current future strategy: %s", sQuote(paste(class(cleanup), collapse = ", ")))))
+      }
+    }
   }
 
   plan_init <- function() {
@@ -266,11 +275,11 @@ plan <- local({
     ## Warn about 'multicore' on certain systems
     warn_about_multicore(newStack)
 
-    stack <<- newStack
-
-    ## Stop any (implicitly started) clusters?
+    ## Stop/cleanup any previously registered backends?
     if (cleanup) plan_cleanup()
 
+    stack <<- newStack
+
     ## Initiate future workers?
     if (init) plan_init()
 
@@ -575,4 +584,3 @@ resetWorkers.multicore <- function(x, ...) {
   FutureRegistry(reg, action = "collect-all", earlySignal = FALSE)
   stop_if_not(usedCores() == 0L)
 }
-
diff --git a/man/future.options.Rd b/man/future.options.Rd
index 7a533cdb..d606ba80 100644
--- a/man/future.options.Rd
+++ b/man/future.options.Rd
@@ -188,9 +188,9 @@ All of the above \R \option{future.*} options can be set by corresponding
 environment variable \env{R_FUTURE_*} \emph{when the \pkg{future} package is
 loaded}. This means that those environment variables must be set before
 the \pkg{future} package is loaded in order to have an effect.
-For example, if \code{R_FUTURE_RNG_ONMISUSE = "ignore"}, then option
+For example, if \code{R_FUTURE_RNG_ONMISUSE="ignore"}, then option
 \option{future.rng.onMisuse} is set to \code{"ignore"} (character string).
-Similarly, if \code{R_FUTURE_GLOBALS_MAXSIZE = "50000000"}, then option
+Similarly, if \code{R_FUTURE_GLOBALS_MAXSIZE="50000000"}, then option
 \option{future.globals.maxSize} is set to \code{50000000} (numeric).
 }
 
diff --git a/man/multisession.Rd b/man/multisession.Rd
index 1dd50bc5..c309d638 100644
--- a/man/multisession.Rd
+++ b/man/multisession.Rd
@@ -37,7 +37,7 @@ identified.}
 }
 \value{
 A \link{MultisessionFuture}.
-If \code{workers == 1}, then all processing using done in the
+If \code{workers == 1}, then all processing is done in the
 current/main \R session and we therefore fall back to using a
 lazy future.  To override this fallback, use \code{workers = I(1)}.
 }