Merge pull request #3834 from DataDog/ivoanjo/prof-10241-simplify-at-…

…fork-monkey-patch

[PROF-10241] Fix issue in tests by simplifying `at_fork` monkey patch

lib/datadog/core/utils/at_fork_monkey_patch.rb

@@ -3,9 +3,12 @@
 module Datadog
   module Core
     module Utils
-      # Monkey patches `Kernel#fork` and similar functions, adding a `Process#datadog_at_fork` callback mechanism which
+      # Monkey patches `Kernel#fork` and similar functions, adding an `at_fork` callback mechanism which
       # is used to restart observability after the VM forks (e.g. in multiprocess Ruby apps).
       module AtForkMonkeyPatch
+        AT_FORK_CHILD_BLOCKS = [] # rubocop:disable Style/MutableConstant Used to store blocks to run, mutable by design.
+        private_constant :AT_FORK_CHILD_BLOCKS
+
         def self.supported?
           Process.respond_to?(:fork)
         end
@@ -28,14 +31,28 @@ def self.apply!
           true
         end
 
-        # Adds `datadog_at_fork` behavior; see parent module for details.
+        def self.run_at_fork_blocks(stage)
+          raise(ArgumentError, "Unsupported stage #{stage}") unless stage == :child
+
+          AT_FORK_CHILD_BLOCKS.each(&:call)
+        end
+
+        def self.at_fork(stage, &block)
+          raise(ArgumentError, "Unsupported stage #{stage}") unless stage == :child
+          raise(ArgumentError, 'Missing block argument') unless block
+
+          AT_FORK_CHILD_BLOCKS << block
+
+          true
+        end
+
+        # Adds `at_fork` behavior; see parent module for details.
         module KernelMonkeyPatch
           def fork
             # If a block is provided, it must be wrapped to trigger callbacks.
             child_block = if block_given?
                             proc do
-                              # Trigger :child callback
-                              datadog_at_fork_blocks[:child].each(&:call) if datadog_at_fork_blocks.key?(:child)
+                              AtForkMonkeyPatch.run_at_fork_blocks(:child)
 
                               # Invoke original block
                               yield
@@ -50,30 +67,20 @@ def fork
             # If we're in the fork, result = nil: trigger child callbacks.
             # If we're in the parent, result = pid: we do nothing.
             # (If it gets called with a block, it only returns on the parent)
-            datadog_at_fork_blocks[:child].each(&:call) if result.nil? && datadog_at_fork_blocks.key?(:child)
+            AtForkMonkeyPatch.run_at_fork_blocks(:child) if result.nil?
 
             result
           end
-
-          module_function
-
-          def datadog_at_fork_blocks
-            # Blocks are shared across all users of this module,
-            # e.g. Process#fork, Kernel#fork, etc. should all invoke the same callbacks.
-            @@datadog_at_fork_blocks ||= {} # rubocop:disable Style/ClassVars
-          end
         end
 
-        # Adds `datadog_at_fork` behavior; see parent module for details.
+        # Adds `at_fork` behavior; see parent module for details.
         module ProcessMonkeyPatch
           # Hook provided by Ruby 3.1+ for observability libraries that want to know about fork, see
           # https://github.com/ruby/ruby/pull/5017 and https://bugs.ruby-lang.org/issues/17795
           def _fork
-            datadog_at_fork_blocks = Datadog::Core::Utils::AtForkMonkeyPatch::KernelMonkeyPatch.datadog_at_fork_blocks
-
             pid = super
 
-            datadog_at_fork_blocks[:child].each(&:call) if pid == 0 && datadog_at_fork_blocks.key?(:child)
+            AtForkMonkeyPatch.run_at_fork_blocks(:child) if pid == 0
 
             pid
           end
@@ -82,31 +89,12 @@ def _fork
           # keeps executing code in the child process, killing off the parent, thus effectively replacing it.
           # This is not covered by `_fork` and thus we have some extra code for it.
           def daemon(*args)
-            datadog_at_fork_blocks = Datadog::Core::Utils::AtForkMonkeyPatch::KernelMonkeyPatch.datadog_at_fork_blocks
-
             result = super
 
-            datadog_at_fork_blocks[:child].each(&:call) if datadog_at_fork_blocks.key?(:child)
+            AtForkMonkeyPatch.run_at_fork_blocks(:child)
 
             result
           end
-
-          # NOTE: You probably want to wrap any calls to datadog_at_fork with a OnlyOnce so as to not re-register
-          #       the same block/behavior more than once.
-          def datadog_at_fork(stage, &block)
-            ProcessMonkeyPatch.datadog_at_fork(stage, &block)
-          end
-
-          # Also allow calling without going through Process for tests
-          def self.datadog_at_fork(stage, &block)
-            raise ArgumentError, 'Bad \'stage\' for ::datadog_at_fork' unless stage == :child
-
-            datadog_at_fork_blocks = Datadog::Core::Utils::AtForkMonkeyPatch::KernelMonkeyPatch.datadog_at_fork_blocks
-            datadog_at_fork_blocks[stage] ||= []
-            datadog_at_fork_blocks[stage] << block
-
-            nil
-          end
         end
       end
     end

lib/datadog/profiling/tasks/setup.rb

@@ -27,20 +27,16 @@ def run
         private
 
         def setup_at_fork_hooks
-          if Process.respond_to?(:datadog_at_fork)
-            Process.datadog_at_fork(:child) do
-              begin
-                # Restart profiler, if enabled
-                Profiling.start_if_enabled
-              rescue StandardError => e
-                Datadog.logger.warn do
-                  "Error during post-fork hooks. Cause: #{e.class.name} #{e.message} " \
-                  "Location: #{Array(e.backtrace).first}"
-                end
+          Datadog::Core::Utils::AtForkMonkeyPatch.at_fork(:child) do
+            begin
+              # Restart profiler, if enabled
+              Profiling.start_if_enabled
+            rescue StandardError => e
+              Datadog.logger.warn do
+                "Error during post-fork hooks. Cause: #{e.class.name} #{e.message} " \
+                "Location: #{Array(e.backtrace).first}"
               end
             end
-          else
-            Datadog.logger.debug 'Unexpected: At fork hooks not available'
           end
         end
       end

spec/datadog/core/utils/at_fork_monkey_patch_spec.rb

@@ -6,15 +6,15 @@
   describe '::apply!' do
     subject(:apply!) { described_class.apply! }
 
-    let(:toplevel_receiver) { TOPLEVEL_BINDING.receiver }
-
     context 'when forking is supported' do
       before do
         if ::Process.singleton_class.ancestors.include?(Datadog::Core::Utils::AtForkMonkeyPatch::ProcessMonkeyPatch)
           skip 'Monkey patch already applied (unclean state)'
         end
       end
 
+      let(:toplevel_receiver) { TOPLEVEL_BINDING.receiver }
+
       context 'on Ruby 3.0 or below' do
         before { skip 'Test applies only to Ruby 3.0 or below' if RUBY_VERSION >= '3.1' }
 
@@ -47,6 +47,20 @@
             expect(::Process.method(:_fork).source_location.first).to match(/.*at_fork_monkey_patch.rb/)
           end
         end
+
+        it 'does not monkey patch Kernel/Object' do
+          expect_in_fork do
+            apply!
+
+            expect(::Process.ancestors).to_not include(described_class::KernelMonkeyPatch)
+            expect(::Kernel.ancestors).to_not include(described_class::KernelMonkeyPatch)
+            expect(toplevel_receiver.class.ancestors).to_not include(described_class::KernelMonkeyPatch)
+
+            expect(::Process.method(:fork).source_location&.first).to_not match(/.*at_fork_monkey_patch.rb/)
+            expect(::Kernel.method(:fork).source_location&.first).to_not match(/.*at_fork_monkey_patch.rb/)
+            expect(toplevel_receiver.method(:fork).source_location&.first).to_not match(/.*at_fork_monkey_patch.rb/)
+          end
+        end
       end
     end
 
@@ -57,7 +71,7 @@
           .and_return(false)
       end
 
-      it 'skips the Kernel patch' do
+      it 'skips the monkey patch' do
         is_expected.to be false
       end
     end
@@ -91,15 +105,15 @@ def fork(&block)
       end
     end
 
-    shared_context 'datadog_at_fork callbacks' do
+    shared_context 'at_fork callbacks' do
       let(:child) { double('child') }
 
       before do
-        Datadog::Core::Utils::AtForkMonkeyPatch::ProcessMonkeyPatch.datadog_at_fork(:child) { child.call }
+        Datadog::Core::Utils::AtForkMonkeyPatch.at_fork(:child) { child.call }
       end
 
       after do
-        described_class.datadog_at_fork_blocks.clear
+        Datadog::Core::Utils::AtForkMonkeyPatch.const_get(:AT_FORK_CHILD_BLOCKS).clear
       end
     end
 
@@ -112,14 +126,12 @@ def fork(&block)
 
       describe '#fork' do
         context 'when a block is not provided' do
-          include_context 'datadog_at_fork callbacks'
+          include_context 'at_fork callbacks'
 
           subject(:fork) { fork_class.fork }
 
           context 'and returns from the parent context' do
-            # By setting the fork result = integer, we're
-            # simulating #fork running in the parent process.
-            let(:fork_result) { rand(100) }
+            let(:fork_result) { 1234 } # simulate parent: result is a pid
 
             it do
               expect(child).to_not receive(:call)
@@ -129,9 +141,7 @@ def fork(&block)
           end
 
           context 'and returns from the child context' do
-            # By setting the fork result = nil, we're
-            # simulating #fork running in the child process.
-            let(:fork_result) { nil }
+            let(:fork_result) { nil } # simulate child: result is a nil
 
             it do
               expect(child).to receive(:call)
@@ -154,7 +164,7 @@ def fork(&block)
           end
 
           context 'when callbacks are configured' do
-            include_context 'datadog_at_fork callbacks'
+            include_context 'at_fork callbacks'
 
             it 'invokes all the callbacks in order' do
               expect(child).to receive(:call)
@@ -164,48 +174,6 @@ def fork(&block)
           end
         end
       end
-
-      describe '#datadog_at_fork' do
-        include_context 'datadog_at_fork callbacks'
-
-        let(:callback) { double('callback') }
-        let(:block) { proc { callback.call } }
-
-        subject(:datadog_at_fork) do
-          Datadog::Core::Utils::AtForkMonkeyPatch::ProcessMonkeyPatch.datadog_at_fork(:child, &block)
-        end
-
-        it 'adds a child callback' do
-          datadog_at_fork
-
-          expect(child).to receive(:call).ordered
-          expect(callback).to receive(:call).ordered
-
-          fork_class.fork {}
-        end
-      end
-    end
-
-    context 'when applied to multiple classes with forking' do
-      include_context 'fork class'
-
-      let(:other_fork_class) { new_fork_class }
-
-      context 'and #datadog_at_fork is called in one' do
-        include_context 'datadog_at_fork callbacks'
-
-        it 'applies the callback to the original class' do
-          expect(child).to receive(:call)
-
-          fork_class.fork {}
-        end
-
-        it 'applies the callback to the other class' do
-          expect(child).to receive(:call)
-
-          other_fork_class.fork {}
-        end
-      end
     end
   end
 
@@ -215,49 +183,42 @@ def fork(&block)
       result = _fork_result
 
       Module.new do
-        def self.daemon(nochdir = nil, noclose = nil); end
+        def self.daemon(nochdir = nil, noclose = nil)
+          [nochdir, noclose]
+        end
         define_singleton_method(:_fork) { result }
       end
     end
     let(:child_callback) { double('child', call: true) }
 
     before do
-      allow(process_module).to receive(:daemon)
-
-      process_module.singleton_class.prepend(Datadog::Core::Utils::AtForkMonkeyPatch::KernelMonkeyPatch)
       process_module.singleton_class.prepend(described_class)
 
-      process_module.datadog_at_fork(:child) { child_callback.call }
+      Datadog::Core::Utils::AtForkMonkeyPatch.at_fork(:child) { child_callback.call }
     end
 
     after do
-      Datadog::Core::Utils::AtForkMonkeyPatch::KernelMonkeyPatch.datadog_at_fork_blocks.clear
-    end
-
-    it 'calls the child datadog_at_fork callbacks after calling Process.daemon' do
-      expect(process_module).to receive(:daemon).ordered
-      expect(child_callback).to receive(:call).ordered
-
-      process_module.daemon
+      Datadog::Core::Utils::AtForkMonkeyPatch.const_get(:AT_FORK_CHILD_BLOCKS).clear
     end
 
-    it 'passes any arguments to Process.daemon' do
-      expect(process_module).to receive(:daemon).with(true, true)
+    describe '.daemon' do
+      it 'calls the child at_fork callbacks after calling Process.daemon' do
+        expect(process_module).to receive(:daemon).ordered.and_call_original
+        expect(child_callback).to receive(:call).ordered
 
-      process_module.daemon(true, true)
-    end
-
-    it 'returns the result of calling Process.daemon' do
-      expect(process_module).to receive(:daemon).and_return(:process_daemon_result)
+        process_module.daemon
+      end
 
-      expect(process_module.daemon).to be :process_daemon_result
+      it 'passes any arguments to Process.daemon and returns its results' do
+        expect(process_module.daemon(:arg1, :arg2)).to eq([:arg1, :arg2])
+      end
     end
 
-    describe 'Process._fork monkey patch' do
+    describe '_fork' do
       context 'in the child process' do
         let(:_fork_result) { 0 }
 
-        it 'monkey patches _fork to call the child datadog_at_fork callbacks on the child process' do
+        it 'triggers the child callbacks' do
           expect(child_callback).to receive(:call)
 
           expect(process_module._fork).to be 0
@@ -271,7 +232,7 @@ def self.daemon(nochdir = nil, noclose = nil); end
       context 'in the parent process' do
         let(:_fork_result) { 1234 }
 
-        it 'does not trigger any callbacks' do
+        it 'does not trigger the child callbacks' do
           expect(child_callback).to_not receive(:call)
 
           process_module._fork