fix: end tracing when receiving last chunk

worldia · Dec 29, 2023 · 57e6ba0 · 57e6ba0
1 parent ff4d2b1
commit 57e6ba0
Showing 1 changed file with 20 additions and 12 deletions.
diff --git 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;
 
@@ -91,22 +92,21 @@ public function getInfo(string $type = null): mixed
         return $this->response->getInfo($type);
     }
 
+    /**
+     * Do not end tracing when calling toStream(). It will end when self::stream() will receive the last chunk.
+     */
     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 +129,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;
         }
     }