Handle push cancellation

Author: nicolas-grekas

src/Symfony/Component/HttpClient/AmpHttpClient.php

@@ -12,6 +12,7 @@
 namespace Symfony\Component\HttpClient;
 
 use Amp\CancelledException;
+use Amp\DeferredFuture;
 use Amp\Http\Client\DelegateHttpClient;
 use Amp\Http\Client\InterceptedHttpClient;
 use Amp\Http\Client\PooledHttpClient;
@@ -154,8 +155,9 @@ public function reset(): void
         $this->multi->dnsCache = [];
 
         foreach ($this->multi->pushedResponses as $authority => $pushedResponses) {
-            foreach ($pushedResponses as [$pushedUrl, $pushedResponse]) {
-                $pushedResponse->error(new CancelledException());
+            foreach ($pushedResponses as [$pushedUrl, $pushDeferred]) {
+                /** @var DeferredFuture $pushDeferred */
+                $pushDeferred->error(new CancelledException());
 
                 $this->logger?->debug(sprintf('Unused pushed response: "%s"', $pushedUrl));
             }

src/Symfony/Component/HttpClient/Internal/AmpClientState.php

@@ -12,6 +12,7 @@
 namespace Symfony\Component\HttpClient\Internal;
 
 use Amp\Cancellation;
+use Amp\DeferredFuture;
 use Amp\Future;
 use Amp\Http\Client\Connection\ConnectionLimitingPool;
 use Amp\Http\Client\Connection\DefaultConnectionFactory;
@@ -179,6 +180,7 @@ public function connect(SocketAddress|string $uri, ConnectContext $context = nul
 
     private function handlePush(Request $request, Future $response, array $options): void
     {
+        $deferred = new DeferredFuture();
         $authority = $request->getUri()->getAuthority();
 
         if ($this->maxPendingPushes <= \count($this->pushedResponses[$authority] ?? [])) {
@@ -189,11 +191,13 @@ private function handlePush(Request $request, Future $response, array $options):
 
         $url = (string) $request->getUri();
         $this->logger?->debug(sprintf('Queueing pushed response: "%s"', $url));
-        $this->pushedResponses[$authority][] = [$url, $response, $request, [
+        $this->pushedResponses[$authority][] = [$url, $deferred, $request, $response, [
             'proxy' => $options['proxy'],
             'bindto' => $options['bindto'],
             'local_cert' => $options['local_cert'],
             'local_pk' => $options['local_pk'],
         ]];
+
+        $deferred->getFuture()->await();
     }
 }

src/Symfony/Component/HttpClient/Response/AmpResponse.php

@@ -13,12 +13,12 @@
 
 use Amp\ByteStream\StreamException;
 use Amp\DeferredCancellation;
+use Amp\DeferredFuture;
+use Amp\Future;
 use Amp\Http\Client\HttpException;
 use Amp\Http\Client\Request;
 use Amp\Http\Client\Response;
 use Psr\Log\LoggerInterface;
-use Revolt\EventLoop;
-use Revolt\EventLoop\Suspension;
 use Symfony\Component\HttpClient\Chunk\FirstChunk;
 use Symfony\Component\HttpClient\Chunk\InformationalChunk;
 use Symfony\Component\HttpClient\Exception\InvalidArgumentException;
@@ -196,7 +196,7 @@ private static function generateResponse(Request $request, AmpClientState $multi
         });
 
         try {
-            if (null === $response = self::getPushedResponse($request, $multi, $info, $headers, $options, $logger)) {
+            if (null === $response = self::getPushedResponse($request, $multi, $info, $headers, $canceller, $options, $logger)) {
                 $logger?->info(sprintf('Request: "%s %s"', $info['http_method'], $info['url']));
 
                 $response = self::followRedirects($request, $multi, $info, $headers, $canceller, $options, $onProgress, $handle, $logger, $pause);
@@ -364,15 +364,16 @@ private static function addResponseHeaders(Response $response, array &$info, arr
     /**
      * Accepts pushed responses only if their headers related to authentication match the request.
      */
-    private static function getPushedResponse(Request $request, AmpClientState $multi, array &$info, array &$headers, array $options, ?LoggerInterface $logger): ?Response
+    private static function getPushedResponse(Request $request, AmpClientState $multi, array &$info, array &$headers, DeferredCancellation $canceller, array $options, ?LoggerInterface $logger): ?Response
     {
         if ('' !== $options['body']) {
             return null;
         }
 
         $authority = $request->getUri()->getAuthority();
+        $cancellation = $canceller->getCancellation();
 
-        foreach ($multi->pushedResponses[$authority] ?? [] as $i => [$pushedUrl, $pushedResponse, $pushedRequest, $parentOptions]) {
+        foreach ($multi->pushedResponses[$authority] ?? [] as $i => [$pushedUrl, $pushDeferred, $pushedRequest, $pushedResponse, $parentOptions]) {
             if ($info['url'] !== $pushedUrl || $info['http_method'] !== $pushedRequest->getMethod()) {
                 continue;
             }
@@ -383,13 +384,23 @@ private static function getPushedResponse(Request $request, AmpClientState $mult
                 }
             }
 
+            /** @var DeferredFuture $pushDeferred */
+            $id = $cancellation->subscribe(static fn ($e) => $pushDeferred->error($e));
+
+            try {
+                /** @var Future $pushedResponse */
+                $response = $pushedResponse->await($cancellation);
+            } finally {
+                $cancellation->unsubscribe($id);
+            }
+
             foreach (['authorization', 'cookie', 'range', 'proxy-authorization'] as $k) {
-                if ($pushedRequest->getHeaderArray($k) !== $request->getHeaderArray($k)) {
+                if ($response->getHeaderArray($k) !== $request->getHeaderArray($k)) {
                     continue 2;
                 }
             }
 
-            foreach ($pushedResponse->getHeaderArray('vary') as $vary) {
+            foreach ($response->getHeaderArray('vary') as $vary) {
                 foreach (preg_split('/\s*+,\s*+/', $vary) as $v) {
                     if ('*' === $v || ($pushedRequest->getHeaderArray($v) !== $request->getHeaderArray($v) && 'accept-encoding' !== strtolower($v))) {
                         $logger?->debug(sprintf('Skipping pushed response: "%s"', $info['url']));
@@ -398,8 +409,9 @@ private static function getPushedResponse(Request $request, AmpClientState $mult
                 }
             }
 
+            $pushDeferred->complete();
             $logger?->debug(sprintf('Accepting pushed response: "%s %s"', $info['http_method'], $info['url']));
-            self::addResponseHeaders($pushedResponse, $info, $headers);
+            self::addResponseHeaders($response, $info, $headers);
             unset($multi->pushedResponses[$authority][$i]);
 
             if (!$multi->pushedResponses[$authority]) {