From 55f2bc67dd703625cb72fefa4c5a914431639a6c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleksey=20Kliger=20=28=CE=BBgeek=29?= <alklig@microsoft.com>
Date: Thu, 27 Jun 2024 15:39:59 -0400
Subject: [PATCH] [mini] Use atomics instead of loader lock for JIT wrappers
 and trampolines (#104038)

* [mini] Use atomics, instead of loader lock, for JIT wrappers

   Related to https://github.com/dotnet/runtime/issues/93686

   While this doesn't eliminate all deadlocks related to the global loader lock and managed locks, it removes one unneeded use of the loader lock.

   The wrapper (and trampoline) of a JIT icall are only ever set from NULL to non-NULL.  We can use atomics to deal with races instead of double checked locking.  This was not the case historically, because the JIT info was dynamically allocated - so we used the loader lock to protect the integrity of the hash table
---
 src/mono/mono/metadata/class-internals.h |  2 ++
 src/mono/mono/metadata/icall.c           |  3 +++
 src/mono/mono/mini/mini-runtime.c        | 25 +++++++-----------------
 3 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/src/mono/mono/metadata/class-internals.h b/src/mono/mono/metadata/class-internals.h
index b423350d29810..a6a908e8bab10 100644
--- a/src/mono/mono/metadata/class-internals.h
+++ b/src/mono/mono/metadata/class-internals.h
@@ -588,7 +588,9 @@ typedef struct {
 	// have both of them to be non-NULL.
 	const char *name;
 	gconstpointer func;
+	// use CAS to write
 	gconstpointer wrapper;
+	// use CAS to write
 	gconstpointer trampoline;
 	MonoMethodSignature *sig;
 	const char *c_symbol;
diff --git a/src/mono/mono/metadata/icall.c b/src/mono/mono/metadata/icall.c
index e1d9bacb5f7ee..67cd6470df19d 100644
--- a/src/mono/mono/metadata/icall.c
+++ b/src/mono/mono/metadata/icall.c
@@ -7205,6 +7205,8 @@ mono_create_icall_signatures (void)
 	}
 }
 
+/* LOCKING: does not take locks. does not use an atomic write to info->wrapper
+ */
 void
 mono_register_jit_icall_info (MonoJitICallInfo *info, gconstpointer func, const char *name, MonoMethodSignature *sig, gboolean avoid_wrapper, const char *c_symbol)
 {
@@ -7218,6 +7220,7 @@ mono_register_jit_icall_info (MonoJitICallInfo *info, gconstpointer func, const
 	// Fill in wrapper ahead of time, to just be func, to avoid
 	// later initializing it to anything else. So therefore, no wrapper.
 	if (avoid_wrapper) {
+		// not using CAS, because its idempotent
 		info->wrapper = func;
 	} else {
 		// Leave it alone in case of a race.
diff --git a/src/mono/mono/mini/mini-runtime.c b/src/mono/mono/mini/mini-runtime.c
index e99e464fb8c85..240aa1f8a9b8b 100644
--- a/src/mono/mono/mini/mini-runtime.c
+++ b/src/mono/mono/mini/mini-runtime.c
@@ -663,8 +663,8 @@ mono_icall_get_wrapper_full (MonoJitICallInfo* callinfo, gboolean do_compile)
 		addr = mono_compile_method_checked (wrapper, error);
 		mono_error_assert_ok (error);
 		mono_memory_barrier ();
-		callinfo->wrapper = addr;
-		return addr;
+		mono_atomic_cas_ptr ((volatile gpointer*)&callinfo->wrapper, (gpointer)addr, NULL);
+		return (gconstpointer)mono_atomic_load_ptr((volatile gpointer*)&callinfo->wrapper);
 	} else {
 		if (callinfo->trampoline)
 			return callinfo->trampoline;
@@ -672,13 +672,10 @@ mono_icall_get_wrapper_full (MonoJitICallInfo* callinfo, gboolean do_compile)
 		mono_error_assert_ok (error);
 		trampoline = mono_create_ftnptr ((gpointer)trampoline);
 
-		mono_loader_lock ();
-		if (!callinfo->trampoline) {
-			callinfo->trampoline = trampoline;
-		}
-		mono_loader_unlock ();
+		
+		mono_atomic_cas_ptr ((volatile gpointer*)&callinfo->trampoline, (gpointer)trampoline, NULL);
 
-		return callinfo->trampoline;
+		return (gconstpointer)mono_atomic_load_ptr ((volatile gpointer*)&callinfo->trampoline);
 	}
 }
 
@@ -2856,18 +2853,10 @@ mono_jit_compile_method_with_opt (MonoMethod *method, guint32 opt, gboolean jit_
 	p = mono_create_ftnptr (code);
 
 	if (callinfo) {
-		// FIXME Locking here is somewhat historical due to mono_register_jit_icall_wrapper taking loader lock.
-		// atomic_compare_exchange should suffice.
-		mono_loader_lock ();
-		mono_jit_lock ();
-		if (!callinfo->wrapper) {
-			callinfo->wrapper = p;
-		}
-		mono_jit_unlock ();
-		mono_loader_unlock ();
+		mono_atomic_cas_ptr ((volatile void*)&callinfo->wrapper, p, NULL);
+		p = mono_atomic_load_ptr((volatile gpointer*)&callinfo->wrapper);
 	}
 
-	// FIXME p or callinfo->wrapper or does not matter?
 	return p;
 }