Skip to content

Commit

Permalink
Merge pull request #3834 from DataDog/ivoanjo/prof-10241-simplify-at-…
Browse files Browse the repository at this point in the history
…fork-monkey-patch

[PROF-10241] Fix issue in tests by simplifying `at_fork` monkey patch
  • Loading branch information
ivoanjo authored Aug 12, 2024
2 parents aa6550c + 7642788 commit 46c8c8b
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
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
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
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 46c8c8b

Please sign in to comment.