Add context inheritance via prepend instead of changing Thread

lib/logging/diagnostic_context.rb

@@ -412,35 +412,9 @@ def self.clear_diagnostic_contexts( all = false )
   end
 
   DIAGNOSTIC_MUTEX = Mutex.new
-end
-
-# :stopdoc:
-Logging::INHERIT_CONTEXT =
-  if ENV.key?("LOGGING_INHERIT_CONTEXT")
-    case ENV["LOGGING_INHERIT_CONTEXT"].downcase
-    when 'false', 'no', '0'; false
-    when false, nil; false
-    else true end
-  else
-    true
-  end
-
-if Logging::INHERIT_CONTEXT
-  class Thread
-    class << self
-
-      %w[new start fork].each do |m|
-        class_eval <<-__, __FILE__, __LINE__
-          alias_method :_orig_#{m}, :#{m}
-          private :_orig_#{m}
-          def #{m}( *a, &b )
-            create_with_logging_context(:_orig_#{m}, *a ,&b)
-          end
-        __
-      end
-
-    private
 
+  module InheritContext
+    %w[new start fork].each do |m|
       # In order for the diagnostic contexts to behave properly we need to
       # inherit state from the parent thread. The only way I have found to do
       # this in Ruby is to override `new` and capture the contexts from the
@@ -453,30 +427,49 @@ def #{m}( *a, &b )
       # being executed in the child thread. The same is true for the parent
       # thread's mdc and ndc. If any of those references end up in the binding,
       # then they cannot be garbage collected until the child thread exits.
-      #
-      def create_with_logging_context( m, *a, &b )
-        mdc, ndc = nil
+      class_eval <<-__, __FILE__, __LINE__
+        def #{m}( *args )
+          mdc, ndc = nil
 
-        if Thread.current.thread_variable_get(Logging::MappedDiagnosticContext::STACK_NAME)
-          mdc = Logging::MappedDiagnosticContext.context.dup
-        end
+          if Thread.current.thread_variable_get(Logging::MappedDiagnosticContext::STACK_NAME)
+            mdc = Logging::MappedDiagnosticContext.context.dup
+          end
 
-        if Thread.current.thread_variable_get(Logging::NestedDiagnosticContext::NAME)
-          ndc = Logging::NestedDiagnosticContext.context.dup
-        end
+          if Thread.current.thread_variable_get(Logging::NestedDiagnosticContext::NAME)
+            ndc = Logging::NestedDiagnosticContext.context.dup
+          end
 
-        # This calls the actual `Thread#new` method to create the Thread instance.
-        # If your memory profiling tool says this method is leaking memory, then
-        # you are leaking Thread instances somewhere.
-        self.send(m, *a) { |*args|
-          Logging::MappedDiagnosticContext.inherit(mdc)
-          Logging::NestedDiagnosticContext.inherit(ndc)
-          b.call(*args)
-        }
-      end
+          wrapped_block = proc { |*thread_args|
+            Logging::MappedDiagnosticContext.inherit(mdc)
+            Logging::NestedDiagnosticContext.inherit(ndc)
 
+            yield(*thread_args)
+          }
+          wrapped_block.ruby2_keywords if wrapped_block.respond_to?(:ruby2_keywords, true)
+
+          # This calls the actual method to create the Thread instance.
+          # If your memory profiling tool says this method is leaking memory, then
+          # you are leaking Thread instances somewhere.
+          super(*args, &wrapped_block)
+        end
+        ruby2_keywords :#{m} if respond_to?(:ruby2_keywords, true)
+      __
     end
   end
 end
-# :startdoc:
 
+# :stopdoc:
+Logging::INHERIT_CONTEXT =
+  if ENV.key?("LOGGING_INHERIT_CONTEXT")
+    case ENV["LOGGING_INHERIT_CONTEXT"].downcase
+    when 'false', 'no', '0'; false
+    when false, nil; false
+    else true end
+  else
+    true
+  end
+
+if Logging::INHERIT_CONTEXT
+  ::Thread.singleton_class.prepend(Logging::InheritContext)
+end
+# :startdoc:

test/test_mapped_diagnostic_context.rb

@@ -131,5 +131,13 @@ def test_thread_inheritance
       t.join
     end
 
+    def test_thread_patches_work_correctly_on_ruby3
+      the_args, the_kwargs = Thread.new(1, 2, 3, four: 4, five: 5) { |*args, **kwargs|
+        [args, kwargs]
+      }.value
+
+      assert_equal [1, 2, 3], the_args
+      assert_equal ({four: 4, five: 5}), the_kwargs
+    end
   end  # class TestMappedDiagnosticContext
 end  # module TestLogging