From 599550f8e732347f59237ee9c60a11426b397103 Mon Sep 17 00:00:00 2001
From: Sutou Kouhei <kou@clear-code.com>
Date: Fri, 25 Apr 2025 07:01:39 +0900
Subject: [PATCH] Avoid use of id2ref for weak mapping

[Bugs #15711]

https://bugs.ruby-lang.org/issues/15711

The implementation of `ObjectSpace._id2ref` on JRuby (and TruffleRuby)
is very expensive, and on JRuby we normally do not have it enabled
because it impacts performance of the entire runtime.

This uses `ObjectSpace::WeakMap` instead of `ObjectSpace._id2ref` by
default and deprecates `DRb::WeakIdConv`.

Authored-by: Masatoshi SEKI <m_seki@mva.biglobe.ne.jp>
---
 lib/drb/drb.rb             | 42 +++++++++++++++++++++-----
 lib/drb/weakidconv.rb      | 62 +++++++++++---------------------------
 test/drb/test_drbobject.rb |  4 ++-
 3 files changed, 55 insertions(+), 53 deletions(-)

diff --git a/lib/drb/drb.rb b/lib/drb/drb.rb
index 2c9997c..af35745 100644
--- a/lib/drb/drb.rb
+++ b/lib/drb/drb.rb
@@ -348,6 +348,39 @@ class DRbError < RuntimeError; end
   # protocol.
   class DRbConnError < DRbError; end
 
+  class DRbObjectSpace # :nodoc:
+    # This is an internal class for DRbIdConv. This must not be used
+    # by users.
+
+    include MonitorMixin
+
+    def initialize
+      super()
+      @map = ObjectSpace::WeakMap.new
+    end
+
+    def to_id(obj)
+      synchronize do
+        @map[obj.__id__] = obj
+        obj.__id__
+      end
+    end
+
+    def to_obj(ref)
+      synchronize do
+        obj = @map[ref]
+        raise RangeError.new("invalid reference") unless obj.__id__ == ref
+        obj
+      end
+    end
+  end
+
+  # :nodoc:
+  #
+  # This is an internal singleton instance. This must not be used
+  # by users.
+  DRB_OBJECT_SPACE = DRbObjectSpace.new
+
   # Class responsible for converting between an object and its id.
   #
   # This, the default implementation, uses an object's local ObjectSpace
@@ -364,7 +397,7 @@ class DRbIdConv
     # This implementation looks up the reference id in the local object
     # space and returns the object it refers to.
     def to_obj(ref)
-      ObjectSpace._id2ref(ref)
+      DRB_OBJECT_SPACE.to_obj(ref)
     end
 
     # Convert an object into a reference id.
@@ -372,12 +405,7 @@ def to_obj(ref)
     # 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__
-      end
+      (nil == obj) ? nil : DRB_OBJECT_SPACE.to_id(obj)
     end
   end
 
diff --git a/lib/drb/weakidconv.rb b/lib/drb/weakidconv.rb
index ecf0bf5..951af8a 100644
--- a/lib/drb/weakidconv.rb
+++ b/lib/drb/weakidconv.rb
@@ -1,57 +1,29 @@
 # frozen_string_literal: false
 require_relative 'drb'
-require 'monitor'
 
 module DRb
 
+  # DRb::WeakIdConv is deprecated since 2.2.1. You don't need to use
+  # DRb::WeakIdConv instead of DRb::DRbIdConv. It's the same class.
+  #
+  # This file still exists for backward compatibility.
+  #
   # To use WeakIdConv:
   #
   #  DRb.start_service(nil, nil, {:idconv => DRb::WeakIdConv.new})
 
-  class WeakIdConv < DRbIdConv
-    class WeakSet
-      include MonitorMixin
-      def initialize
-        super()
-        @immutable = {}
-        @map = ObjectSpace::WeakMap.new
-      end
-
-      def add(obj)
-        synchronize do
-          begin
-            @map[obj] = self
-          rescue ArgumentError
-            @immutable[obj.__id__] = obj
-          end
-          return obj.__id__
-        end
-      end
-
-      def fetch(ref)
-        synchronize do
-          @immutable.fetch(ref) {
-            @map.each { |key, _|
-              return key if key.__id__ == ref
-            }
-            raise RangeError.new("invalid reference")
-          }
-        end
-      end
-    end
-
-    def initialize()
-      super()
-      @weak_set = WeakSet.new
-    end
-
-    def to_obj(ref) # :nodoc:
-      return super if ref.nil?
-      @weak_set.fetch(ref)
-    end
-
-    def to_id(obj) # :nodoc:
-      return @weak_set.add(obj)
+  def self.const_missing(name) # :nodoc:
+    case name
+    when :WeakIdConv
+      warn("DRb::WeakIdConv is deprecated. " +
+           "You can use the DRb::DRbIdConv. " +
+           "You don't need to use this.",
+           uplevel: 1)
+      const_set(:WeakIdConv, DRbIdConv)
+      singleton_class.remove_method(:const_missing)
+      DRbIdConv
+    else
+      super
     end
   end
 end
diff --git a/test/drb/test_drbobject.rb b/test/drb/test_drbobject.rb
index 2b0e206..a79dd89 100644
--- a/test/drb/test_drbobject.rb
+++ b/test/drb/test_drbobject.rb
@@ -1,4 +1,5 @@
 require 'test/unit'
+require 'envutil'
 require 'drb'
 require 'drb/timeridconv'
 require 'drb/weakidconv'
@@ -64,6 +65,7 @@ class TestDRbObjectWeakIdConv < Test::Unit::TestCase
   include DRbObjectTest
 
   def setup
-    DRb.start_service(nil, nil, {:idconv => DRb::WeakIdConv.new})
+    idconv = EnvUtil.suppress_warning {DRb::WeakIdConv.new}
+    DRb.start_service(nil, nil, {:idconv => idconv})
   end
 end