From 2b393cdde12e303cab71aee0966feff34c5075e0 Mon Sep 17 00:00:00 2001
From: Charles Oliver Nutter <headius@headius.com>
Date: Thu, 18 Nov 2021 06:00:16 -0600
Subject: [PATCH] Use a weak equality map instead of _id2ref

The ObjectSpace._id2ref method is considered an internal API of
CRuby and is difficult to implement efficiently on implementations
that do not have direct control over the garbage-collected heap.
On JRuby, for example, arbitrary _id2ref object lookup is normally
disabled, as the implementation requires maintaining a parallel
mapping of object IDs to weak object references, adding a very
large amount of overhead to all object ID uses. As a result, Ruby
standard libraries should try to avoid using _id2ref.

This patch introduces a weak map into DRb for managing object
references without _id2ref. The map is weak-valued, with a clean
subroutine to scrub out defunct entries. Puts initiate a clean,
hopefully ensuring that very few stale entries accumulate.

This patch does not use WeakMap do to limitations of its API.
WeakMap specifies that keys are identity-based, and as of Ruby 3.0
allows immediate values like Symbols and Integers. In order to
support this behavior, those values must be idempotent, which on
JRuby is impossible to support for fixnum-ranged Integers and
flonum-ranged Floats, rendering it useless as an object ID-keyed
cache.
---
 lib/drb/drb.rb | 65 ++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 53 insertions(+), 12 deletions(-)

diff --git a/lib/drb/drb.rb b/lib/drb/drb.rb
index 3e23213..8fc02af 100644
--- a/lib/drb/drb.rb
+++ b/lib/drb/drb.rb
@@ -49,6 +49,7 @@
 require 'socket'
 require 'io/wait'
 require 'monitor'
+require 'weakref'
 require_relative 'eq'
 
 #
@@ -357,27 +358,67 @@ class DRbConnError < DRbError; end
   # For alternative mechanisms, see DRb::TimerIdConv in drb/timeridconv.rb
   # and DRbNameIdConv in sample/name.rb in the full drb distribution.
   class DRbIdConv
+    MUTEX = Mutex.new
+
+    def initialize
+      @id2ref = {}
+    end
 
     # Convert an object reference id to an object.
-    #
-    # This implementation looks up the reference id in the local object
-    # space and returns the object it refers to.
+    # and returns the object it refers to.
     def to_obj(ref)
-      ObjectSpace._id2ref(ref)
+      _get(ref)
     end
 
     # Convert an object into a reference id.
-    #
-    # This implementation returns the object's __id__ in the local
-    # object space.
     def to_id(obj)
-      case obj
-      when Object
-        obj.nil? ? nil : obj.__id__
-      when BasicObject
-        obj.__id__
+      obj.nil? ? nil : _put(obj)
+    end
+
+    private def _safe_lock
+      Thread.handle_interrupt(Exception => :never) do
+        MUTEX.lock
+        begin
+          yield
+        ensure
+          MUTEX.unlock
+        end
       end
     end
+
+    private def _put(obj)
+      id = obj.__id__
+      _safe_lock do
+        _clean_locked
+        @id2ref[id] = WeakRef.new(obj)
+      end
+      id
+    end
+
+    private def _get(id)
+      _safe_lock do
+        weakref = @id2ref[id]
+        if weakref
+          result = weakref.__getobj__ rescue nil
+          if result
+            return result
+          else
+            @id2ref.delete id
+          end
+        end
+      end
+      nil
+    end
+
+    private def _clean
+      _safe_lock { _clean_locked }
+    end
+
+    private def _clean_locked
+      dead = []
+      @id2ref.each { |id, weakref| dead << id unless weakref.weakref_alive? }
+      dead.each { |id| @id2ref.delete(id) }
+    end
   end
 
   # Mixin module making an object undumpable or unmarshallable.