From e0727ddc2aee8b4559e0876e793d8bfb1dbadbd7 Mon Sep 17 00:00:00 2001
From: cyve <madislak@yahoo.fr>
Date: Fri, 29 Dec 2023 10:53:53 +0100
Subject: [PATCH] fix: end tracing when receiving last chunk

---
 src/Http/TracedResponse.php    | 30 ++++++++++++++++++------------
 src/Http/TracingHttpClient.php |  7 ++++++-
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/Http/TracedResponse.php b/src/Http/TracedResponse.php
index 835c163..99c23bb 100644
--- a/src/Http/TracedResponse.php
+++ b/src/Http/TracedResponse.php
@@ -15,6 +15,7 @@
 use Symfony\Component\HttpClient\Exception\ServerException;
 use Symfony\Component\HttpClient\Response\StreamableInterface;
 use Symfony\Component\HttpClient\Response\StreamWrapper;
+use Symfony\Contracts\HttpClient\Exception\TransportExceptionInterface;
 use Symfony\Contracts\HttpClient\HttpClientInterface;
 use Symfony\Contracts\HttpClient\ResponseInterface;
 
@@ -93,20 +94,16 @@ public function getInfo(string $type = null): mixed
 
     public function toStream(bool $throw = true)
     {
-        try {
-            if ($throw) {
-                // Ensure headers arrived
-                $this->response->getHeaders();
-            }
-
-            if ($this->response instanceof StreamableInterface) {
-                return $this->stream = $this->response->toStream(false);
-            }
+        if ($throw) {
+            // Ensure headers arrived
+            $this->response->getHeaders();
+        }
 
-            return $this->stream = StreamWrapper::createResource($this->response);
-        } finally {
-            $this->endTracing();
+        if ($this->response instanceof StreamableInterface) {
+            return $this->stream = $this->response->toStream(false);
         }
+
+        return $this->stream = StreamWrapper::createResource($this->response);
     }
 
     /**
@@ -129,6 +126,14 @@ public static function stream(HttpClientInterface $client, iterable $responses,
         }
 
         foreach ($client->stream($wrappedResponses, $timeout) as $r => $chunk) {
+            try {
+                if ($chunk->isLast() || $chunk->isTimeout()) {
+                    $traceableMap[$r]->endTracing();
+                }
+            } catch (TransportExceptionInterface) {
+                $traceableMap[$r]->endTracing();
+            }
+
             yield $traceableMap[$r] => $chunk;
         }
     }
@@ -150,6 +155,7 @@ protected function endTracing(): void
             if (\in_array('response.body', $info['user_data']['span_attributes'] ?? [])) {
                 if (empty($this->content) && \is_resource($this->stream)) {
                     $this->content = stream_get_contents($this->stream) ?: null;
+                    rewind($this->stream);
                 }
                 $this->span->setAttribute('response.body', $this->content);
             }
diff --git a/src/Http/TracingHttpClient.php b/src/Http/TracingHttpClient.php
index 13919ee..cc8e35c 100644
--- a/src/Http/TracingHttpClient.php
+++ b/src/Http/TracingHttpClient.php
@@ -16,6 +16,7 @@
 use Instrumentation\Tracing\Tracing;
 use OpenTelemetry\API\Trace\SpanKind;
 use OpenTelemetry\API\Trace\StatusCode;
+use OpenTelemetry\SDK\Common\Time\ClockInterface;
 use OpenTelemetry\SemConv\TraceAttributes;
 use Symfony\Component\HttpClient\DecoratorTrait;
 use Symfony\Component\HttpClient\HttpClient;
@@ -119,7 +120,11 @@ public function request(string $method, string $url, array $options = []): Respo
 
                     $span->setAttribute(TraceAttributes::HTTP_STATUS_CODE, $info['http_code']);
                     $span->setAttribute(TraceAttributes::HTTP_URL, $info['url']);
-                    $span->addEvent('HTTP headers received');
+
+                    if (array_key_exists('total_time', $info)) {
+                        $timestamp = (int) (($info['start_time'] + $info['total_time']) * ClockInterface::NANOS_PER_SECOND);
+                    }
+                    $span->addEvent('http.response.headers', [], $timestamp ?? null);
 
                     if ($info['http_code'] >= 400) {
                         $span->setStatus(StatusCode::STATUS_ERROR);