Give each Ractor its own WeakMap

lib/weakref.rb

@@ -15,6 +15,9 @@
 #   GC.start                    # start the garbage collector
 #   p foo.to_s                  # should raise exception (recycled)
 #
+# A WeakRef belongs to the Ractor that created it: the referenced object is
+# usually not shareable, so it may not be reached from another Ractor.
+# Consequently a WeakRef is neither shareable nor movable.
 
 class WeakRef < Delegator
   # The version string
@@ -27,23 +30,51 @@ class WeakRef < Delegator
   class RefError < StandardError
   end
 
-  @@__map = ::ObjectSpace::WeakMap.new
+  # ObjectSpace::WeakMap is neither shareable nor thread safe, and its values
+  # are the referenced objects, which are not shareable either. So each Ractor
+  # gets its own map instead of one map shared by the whole process.
+  if defined?(::Ractor) && ::Ractor.respond_to?(:store_if_absent)
+    def self.__map__ # :nodoc:
+      ::Ractor.store_if_absent(:__weakref_map__) { ::ObjectSpace::WeakMap.new }
+    end
+  elsif defined?(::Ractor)
+    # Ractor.store_if_absent is Ruby 3.4 and later. Two threads of the same
+    # Ractor can race here and each build a map, but that only leaves one empty
+    # map behind: every WeakRef keeps the map it registered itself in.
+    def self.__map__ # :nodoc:
+      current = ::Ractor.current
+      current[:__weakref_map__] ||= ::ObjectSpace::WeakMap.new
+    end
+  else
+    @__map = ::ObjectSpace::WeakMap.new
+
+    def self.__map__ # :nodoc:
+      @__map
+    end
+  end
 
   ##
   # Creates a weak reference to +orig+
 
   def initialize(orig)
     case orig
     when true, false, nil
+      @map = nil
       @delegate_sd_obj = orig
     else
-      @@__map[self] = orig
+      # Holding the map in an instance variable keeps lookups to a single ivar
+      # read, and makes the WeakRef itself non-shareable and non-movable, which
+      # is what we want: its entry lives in this Ractor's map only.
+      # Delegator does not inherit from Object, so a bare constant here would go
+      # through Delegator.const_missing on every call.
+      @map = ::WeakRef.__map__
+      @map[self] = orig
     end
     super
   end
 
   def __getobj__(&_block) # :nodoc:
-    @@__map[self] or defined?(@delegate_sd_obj) ? @delegate_sd_obj :
+    @map && @map[self] or defined?(@delegate_sd_obj) ? @delegate_sd_obj :
       Kernel::raise(RefError, "Invalid Reference - probably recycled", Kernel::caller(2))
   end
 
@@ -54,6 +85,6 @@ def __setobj__(obj) # :nodoc:
   # Returns true if the referenced object is still alive.
 
   def weakref_alive?
-    @@__map.key?(self) or defined?(@delegate_sd_obj)
+    (@map ? @map.key?(self) : false) or defined?(@delegate_sd_obj)
   end
 end

test/test_weakref.rb

@@ -69,4 +69,41 @@ def test_repeated_object_memory_leak
       150_000.times { WeakRef.new(a) }
     end;
   end
+
+  if defined?(Ractor)
+    def test_weakref_in_ractor
+      bug22105 = '[ruby-core:125705]'
+      assert_separately(%w(-rweakref), <<-'end;', ignore_stderr: true)
+        Warning[:experimental] = false
+        ractor = Ractor.new do
+          str = "referenced"
+          ref = WeakRef.new(str)
+          [ref.__getobj__, ref.weakref_alive?]
+        end
+        # Ractor#value is 4.0 and later, Ractor#take before that.
+        obj, alive = ractor.respond_to?(:value) ? ractor.value : ractor.take
+        assert_equal("referenced", obj)
+        assert_equal(true, alive)
+      end;
+    end
+
+    def test_map_is_per_ractor
+      assert_separately(%w(-rweakref), <<-'end;', ignore_stderr: true)
+        Warning[:experimental] = false
+        ractor = Ractor.new { WeakRef.__map__.object_id }
+        other = ractor.respond_to?(:value) ? ractor.value : ractor.take
+        assert_not_equal(WeakRef.__map__.object_id, other)
+      end;
+    end
+
+    def test_weakref_is_not_shareable
+      # The referenced object belongs to the Ractor that created the WeakRef,
+      # so the WeakRef must not escape it.
+      assert_separately(%w(-rweakref), <<-'end;', ignore_stderr: true)
+        Warning[:experimental] = false
+        ref = WeakRef.new(Object.new)
+        assert_raise(Ractor::Error) { Ractor.make_shareable(ref) }
+      end;
+    end
+  end
 end