From dcb0947114122ab1f7eab134bb7416daf2d038a7 Mon Sep 17 00:00:00 2001
From: iamluc <luc.vieillescazes@datadoghq.com>
Date: Tue, 10 Sep 2024 09:21:53 +0200
Subject: [PATCH] Work around race conditions leading to different orders or
 numbers of requests

---
 tests/Common/TracerTestTrait.php              | 36 +++++++++++++++++--
 tests/Integrations/PCNTL/PCNTLTest.php        | 35 +++++++++---------
 .../PCNTL/scripts/long-running-autoflush.php  | 11 ++----
 3 files changed, 53 insertions(+), 29 deletions(-)

diff --git a/tests/Common/TracerTestTrait.php b/tests/Common/TracerTestTrait.php
index e65cbb7d61e..ce73cb08d88 100644
--- a/tests/Common/TracerTestTrait.php
+++ b/tests/Common/TracerTestTrait.php
@@ -455,7 +455,7 @@ public function retrieveDumpedMetrics(callable $until = null, $throw = false)
     }
 
     public function retrieveAnyDumpedData(callable $until = null, $throw, $metrics = false) {
-        $until = $until ?? $this->untilFirstTelemetryRequest();
+        $until = $until ?? $this->untilFirstTraceRequest();
 
         $allResponses = [];
 
@@ -512,9 +512,11 @@ function untilNumberOfTraces($number) {
         };
     }
 
-    function untilFirstTelemetryRequest() {
+    function untilFirstTraceRequest() {
         return function ($request) {
-            return (strpos($request["uri"] ?? "", "/telemetry/") !== 0);
+            return (strpos($request["uri"] ?? "", "/v0.4/traces") === 0)
+                || (strpos($request["uri"] ?? "", "/v0.7/traces") === 0)
+            ;
         };
     }
 
@@ -526,6 +528,34 @@ function untilTelemetryRequest($metricName) {
         };
     }
 
+    function untilSpan(SpanAssertion $assertion) {
+        return function ($request) use ($assertion) {
+            if (strpos($request["uri"] ?? "", "/telemetry/") === 0 || !isset($request['body'])) {
+                return false;
+            }
+            $trace = $this->parseRawDumpedTraces(json_decode($request['body'], true));
+            try {
+                (new SpanChecker())->assertFlameGraph($trace, [$assertion]);
+            } catch (\Exception $e) {
+                return false;
+            }
+
+            return true;
+        };
+    }
+
+    function until(...$expectations) {
+        return function ($request) use (&$expectations) {
+            foreach ($expectations as $key => $expect) {
+                if ($expect($request)) {
+                    unset($expectations[$key]);
+                }
+            }
+
+            return !count($expectations);
+        };
+    }
+
     /**
      * Set a property into an object from an array optionally applying a converter.
      *
diff --git a/tests/Integrations/PCNTL/PCNTLTest.php b/tests/Integrations/PCNTL/PCNTLTest.php
index 14d84aef12a..6aa0d577756 100644
--- a/tests/Integrations/PCNTL/PCNTLTest.php
+++ b/tests/Integrations/PCNTL/PCNTLTest.php
@@ -250,7 +250,7 @@ public function testCliShortRunningMultipleNestedForks()
 
     public function testCliLongRunningMultipleForksAutoFlush()
     {
-        list($requests) = $this->inCli(
+        $this->inCli(
             __DIR__ . '/scripts/long-running-autoflush.php',
             [
                 'DD_TRACE_CLI_ENABLED' => 'true',
@@ -262,24 +262,23 @@ public function testCliLongRunningMultipleForksAutoFlush()
             [],
             '',
             false,
-            $this->untilNumberOfTraces(4)
-        );
-        $this->assertCount(4, $requests);
-
-        for ($i = 0; $i < 4; $i += 2) {
-            $this->assertFlameGraph([$requests[$i]], [
-                SpanAssertion::exists('curl_exec', '/httpbin_integration/child-'.($i/2)),
-            ]);
-
-            $this->assertFlameGraph([$requests[$i + 1]], [
-                SpanAssertion::exists('long_running_entry_point')->withChildren([
-                    SpanAssertion::exists('curl_exec', '/httpbin_integration/entry_point'),
-                    SpanAssertion::exists('curl_exec', '/httpbin_integration/main_process'),
-                    SpanAssertion::exists('curl_exec', '/httpbin_integration/end_entry_point'),
+            $this->until(
+                $this->untilSpan(SpanAssertion::exists('curl_exec', '/httpbin_integration/child-0')),
+                $this->untilSpan(SpanAssertion::exists('curl_exec', '/httpbin_integration/child-1')),
+                $this->untilSpan(SpanAssertion::exists('long_running_entry_point')->withChildren([
+                    SpanAssertion::exists('curl_exec', '/httpbin_integration/entry_point-0'),
+                    SpanAssertion::exists('curl_exec', '/httpbin_integration/main_process-0'),
+                    SpanAssertion::exists('curl_exec', '/httpbin_integration/end_entry_point-0'),
                     SpanAssertion::exists('pcntl_fork'),
-                ]),
-            ]);
-        }
+                ])),
+                $this->untilSpan(SpanAssertion::exists('long_running_entry_point')->withChildren([
+                    SpanAssertion::exists('curl_exec', '/httpbin_integration/entry_point-1'),
+                    SpanAssertion::exists('curl_exec', '/httpbin_integration/main_process-1'),
+                    SpanAssertion::exists('curl_exec', '/httpbin_integration/end_entry_point-1'),
+                    SpanAssertion::exists('pcntl_fork'),
+                ])),
+            )
+        );
     }
 
     public function testCliLongRunningMultipleForksManualFlush()
diff --git a/tests/Integrations/PCNTL/scripts/long-running-autoflush.php b/tests/Integrations/PCNTL/scripts/long-running-autoflush.php
index d9f713b25ac..876ff59dd96 100644
--- a/tests/Integrations/PCNTL/scripts/long-running-autoflush.php
+++ b/tests/Integrations/PCNTL/scripts/long-running-autoflush.php
@@ -11,30 +11,25 @@
 
 for ($iteration = 0; $iteration < ITERATIONS; $iteration++) {
     long_running_entry_point($iteration);
-
-    // Add a delay to ensure the spans from each iteration are sent separately
-    // if execution is too fast, they can be grouped in 1 "request", but we expect 2.
-    usleep(1000);
 }
 
 function long_running_entry_point($iteration)
 {
-    call_httpbin('entry_point');
+    call_httpbin('entry_point-'.$iteration);
 
     $forkPid = pcntl_fork();
 
     if ($forkPid > 0) {
         // Main
-        call_httpbin('main_process');
+        call_httpbin('main_process-'.$iteration);
     } else if ($forkPid === 0) {
         // Child
-        usleep(1000);
         call_httpbin('child-'.$iteration);
         exit(0);
     } else {
         error_log('Error');
         exit(-1);
     }
-    call_httpbin('end_entry_point');
+    call_httpbin('end_entry_point-'.$iteration);
     pcntl_waitpid($forkPid, $childStatus);
 }