/
forking.rb
98 lines (81 loc) · 3.82 KB
/
forking.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
# frozen_string_literal: true
module Datadog
module Profiling
module Ext
# Monkey patches `Kernel#fork`, adding a `Kernel#at_fork` callback mechanism which is used to restore
# profiling abilities after the VM forks.
#
# Known limitations: Does not handle `BasicObject`s that include `Kernel` directly; e.g.
# `Class.new(BasicObject) { include(::Kernel); def call; fork { }; end }.new.call`.
#
# This will be fixed once we moved to hooking into `Process._fork`
module Forking
def self.supported?
Process.respond_to?(:fork)
end
def self.apply!
return false unless supported?
[
::Process.singleton_class, # Process.fork
::Kernel.singleton_class, # Kernel.fork
::Object, # fork without explicit receiver (it's defined as a method in ::Kernel)
# Note: Modifying Object as we do here is irreversible. During tests, this
# change will stick around even if we otherwise stub `Process` and `Kernel`
].each { |target| target.prepend(Kernel) }
::Process.singleton_class.prepend(ProcessDaemonPatch)
end
# Extensions for kernel
#
# TODO: Consider hooking into `Process._fork` on Ruby 3.1+ instead, see
# https://github.com/ruby/ruby/pull/5017 and https://bugs.ruby-lang.org/issues/17795
module Kernel
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)
# Invoke original block
yield
end
end
# Start fork
# If a block is provided, use the wrapped version.
result = child_block.nil? ? super : super(&child_block)
# Trigger correct callbacks depending on whether we're in the parent or child.
# If we're in the fork, result = nil: trigger child callbacks.
# If we're in the parent, result = fork PID: trigger parent callbacks.
datadog_at_fork_blocks[:child].each(&:call) if result.nil? && datadog_at_fork_blocks.key?(:child)
# Return PID from #fork
result
end
def at_fork(stage, &block)
raise ArgumentError, 'Bad \'stage\' for ::at_fork' unless stage == :child
datadog_at_fork_blocks[stage] = [] unless datadog_at_fork_blocks.key?(stage)
datadog_at_fork_blocks[stage] << block
end
module_function
def datadog_at_fork_blocks
# Blocks should be shared across all users of this module,
# e.g. Process#fork, Kernel#fork, etc. should all invoke the same callbacks.
# rubocop:disable Style/ClassVars
@@datadog_at_fork_blocks ||= {}
# rubocop:enable Style/ClassVars
end
end
# A call to Process.daemon ( https://rubyapi.org/3.1/o/process#method-c-daemon ) forks the current process and
# keeps executing code in the child process, killing off the parent, thus effectively replacing it.
#
# This monkey patch makes the `Kernel#at_fork` mechanism defined above also work in this situation.
module ProcessDaemonPatch
def daemon(*args)
datadog_at_fork_blocks = Datadog::Profiling::Ext::Forking::Kernel.datadog_at_fork_blocks
result = super
datadog_at_fork_blocks[:child].each(&:call) if datadog_at_fork_blocks.key?(:child)
result
end
end
end
end
end
end