Fix race condition in Remote::Component::Barrier causing flaky boot tags (#5545)

Author: p-datadog

lib/datadog/core/remote/component.rb

@@ -105,44 +105,61 @@ def after_fork
         class Barrier
           def initialize(timeout = nil)
             @once = false
+            @waited = false
             @timeout = timeout
 
             @mutex = Mutex.new
             @condition = ConditionVariable.new
           end
 
           # Wait for first lift to happen, otherwise don't wait
+          #
+          # Returns:
+          # - :lift if the barrier was lifted (worker completed a cycle)
+          # - :timeout if the wait timed out before the barrier was lifted
+          # - :pass if wait_once was already called previously
+          #
+          # Uses a separate @waited flag to distinguish "already waited" (:pass)
+          # from "worker lifted before we could wait" (:lift). Without this,
+          # a race between Worker#start and wait_once can cause the first call
+          # to return :pass if the worker completes before wait_once runs.
           def wait_once(timeout = nil)
-            # TTAS (Test and Test-And-Set) optimisation
-            # Since @once only ever goes from false to true, this is semantically valid
-            return :pass if @once
-
-            begin
-              @mutex.lock
+            # TTAS (Test and Test-And-Set) optimisation for subsequent calls.
+            # @waited is only set inside the mutex and only transitions false -> true,
+            # so an unsynchronized read is safe: a stale `false` just falls through
+            # to the synchronized path which re-checks.
+            return :pass if @waited
 
-              return :pass if @once
-
-              timeout ||= @timeout
+            @mutex.synchronize do
+              return :pass if @waited
 
-              # - starting with Ruby 3.2, ConditionVariable#wait returns nil on
-              #   timeout and an integer otherwise
-              # - before Ruby 3.2, ConditionVariable returns itself
-              # so we have to rely on @once having been set
-              if RUBY_VERSION >= '3.2'
-                lifted = @condition.wait(@mutex, timeout)
+              if @once
+                # Worker lifted the barrier before we could wait.
+                # This is still the first call, so return :lift not :pass.
+                lifted = true
               else
-                @condition.wait(@mutex, timeout)
-                lifted = @once
+                timeout ||= @timeout
+
+                # - starting with Ruby 3.2, ConditionVariable#wait returns nil on
+                #   timeout and an integer otherwise
+                # - before Ruby 3.2, ConditionVariable returns itself
+                # so we have to rely on @once having been set
+                if RUBY_VERSION >= '3.2'
+                  lifted = @condition.wait(@mutex, timeout)
+                else
+                  @condition.wait(@mutex, timeout)
+                  lifted = @once
+                end
               end
 
+              @waited = true
+
               if lifted
                 :lift
               else
                 @once = true
                 :timeout
               end
-            ensure
-              @mutex.unlock
             end
           end
 

sig/datadog/core/remote/component.rbs

@@ -41,6 +41,7 @@ module Datadog
           type timeout_s = ::Integer | ::Float
 
           @once: bool
+          @waited: bool
           @timeout: timeout_s?
 
           @mutex: ::Thread::Mutex

spec/datadog/core/remote/component_spec.rb

@@ -430,5 +430,15 @@
         expect(record).to eq [:one, :two, :three]
       end
     end
+
+    context 'when lift happens before wait_once' do
+      it 'returns :lift on first call, not :pass' do
+        # Simulate the race: worker lifts barrier before wait_once is called
+        barrier.lift
+
+        expect(barrier.wait_once).to eq :lift
+        expect(barrier.wait_once).to eq :pass
+      end
+    end
   end
 end