From ed95bcfc93a9bd0fc61be84b68b1d1938e128cb5 Mon Sep 17 00:00:00 2001
From: Stephen Dolan <sdolan@janestreet.com>
Date: Tue, 26 Sep 2023 12:12:39 +0100
Subject: [PATCH] Replace thread save/restore API with switch

---
 ocaml/otherlibs/systhreads/st_stubs.c | 32 +++++++++++++--------------
 ocaml/otherlibs/systhreads/threads.h  | 10 +++++----
 2 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/ocaml/otherlibs/systhreads/st_stubs.c b/ocaml/otherlibs/systhreads/st_stubs.c
index 6ea13152822..9c20a84e57b 100644
--- a/ocaml/otherlibs/systhreads/st_stubs.c
+++ b/ocaml/otherlibs/systhreads/st_stubs.c
@@ -243,9 +243,7 @@ static void memprof_ctx_iter(th_ctx_action f, void* data)
   } while (th != curr_thread);
 }
 
-/* Saving and restoring runtime state in curr_thread */
-
-CAMLexport void caml_thread_save_runtime_state(void)
+static void thread_save_runtime_state(void)
 {
   if (Caml_state->in_minor_collection)
     caml_fatal_error("Thread switch from inside minor GC");
@@ -273,12 +271,8 @@ CAMLexport void caml_thread_save_runtime_state(void)
   caml_memprof_leave_thread();
 }
 
-CAMLexport void caml_thread_restore_runtime_state(void)
+static void thread_restore_runtime_state(void)
 {
-  /* Update curr_thread to point to the thread descriptor corresponding
-     to the thread currently executing */
-  curr_thread = st_tls_get(thread_descriptor_key);
-
 #ifdef NATIVE_CODE
   Caml_state->top_of_stack = curr_thread->top_of_stack;
   Caml_state->bottom_of_stack= curr_thread->bottom_of_stack;
@@ -303,16 +297,24 @@ CAMLexport void caml_thread_restore_runtime_state(void)
   caml_memprof_enter_thread(curr_thread->memprof_ctx);
 }
 
+CAMLexport void caml_thread_switch_runtime_state(void)
+{
+  caml_thread_t new_thread = st_tls_get(thread_descriptor_key);
+  if (new_thread == curr_thread) return;
+  thread_save_runtime_state();
+  curr_thread = new_thread;
+  thread_restore_runtime_state();
+}
+
 CAMLexport void caml_switch_runtime_locking_scheme(struct caml_locking_scheme* new)
 {
   struct caml_locking_scheme* old;
 
-  caml_thread_save_runtime_state();
   old = atomic_exchange(&caml_locking_scheme, new);
   /* We hold 'old', but it is no longer the runtime lock */
   old->unlock(old->context);
   acquire_runtime_lock();
-  caml_thread_restore_runtime_state();
+  caml_thread_switch_runtime_state();
 }
 
 
@@ -321,9 +323,6 @@ CAMLexport void caml_switch_runtime_locking_scheme(struct caml_locking_scheme* n
 
 static void caml_thread_enter_blocking_section(void)
 {
-  /* Save the current runtime state in the thread descriptor
-     of the current thread */
-  caml_thread_save_runtime_state();
   /* Tell other threads that the runtime is free */
   release_runtime_lock();
 }
@@ -338,7 +337,7 @@ static void caml_thread_leave_blocking_section(void)
 #endif
   /* Wait until the runtime is free */
   acquire_runtime_lock();
-  caml_thread_restore_runtime_state();
+  caml_thread_switch_runtime_state();
 #ifdef _WIN32
   SetLastError(error);
 #endif
@@ -630,7 +629,7 @@ static void caml_thread_stop(void)
      changed as the thread was running, so we save it in the
      curr_thread data to make sure that the cleanup logic
      below uses accurate information. */
-  caml_thread_save_runtime_state();
+  thread_save_runtime_state();
   /* Tell memprof that this thread is terminating. */
   caml_memprof_delete_th_ctx(curr_thread->memprof_ctx);
   /* Signal that the thread has terminated */
@@ -881,7 +880,6 @@ CAMLprim value caml_thread_yield(value unit)        /* ML */
   */
   caml_raise_async_if_exception(caml_process_pending_signals_exn(),
                                 "signal handler");
-  caml_thread_save_runtime_state();
   /* caml_locking_scheme may have changed in caml_process_pending_signals_exn */
   s = atomic_load(&caml_locking_scheme);
   s->yield(s->context);
@@ -890,7 +888,7 @@ CAMLprim value caml_thread_yield(value unit)        /* ML */
     s->unlock(s->context);
     acquire_runtime_lock();
   }
-  caml_thread_restore_runtime_state();
+  caml_thread_switch_runtime_state();
   caml_raise_async_if_exception(caml_process_pending_signals_exn(),
                                 "signal handler");
 
diff --git a/ocaml/otherlibs/systhreads/threads.h b/ocaml/otherlibs/systhreads/threads.h
index df6f7aa0e83..f4b8ee1a749 100644
--- a/ocaml/otherlibs/systhreads/threads.h
+++ b/ocaml/otherlibs/systhreads/threads.h
@@ -102,11 +102,13 @@ CAMLextern_libthreads
 void caml_switch_runtime_locking_scheme(struct caml_locking_scheme*);
 
 CAMLextern_libthreads
-void caml_thread_save_runtime_state(void);
-
-CAMLextern_libthreads
-void caml_thread_restore_runtime_state(void);
+void caml_thread_switch_runtime_state(void);
 
+/* A prior version of this API used save/restore rather than a single switch.
+   For compatibility, aliases are defined for the old API.
+   (These will be removed when the lone user of this API is updated) */
+#define caml_thread_save_runtime_state()
+#define caml_thread_restore_runtime_state caml_thread_switch_runtime_state
 
 #ifdef __cplusplus
 }