WIP: Use Primitive.use_cext_lock? to find out if the C ext lock should be used

Author: eregon

doc/user/thread-safe-extensions.md

@@ -8,7 +8,7 @@ permalink: /reference-manual/ruby/ThreadSafeExtensions/
 
 Native extensions are by default considered thread-unsafe for maximum compatibility with CRuby and use the global extension lock (unless `--cexts-lock=false` is used).
 
-Extensions can mark themselves as thread-safe either by using `rb_ext_ractor_safe()` or `rb_ext_thread_safe()` (TruffleRuby-specific).
+Extensions can mark themselves as thread-safe either by using `rb_ext_ractor_safe()` or `rb_ext_thread_safe()` (the latter is TruffleRuby-specific).
 Such extensions are then run by TruffleRuby without a global extension lock, i.e. in parallel.
 
 Here is an example of an extension marking itself as Ractor-safe.
@@ -34,6 +34,27 @@ void Init_my_extension(void) {
 }
 ```
 
+It is possible to mark individual methods as Ractor/Thread-safe by using `rb_ext_ractor_safe/rb_ext_thread_safe(true/false)` around them (these functions actually set a Fiber-local flag):
+```c
+void Init_my_extension(void) {
+    #ifdef HAVE_RB_EXT_RACTOR_SAFE
+    rb_ext_ractor_safe(true);
+    #endif
+    rb_define_method(myClass, "ractor_safe_method", foo_impl, 0); // The C function foo_impl can be called from multiple threads in parallel
+    // more Ractor-safe methods
+
+    #ifdef HAVE_RB_EXT_RACTOR_SAFE
+    rb_ext_ractor_safe(false);
+    #endif
+    rb_define_method(myClass, "ractor_unsafe_method", bar_impl, 0); // The C function bar_impl needs a global extension lock for correctness
+    // more Ractor-unsafe methods
+}
+```
+
+Other Ruby C API functions taking a C function like `rb_proc_new()` do not use the global extension lock if:
+* Called inside the `Init_my_extension` and `rb_ext_ractor_safe(true)` / `rb_ext_thread_safe(true)` are used.
+* Called outside the `Init_my_extension` and the calling function does not hold the global extension lock.
+
 The conditions for an extension to be thread-safe are the following.
 This is similar to [the conditions for Ractor-safe extensions](https://github.com/ruby/ruby/blob/master/doc/extension.rdoc#appendix-f-ractor-support-) but not all conditions are necessary.
 1. The extension should make it clear in its documentation which objects are safe to share between threads and which are not.

lib/truffle/truffle/cext.rb

@@ -27,6 +27,8 @@ module Truffle::CExt
   SULONG = Truffle::Boot.get_option('cexts-sulong')
   NFI = !SULONG
 
+  CEXT_LOCK = Truffle::Boot.get_option('cexts-lock')
+
   SET_LIBTRUFFLERUBY = -> libtruffleruby do
     LIBTRUFFLERUBY = libtruffleruby
   end
@@ -243,7 +245,7 @@ def init_extension(library, library_path)
         # Resolve while inside the ExtensionCallStackEntry to ensure the preservedObjects are still all alive
         resolve_registered_addresses
       end
-    }, [], nil, nil)
+    }, [], nil, nil, CEXT_LOCK)
   end
 
   def supported?
@@ -856,12 +858,14 @@ def rb_str_new_frozen(value)
   end
 
   def rb_tracepoint_new(events, func, data)
+    use_cext_lock = Primitive.use_cext_lock?
     TracePoint.new(*events_to_events_array(events)) do |tp|
       Primitive.call_with_cext_lock_and_frame(
         POINTER2_TO_VOID_WRAPPER,
         [func, Primitive.cext_wrap(tp), data],
         Primitive.caller_special_variables_if_available,
-        nil)
+        nil,
+        use_cext_lock)
     end
   end
 
@@ -1204,6 +1208,7 @@ def rb_path_to_class(path)
   end
 
   def rb_proc_new(function, value)
+    use_cext_lock = Primitive.use_cext_lock?
     Proc.new do |*args, &block|
       Primitive.call_with_cext_lock_and_frame_and_unwrap(RB_BLOCK_CALL_FUNC_WRAPPER, [
         function,
@@ -1212,7 +1217,7 @@ def rb_proc_new(function, value)
         args.size, # argc
         Truffle::CExt.RARRAY_PTR(args), # argv
         Primitive.cext_wrap(block), # blockarg
-      ], Primitive.caller_special_variables_if_available, nil)
+      ], Primitive.caller_special_variables_if_available, nil, use_cext_lock)
     end
   end
 
@@ -1475,12 +1480,14 @@ def rb_enumeratorize(obj, meth, args)
 
   def rb_enumeratorize_with_size(obj, meth, args, size_fn)
     return rb_enumeratorize(obj, meth, args) if Primitive.interop_null?(size_fn)
+    use_cext_lock = Primitive.use_cext_lock?
     enum = obj.to_enum(meth, *args) do
       Primitive.call_with_cext_lock_and_frame_and_unwrap(
         POINTER3_TO_POINTER_WRAPPER,
         [size_fn, Primitive.cext_wrap(obj), Primitive.cext_wrap(args), Primitive.cext_wrap(enum)],
         Primitive.caller_special_variables_if_available,
-        nil)
+        nil,
+        use_cext_lock)
     end
     enum
   end
@@ -1494,12 +1501,14 @@ def rb_newobj_of(ruby_class)
   end
 
   def rb_define_alloc_func(ruby_class, function)
+    use_cext_lock = Primitive.use_cext_lock?
     ruby_class.singleton_class.define_method(:__allocate__) do
       Primitive.call_with_cext_lock_and_frame_and_unwrap(
         POINTER_TO_POINTER_WRAPPER,
         [function, Primitive.cext_wrap(self)],
         Primitive.caller_special_variables_if_available,
-        nil)
+        nil,
+        use_cext_lock)
     end
     class << ruby_class
       private :__allocate__
@@ -1697,10 +1706,11 @@ def rb_nativethread_lock_destroy(lock)
   end
 
   def rb_set_end_proc(func, data)
+    use_cext_lock = Primitive.use_cext_lock?
     at_exit do
       Primitive.call_with_cext_lock_and_frame(
         POINTER_TO_VOID_WRAPPER, [func, data],
-        Primitive.caller_special_variables_if_available, nil)
+        Primitive.caller_special_variables_if_available, nil, use_cext_lock)
     end
   end
 
@@ -1736,10 +1746,11 @@ def RTYPEDDATA(object)
 
   private def data_sizer(sizer_function, rtypeddata)
     raise unless sizer_function.respond_to?(:call)
+    use_cext_lock = Primitive.use_cext_lock?
     proc {
       Primitive.call_with_cext_lock_and_frame(
         POINTER_TO_SIZE_T_WRAPPER, [sizer_function, rtypeddata],
-        Primitive.caller_special_variables_if_available, nil)
+        Primitive.caller_special_variables_if_available, nil, use_cext_lock)
     }
   end
 
@@ -1779,7 +1790,9 @@ def rb_data_typed_object_wrap(ruby_class, data, data_type, mark, free, size)
   end
 
   def run_data_finalizer(function, data)
-    Primitive.call_with_cext_lock_and_frame POINTER_TO_VOID_WRAPPER, [function, data], nil, nil
+    # TODO: Maybe we should not hold the C extension lock here, and instead save the result of Primitive.use_cext_lock?
+    # in DataObjectFinalizerReference from rb_data_object_wrap/rb_data_typed_object_wrap.
+    Primitive.call_with_cext_lock_and_frame POINTER_TO_VOID_WRAPPER, [function, data], nil, nil, CEXT_LOCK
   end
 
   def run_marker(obj)
@@ -2006,8 +2019,9 @@ def rb_time_interval_acceptable(time_val)
   end
 
   def rb_thread_create(fn, args)
+    use_cext_lock = Primitive.use_cext_lock?
     Thread.new do
-      Primitive.call_with_cext_lock_and_frame(POINTER_TO_POINTER_WRAPPER, [fn, args], Primitive.caller_special_variables_if_available, nil)
+      Primitive.call_with_cext_lock_and_frame(POINTER_TO_POINTER_WRAPPER, [fn, args], Primitive.caller_special_variables_if_available, nil, use_cext_lock)
     end
   end
 
@@ -2035,6 +2049,7 @@ def rb_thread_call_without_gvl(function, data1, unblock, data2)
   end
 
   def rb_iterate(iteration, iterated_object, callback, callback_arg)
+    use_cext_lock = Primitive.use_cext_lock?
     block = rb_block_proc
     wrapped_callback = proc do |block_arg|
       Primitive.call_with_cext_lock_and_frame_and_unwrap(RB_BLOCK_CALL_FUNC_WRAPPER, [
@@ -2044,13 +2059,13 @@ def rb_iterate(iteration, iterated_object, callback, callback_arg)
         0, # argc
         nil, # argv
         nil, # blockarg
-      ], Primitive.cext_special_variables_from_stack, block)
+      ], Primitive.cext_special_variables_from_stack, block, use_cext_lock)
     end
     Primitive.cext_unwrap(
       Primitive.call_with_cext_lock_and_frame(POINTER_TO_POINTER_WRAPPER, [
         iteration,
         Primitive.cext_wrap(iterated_object)
-      ], Primitive.cext_special_variables_from_stack, wrapped_callback))
+      ], Primitive.cext_special_variables_from_stack, wrapped_callback, use_cext_lock))
   end
 
   # From ruby.h
@@ -2147,21 +2162,24 @@ def rb_hash_aref(object, key)
   def rb_define_hooked_variable(name, gvar, getter, setter)
     name = "$#{name}" unless name.start_with?('$')
     id = name.to_sym
+    use_cext_lock = Primitive.use_cext_lock?
 
     getter_proc = -> {
       Primitive.call_with_cext_lock_and_frame_and_unwrap(
         POINTER2_TO_POINTER_WRAPPER,
         [getter, Primitive.cext_wrap(id), gvar],
         Primitive.caller_special_variables_if_available,
-        nil)
+        nil,
+        use_cext_lock)
     }
 
     setter_proc = -> value {
       Primitive.call_with_cext_lock_and_frame(
         POINTER3_TO_VOID_WRAPPER,
         [setter, Primitive.cext_wrap(value), Primitive.cext_wrap(id), gvar],
         Primitive.caller_special_variables_if_available,
-        nil)
+        nil,
+        use_cext_lock)
     }
 
     Truffle::KernelOperations.define_hooked_variable id, getter_proc, setter_proc
@@ -2285,6 +2303,7 @@ def rb_fiber_current
   end
 
   def rb_fiber_new(function, value)
+    use_cext_lock = Primitive.use_cext_lock?
     Fiber.new do |*args|
       Primitive.call_with_cext_lock_and_frame_and_unwrap(RB_BLOCK_CALL_FUNC_WRAPPER, [
         function,
@@ -2293,7 +2312,7 @@ def rb_fiber_new(function, value)
         0, # argc
         nil, # argv
         nil, # blockarg
-      ], nil, nil)
+      ], nil, nil, use_cext_lock)
     end
   end
 

lib/truffle/truffle/cext_ruby.rb

@@ -21,7 +21,7 @@ def rb_define_method(mod, name, function, argc)
     end
 
     wrapper = RB_DEFINE_METHOD_WRAPPERS[argc]
-    thread_safe = Primitive.cext_thread_safe?
+    use_cext_lock = Primitive.use_cext_lock?
     method_body = Truffle::Graal.copy_captured_locals -> *args, &block do
       if argc == -1 # (int argc, VALUE *argv, VALUE obj)
         args = [function, args.size, Truffle::CExt.RARRAY_PTR(args), Primitive.cext_wrap(self)]
@@ -36,10 +36,10 @@ def rb_define_method(mod, name, function, argc)
 
       # We must set block argument if given here so that the
       # `rb_block_*` functions will be able to find it by walking the stack.
-      if thread_safe
-        Primitive.call_with_frame_and_unwrap(wrapper, args, Primitive.caller_special_variables_if_available, block)
-      else
+      if use_cext_lock
         Primitive.call_with_cext_lock_and_frame_and_unwrap(wrapper, args, Primitive.caller_special_variables_if_available, block)
+      else
+        Primitive.call_with_frame_and_unwrap(wrapper, args, Primitive.caller_special_variables_if_available, block)
       end
     end
 

src/main/java/org/truffleruby/cext/CExtNodes.java

@@ -183,7 +183,12 @@ public abstract static class CallWithCExtLockAndFrameNode extends PrimitiveArray
 
         @Specialization
         Object callWithCExtLockAndFrame(
-                VirtualFrame frame, Object receiver, RubyArray argsArray, Object specialVariables, Object block,
+                VirtualFrame frame,
+                Object receiver,
+                RubyArray argsArray,
+                Object specialVariables,
+                Object block,
+                boolean useCExtLock,
                 @CachedLibrary(limit = "getCacheLimit()") InteropLibrary receivers,
                 @Cached ArrayToObjectArrayNode arrayToObjectArrayNode,
                 @Cached TranslateInteropExceptionNode translateInteropExceptionNode,
@@ -197,7 +202,7 @@ Object callWithCExtLockAndFrame(
             try {
                 final Object[] args = arrayToObjectArrayNode.executeToObjectArray(argsArray);
 
-                if (getContext().getOptions().CEXT_LOCK) {
+                if (getContext().getOptions().CEXT_LOCK && useCExtLock) {
                     final ReentrantLock lock = getContext().getCExtensionsLock();
                     boolean owned = ownedProfile.profile(this, lock.isHeldByCurrentThread());
 
@@ -499,6 +504,18 @@ boolean isCExtLockOwned() {
         }
     }
 
+    @CoreMethod(names = "use_cext_lock?", onSingleton = true)
+    public abstract static class UseCExtLockNode extends PrimitiveArrayArgumentsNode {
+        @Specialization
+        boolean useCExtLock() {
+            if (getLanguage().getCurrentFiber().threadSafeExtension) {
+                return false;
+            } else {
+                return getContext().getCExtensionsLock().isHeldByCurrentThread();
+            }
+        }
+    }
+
     @CoreMethod(names = "rb_ulong2num", onSingleton = true, required = 1)
     public abstract static class ULong2NumNode extends CoreMethodArrayArgumentsNode {
 
@@ -2147,12 +2164,4 @@ Object setThreadSafe(int threadSafe) {
         }
     }
 
-    @Primitive(name = "cext_thread_safe?")
-    public abstract static class CExtIsThreadSafe extends PrimitiveArrayArgumentsNode {
-        @Specialization
-        boolean isThreadSafe() {
-            return getLanguage().getCurrentFiber().threadSafeExtension;
-        }
-    }
-
 }

src/main/java/org/truffleruby/core/fiber/RubyFiber.java

@@ -104,6 +104,7 @@ public enum FiberStatus {
     public final ValueWrapperManager.HandleBlockHolder handleData;
     boolean blocking = true;
     public RubyArray cGlobalVariablesDuringInitFunction;
+    /** Always false when not inside a C ext Init_ function */
     public boolean threadSafeExtension = false;
 
     // To pass state between beforeEnter(), fiberMain() and afterLeave()