Fix process spawn

docs/DevelopmentGuide.md

@@ -185,7 +185,7 @@ end
 3. Now let's generate that dependency Gemfile with `rake`. Simply run
 
 > [!IMPORTANT]
-> Ensure you are either using Ruby 3.3 as the current Ruby version (`ruby -v`) or running commands within a Docker container.
+> Ensure you are either using the repo's current Ruby version from `.ruby-version` (`ruby -v`) or running commands within a Docker container.
 
 ```console
 $ bundle exec rake dependency:generate

lib/datadog/core/configuration/components.rb

@@ -133,7 +133,7 @@ def initialize(settings)
           self.class::PATCH_ONLY_ONCE.run do
             Utils::AtForkMonkeyPatch.apply!
             Utils::SpawnMonkeyPatch.apply!(
-              lineage_envs_provider: Core::Environment::Identity.method(:runtime_propagation_envs),
+              env_provider: Core::Environment::Identity.method(:runtime_propagation_envs),
             )
 
             # Register callback that calls Components.after_fork

lib/datadog/core/utils/spawn_monkey_patch.rb

@@ -3,32 +3,46 @@
 module Datadog
   module Core
     module Utils
+      # Applies the Process.spawn wrapper used to merge additional environment variables
+      # into child processes.
       module SpawnMonkeyPatch
-        # @param lineage_envs_provider [#call] returns a Hash of env vars to merge into the child process
-        def self.apply!(lineage_envs_provider:)
-          @lineage_envs_provider = lineage_envs_provider
+        # @param env_provider [#call] returns a Hash of env vars to merge into the child process
+        def self.apply!(env_provider:)
+          @env_provider = env_provider
+
+          return if ::Process.singleton_class.ancestors.include?(ProcessSpawnPatch)
+
           ::Process.singleton_class.prepend(ProcessSpawnPatch)
-          true
         end
 
+        # Vessel for env_provider propagation.
         module ProcessSpawnPatch
-          def spawn(*args, **opts)
-            args.replace(SpawnMonkeyPatch.inject_lineage_envs(args))
-            super
+          def spawn(*args)
+            super(*SpawnMonkeyPatch.inject_envs(args))
           end
         end
 
-        # Process.spawn(env?, cmd, ...): env is optional first arg (Hash). When present, merge
-        # runtime_ids into it; when absent, prepend full ENV + runtime_ids so the child inherits both.
-        def self.inject_lineage_envs(args)
-          runtime_ids = @lineage_envs_provider.call
-          env_provided = Hash === args.first
+        # Merge the env vars from `env_provider` with the `env` argument from {Process.spawn}.
+        #
+        # `env` is the first argument to {Process.spawn}, which is an optional {Hash}
+        # (https://docs.ruby-lang.org/en/4.0/Process.html#method-c-spawn):
+        # `Process.spawn([env, ] *args, options = {})`
+        #
+        # One thing to note is that parent process' (this process') environment variables are
+        # inherited by default by the spawned process. They are merged with (and possibly overwritten by)
+        # the env vars from the argument `env`.
+        # (https://docs.ruby-lang.org/en/4.0/Process.html#module-Process-label-Environment+Variables+-28-3Aunsetenv_others-29)
+        # This doesn't affect the current implementation.
+        def self.inject_envs(args)
+          provided_env = @env_provider.call
 
-          base_env = env_provided ? args.first : DATADOG_ENV.to_h
-          env = base_env.merge(runtime_ids)
-          rest = env_provided ? args.drop(1) : args
+          if ::Hash === args.first
+            args[0] = args.first.merge(provided_env)
+          else
+            args.unshift(provided_env)
+          end
 
-          [env, *rest]
+          args
         end
       end
     end

sig/datadog/core/utils/spawn_monkey_patch.rbs

@@ -3,15 +3,15 @@ module Datadog
     module Utils
       module SpawnMonkeyPatch
         # Set in apply! before Process.spawn is intercepted; internal wiring only.
-        self.@lineage_envs_provider: ^() -> ::Hash[::String, ::String]
+        self.@env_provider: ^() -> ::Hash[::String, ::String]
 
-        def self.apply!: (lineage_envs_provider: ^() -> ::Hash[::String, ::String]) -> true
+        def self.apply!: (env_provider: ^() -> ::Hash[::String, ::String]) -> untyped
 
         module ProcessSpawnPatch
-          def spawn: (*untyped args, **untyped opts) -> untyped
+          def spawn: (*untyped args) -> untyped
         end
 
-        def self.inject_lineage_envs: (untyped args) -> ::Array[untyped]
+        def self.inject_envs: (untyped args) -> ::Array[untyped]
       end
     end
   end

spec/datadog/core/utils/spawn_monkey_patch_spec.rb

@@ -5,8 +5,28 @@
 require 'datadog/core/configuration/settings'
 
 RSpec.describe Datadog::Core::Utils::SpawnMonkeyPatch do
+  let(:envs) do
+    {
+      'ENV1' => 'val1',
+      'ENV2' => 'val2',
+    }
+  end
+
+  def process_spawn(*spawn_args)
+    IO.pipe do |read_io, write_io|
+      process_options = {in: File::NULL, out: write_io, err: write_io}
+      process_options.merge!(spawn_args.pop) if ::Hash === spawn_args.last
+
+      pid = Process.spawn(*spawn_args, process_options)
+      write_io.close
+      Process.wait(pid)
+
+      read_io.read.lines.map { |line| line.chomp.split('=', 2) }.to_h
+    end
+  end
+
   describe '::apply!' do
-    subject(:apply!) { described_class.apply!(lineage_envs_provider: -> { {} }) }
+    subject(:apply!) { described_class.apply!(env_provider: -> { envs }) }
 
     context 'when Process.spawn is supported' do
       before do
@@ -15,12 +35,86 @@
       end
 
       it 'prepends the spawn monkey patch' do
-        expect_in_fork do
-          apply!
-          expect(Process.singleton_class.ancestors).to include(described_class::ProcessSpawnPatch)
-          expect(Process.method(:spawn).source_location.first).to match(/spawn_monkey_patch\.rb/)
-        end
+        apply!
+        expect(Process.singleton_class.ancestors).to include(described_class::ProcessSpawnPatch)
+        expect(Process.method(:spawn).source_location.first).to match(/spawn_monkey_patch\.rb/)
+      end
+
+      it 'does not patch twice' do
+        described_class.apply!(env_provider: -> { {'ENV1' => 'val1'} })
+        described_class.apply!(env_provider: -> { {'ENV1' => 'val2'} })
+
+        expect(Process.singleton_class.ancestors.count(described_class::ProcessSpawnPatch)).to eq(1)
+        expect(process_spawn('/usr/bin/env')).to include('ENV1' => 'val2')
+      end
+    end
+  end
+
+  describe 'on Process.spawn' do
+    subject(:apply!) { described_class.apply!(env_provider: -> { envs }) }
+
+    around do |example|
+      ClimateControl.modify('PARENT1' => 'parent_val') { example.run }
+    end
+
+    before do
+      skip 'Process.spawn not supported' unless Process.respond_to?(:spawn)
+      apply!
+    end
+
+    it 'merges env_provider, parent envs, and env argument' do
+      output = process_spawn({'ARG' => 'arg_val'}, '/usr/bin/env', pgroup: true)
+
+      expect(output).to include('PARENT1' => 'parent_val', 'ENV1' => 'val1', 'ENV2' => 'val2', 'ARG' => 'arg_val')
+    end
+
+    it 'merges env_provider and parent envs when no env argument is provided' do
+      output = process_spawn('/usr/bin/env', pgroup: true)
+
+      expect(output).to include('PARENT1' => 'parent_val', 'ENV1' => 'val1', 'ENV2' => 'val2')
+    end
+
+    it 'respects parent env removal through the value `nil`' do
+      output = process_spawn({'PARENT1' => nil}, '/usr/bin/env')
+
+      expect(output).not_to include('PARENT1')
+      expect(output).to include('ENV1' => 'val1', 'ENV2' => 'val2')
+    end
+
+    it 'respects array-form command variant' do
+      command = 'printf %s "$0:$ARG:$PARENT1:$ENV1:$ENV2"'
+
+      output = IO.pipe do |read_io, write_io|
+        pid = Process.spawn(
+          {'ARG' => 'arg_val'},
+          ['/bin/sh', 'cmd-name'],
+          '-c',
+          command,
+          in: File::NULL,
+          out: write_io,
+          err: write_io,
+        )
+        write_io.close
+        Process.wait(pid)
+
+        read_io.read
       end
+
+      expect(output).to eq('cmd-name:arg_val:parent_val:val1:val2')
+    end
+  end
+
+  describe '::inject_envs' do
+    subject(:inject_envs) { described_class.inject_envs(args.dup) }
+    let(:args) { [env_argument, '/bin/ls', '.', {pgroup: 0}] }
+    let(:env_argument) { {'TZ' => 'UTC'} }
+
+    before do
+      described_class.apply!(env_provider: -> { envs })
+    end
+
+    it 'does not mutate the provided env argument Hash' do
+      expect { inject_envs }.not_to change { env_argument }
     end
   end
 
@@ -33,16 +127,14 @@
     end
 
     it 'applies both fork and spawn patches when Components is initialized' do
-      expect_in_fork do
-        Datadog::Core::Configuration::Components.new(Datadog::Core::Configuration::Settings.new)
+      Datadog::Core::Configuration::Components.new(Datadog::Core::Configuration::Settings.new)
 
-        expect(Process.singleton_class.ancestors).to include(
-          Datadog::Core::Utils::AtForkMonkeyPatch::ProcessMonkeyPatch,
-        )
-        expect(Process.singleton_class.ancestors).to include(
-          Datadog::Core::Utils::SpawnMonkeyPatch::ProcessSpawnPatch,
-        )
-      end
+      expect(Process.singleton_class.ancestors).to include(
+        Datadog::Core::Utils::AtForkMonkeyPatch::ProcessMonkeyPatch,
+      )
+      expect(Process.singleton_class.ancestors).to include(
+        Datadog::Core::Utils::SpawnMonkeyPatch::ProcessSpawnPatch,
+      )
     end
   end
 end

spec/support/synchronization_helpers.rb

@@ -1,6 +1,11 @@
 require 'English'
 
 module SynchronizationHelpers
+  # Runs the given block in a fork, allowing you to perform RSpec assertions in a fork
+  # and have them be reported in the parent process.
+  #
+  # You can alternatively use `execute_in_fork: true` {ForkableExample}
+  # if your whole example or example group should run in a forked process.
   def expect_in_fork(fork_expectations: nil, timeout_seconds: 10, trigger_stacktrace_on_kill: false, debug: false)
     fork_expectations ||= proc { |status:, stdout:, stderr:|
       expect(status && status.success?).to be(true), "STDOUT:`#{stdout}` STDERR:`#{stderr}"