Provide cancellation token source to push handler

Related to #243.
amphp · Jan 13, 2020 · fe17e49 · fe17e49
1 parent e67c725
commit fe17e49
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/Connection/Internal/Http2ConnectionProcessor.php b/src/Connection/Internal/Http2ConnectionProcessor.php
@@ -591,17 +591,20 @@ static function () {
                 $this->releaseStream($streamId, $exception);
             });
 
+            $promise = $stream->pendingResponse->promise();
+            $promise->onResolve(function () use ($cancellationToken, $cancellationId): void {
+                $cancellationToken->unsubscribe($cancellationId);
+            });
+
             $onPush = $stream->request->getPushHandler();
 
             try {
-                yield call($onPush, $stream->request, $stream->pendingResponse->promise());
+                yield call($onPush, $stream->request, $promise, $tokenSource);
             } catch (HttpException | StreamException | CancelledException $exception) {
                 $tokenSource->cancel($exception);
             } catch (\Throwable $exception) {
                 $tokenSource->cancel($exception);
                 throw $exception;
-            } finally {
-                $cancellationToken->unsubscribe($cancellationId);
             }
         });
     }