From afc44dae16635cfa125e82093a378bed7d8a5b07 Mon Sep 17 00:00:00 2001
From: Nicolas Grekas <nicolas.grekas@gmail.com>
Date: Sat, 9 May 2020 22:11:42 +0200
Subject: [PATCH] [HttpClient] preserve the identity of responses streamed by
 TraceableHttpClient

---
 .../HttpClient/Response/TraceableResponse.php | 23 +++++++++++++++++++
 .../Tests/TraceableHttpClientTest.php         |  7 +++---
 .../HttpClient/TraceableHttpClient.php        | 11 ++-------
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/src/Symfony/Component/HttpClient/Response/TraceableResponse.php b/src/Symfony/Component/HttpClient/Response/TraceableResponse.php
index 9305e9be94..2fe78f4574 100644
--- a/src/Symfony/Component/HttpClient/Response/TraceableResponse.php
+++ b/src/Symfony/Component/HttpClient/Response/TraceableResponse.php
@@ -14,6 +14,7 @@ namespace Symfony\Component\HttpClient\Response;
 use Symfony\Component\HttpClient\Exception\ClientException;
 use Symfony\Component\HttpClient\Exception\RedirectionException;
 use Symfony\Component\HttpClient\Exception\ServerException;
+use Symfony\Component\HttpClient\TraceableHttpClient;
 use Symfony\Contracts\HttpClient\Exception\ClientExceptionInterface;
 use Symfony\Contracts\HttpClient\Exception\RedirectionExceptionInterface;
 use Symfony\Contracts\HttpClient\Exception\ServerExceptionInterface;
@@ -105,6 +106,28 @@ class TraceableResponse implements ResponseInterface
         return StreamWrapper::createResource($this->response, $this->client);
     }
 
+    /**
+     * @internal
+     */
+    public static function stream(HttpClientInterface $client, iterable $responses, ?float $timeout): \Generator
+    {
+        $wrappedResponses = [];
+        $traceableMap = new \SplObjectStorage();
+
+        foreach ($responses as $r) {
+            if (!$r instanceof self) {
+                throw new \TypeError(sprintf('"%s::stream()" expects parameter 1 to be an iterable of TraceableResponse objects, "%s" given.', TraceableHttpClient::class, get_debug_type($r)));
+            }
+
+            $traceableMap[$r->response] = $r;
+            $wrappedResponses[] = $r->response;
+        }
+
+        foreach ($client->stream($wrappedResponses, $timeout) as $r => $chunk) {
+            yield $traceableMap[$r] => $chunk;
+        }
+    }
+
     private function checkStatusCode($code)
     {
         if (500 <= $code) {
diff --git a/src/Symfony/Component/HttpClient/Tests/TraceableHttpClientTest.php b/src/Symfony/Component/HttpClient/Tests/TraceableHttpClientTest.php
index 66097d013b..8c0ec0d48c 100755
--- a/src/Symfony/Component/HttpClient/Tests/TraceableHttpClientTest.php
+++ b/src/Symfony/Component/HttpClient/Tests/TraceableHttpClientTest.php
@@ -88,11 +88,12 @@ class TraceableHttpClientTest extends TestCase
         TestHttpServer::start();
 
         $sut = new TraceableHttpClient(new NativeHttpClient());
-        $chunked = $sut->request('GET', 'http://localhost:8057/chunked');
+        $response = $sut->request('GET', 'http://localhost:8057/chunked');
         $chunks = [];
-        foreach ($sut->stream($chunked) as $response) {
-            $chunks[] = $response->getContent();
+        foreach ($sut->stream($response) as $r => $chunk) {
+            $chunks[] = $chunk->getContent();
         }
+        $this->assertSame($response, $r);
         $this->assertGreaterThan(1, \count($chunks));
         $this->assertSame('Symfony is awesome!', implode('', $chunks));
     }
diff --git a/src/Symfony/Component/HttpClient/TraceableHttpClient.php b/src/Symfony/Component/HttpClient/TraceableHttpClient.php
index f7fbfafc3b..b70c544a66 100644
--- a/src/Symfony/Component/HttpClient/TraceableHttpClient.php
+++ b/src/Symfony/Component/HttpClient/TraceableHttpClient.php
@@ -13,6 +13,7 @@ namespace Symfony\Component\HttpClient;
 
 use Psr\Log\LoggerAwareInterface;
 use Psr\Log\LoggerInterface;
+use Symfony\Component\HttpClient\Response\ResponseStream;
 use Symfony\Component\HttpClient\Response\TraceableResponse;
 use Symfony\Contracts\HttpClient\HttpClientInterface;
 use Symfony\Contracts\HttpClient\ResponseInterface;
@@ -70,15 +71,7 @@ final class TraceableHttpClient implements HttpClientInterface, ResetInterface,
             throw new \TypeError(sprintf('"%s()" expects parameter 1 to be an iterable of TraceableResponse objects, "%s" given.', __METHOD__, get_debug_type($responses)));
         }
 
-        return $this->client->stream(\Closure::bind(static function () use ($responses) {
-            foreach ($responses as $k => $r) {
-                if (!$r instanceof TraceableResponse) {
-                    throw new \TypeError(sprintf('"%s()" expects parameter 1 to be an iterable of TraceableResponse objects, "%s" given.', __METHOD__, get_debug_type($r)));
-                }
-
-                yield $k => $r->response;
-            }
-        }, null, TraceableResponse::class)(), $timeout);
+        return new ResponseStream(TraceableResponse::stream($this->client, $responses, $timeout));
     }
 
     public function getTracedRequests(): array