Skip to content

Commit

Permalink
[PROF-10241] Fix issue in tests by simplifying at_fork monkey patch
Browse files Browse the repository at this point in the history 
**What does this PR do?**

This PR fixes an issue that we identified in the
`profiling/tasks/setup_spec.rb` that was triggered by applying the
at_fork monkey patch before running that spec.

This issue could be triggered by adding

```ruby
Datadog::Core::Utils::AtForkMonkeyPatch.apply!
```

to the top of the `setup_spec.rb` (before the `RSpec.describe`).

```
Datadog::Profiling::Tasks::Setup
  #run
    actives the forking extension before setting up the at_fork hooks
    only sets up the extensions and hooks once, even across different instances
  #setup_at_fork_hooks
    when Process#datadog_at_fork is available
      sets up an at_fork hook that restarts the profiler (FAILED - 1)
      when there is an issue starting the profiler
        logs an exception (FAILED - 2)
        does not raise any error (FAILED - 3)
    when #datadog_at_fork is not available
      logs a debug message

Failures:

  1) Datadog::Profiling::Tasks::Setup#setup_at_fork_hooks when Process#datadog_at_fork is available sets up an at_fork hook that restarts the profiler
     Failure/Error:
       example.run.tap do
         tracer_shutdown!
       end

     NameError:
       undefined method `datadog_at_fork' for class `#<Class:Process>'

               object_singleton_class.__send__(@original_visibility, method_name)
                                     ^^^^^^^^^
     # ./spec/spec_helper.rb:230:in `block (2 levels) in <top (required)>'
     # ./spec/spec_helper.rb:115:in `block (2 levels) in <top (required)>'
     # webmock-3.23.1/lib/webmock/rspec.rb:39:in `block (2 levels) in <top (required)>'
     # rspec-wait-0.0.10/lib/rspec/wait.rb:47:in `block (2 levels) in <top (required)>'

  2) Datadog::Profiling::Tasks::Setup#setup_at_fork_hooks when Process#datadog_at_fork is available when there is an issue starting the profiler logs an exception
     Failure/Error: at_fork_hook.call

     NoMethodError:
       undefined method `call' for nil:NilClass

                 at_fork_hook.call
                             ^^^^^
     # ./spec/datadog/profiling/tasks/setup_spec.rb:84:in `block (5 levels) in <top (required)>'
     # ./spec/spec_helper.rb:230:in `block (2 levels) in <top (required)>'
     # ./spec/spec_helper.rb:115:in `block (2 levels) in <top (required)>'
     # webmock-3.23.1/lib/webmock/rspec.rb:39:in `block (2 levels) in <top (required)>'
     # rspec-wait-0.0.10/lib/rspec/wait.rb:47:in `block (2 levels) in <top (required)>'

  3) Datadog::Profiling::Tasks::Setup#setup_at_fork_hooks when Process#datadog_at_fork is available when there is an issue starting the profiler does not raise any error
     Failure/Error: at_fork_hook.call

     NoMethodError:
       undefined method `call' for nil:NilClass

                 at_fork_hook.call
                             ^^^^^
     # ./spec/datadog/profiling/tasks/setup_spec.rb:76:in `block (5 levels) in <top (required)>'
     # ./spec/spec_helper.rb:230:in `block (2 levels) in <top (required)>'
     # ./spec/spec_helper.rb:115:in `block (2 levels) in <top (required)>'
     # webmock-3.23.1/lib/webmock/rspec.rb:39:in `block (2 levels) in <top (required)>'
     # rspec-wait-0.0.10/lib/rspec/wait.rb:47:in `block (2 levels) in <top (required)>'
```

This issue was caused by an incompatibility between rspec-mocks (not
entirely sure if it's a bug) and the monkey patching we were doing.

Rather than fighting with more monkey patch weirdness or rspec,
I decided to fix the issue by simplifying the at fork monkey patch
module by making it as minimal as possible:

* It no longer places extra unneeded functions in Ruby classes, e.g.
  `datadog_at_fork_blocks` and `datadog_at_fork`
* It no longer stores its data in a weird class var

This enabled a bunch of simplifications, and the affected
spec was fixed as a side-effect of this.

**Motivation:**

Get the at_fork monkey patch in shape to be used by the crashtracker as
well.

**Additional Notes:**

I was tempted to do a larger refactoring of this module when I was
extracting it, but I decided to aim for a smaller diff.

Since I needed to dive in again to fix the issue with the spec, I
decided to finally go for it.

**How to test the change?**

These changes include test coverage.

Here's a tiny app that can be used to experiment with this:

```ruby
puts RUBY_DESCRIPTION

require 'datadog/core/utils/at_fork_monkey_patch'

Datadog::Core::Utils::AtForkMonkeyPatch.apply!

Datadog::Core::Utils::AtForkMonkeyPatch.at_fork(:child) { puts "Hello from child process: #{Process.pid} "}

puts "Parent pid: #{Process.pid}"

puts "fork { }"
fork { }
Process.wait

puts "Kernel.fork { }"
Kernel.fork { }
Process.wait

puts "Process.fork { }"
Process.fork { }
Process.wait

puts "Foo.new.call"
class Foo
  def call
    fork { }
  end

  def self.do_fork
    fork {  }
  end
end
Foo.new.call
Process.wait

puts "Foo.do_fork"
Foo.do_fork
Process.wait

puts "Fork from outside"
Foo.new.send(:fork) { }
Process.wait

class BasicFoo < BasicObject
  include ::Kernel

  def call
    fork { }
  end

  def self.do_fork
    fork { }
  end
end

puts "BasicFoo.new.call"
BasicFoo.new.call
Process.wait

puts "BasicObject:"
Class.new(BasicObject) { include(::Kernel); def call; fork { }; end }.new.call
Process.wait

puts "BasicFork.do_fork"
BasicFoo.do_fork
Process.wait
```
  • Loading branch information
@ivoanjo
ivoanjo committed Aug 8, 2024
1 parent 850462b commit 3dc71fb
Show file tree
Hide file tree
Showing 4 changed files with 97 additions and 177 deletions.
62 changes: 25 additions & 37 deletions lib/datadog/core/utils/at_fork_monkey_patch.rb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
20 changes: 8 additions & 12 deletions lib/datadog/profiling/tasks/setup.rb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
121 changes: 41 additions & 80 deletions spec/datadog/core/utils/at_fork_monkey_patch_spec.rb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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' }

Expand Down Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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

Expand All @@ -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
Expand All @@ -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
Expand Down
Loading

0 comments on commit 3dc71fb

Please sign in to comment.