diff --git a/src/TimeoutCancellation.php b/src/TimeoutCancellation.php index a73f737c..21434686 100644 --- a/src/TimeoutCancellation.php +++ b/src/TimeoutCancellation.php @@ -17,10 +17,11 @@ final class TimeoutCancellation implements Cancellation private readonly Cancellation $cancellation; /** - * @param float $timeout Seconds until cancellation is requested. - * @param string $message Message for TimeoutException. Default is "Operation timed out". + * @param float $timeout Seconds until cancellation is requested. + * @param string $message Message for TimeoutException. Default is "Operation timed out". + * @param bool $reference Whether to reference the timer. */ - public function __construct(float $timeout, string $message = "Operation timed out") + public function __construct(float $timeout, string $message = "Operation timed out", bool $reference = false) { $this->cancellation = $source = new Internal\Cancellable; @@ -37,7 +38,9 @@ public function __construct(float $timeout, string $message = "Operation timed o $source->cancel(new TimeoutException($message)); }); - EventLoop::unreference($this->watcher); + if (!$reference) { + EventLoop::unreference($this->watcher); + } } /** diff --git a/test/Cancellation/TimeoutCancellationTest.php b/test/Cancellation/TimeoutCancellationTest.php index 3b8e229f..628fbe44 100644 --- a/test/Cancellation/TimeoutCancellationTest.php +++ b/test/Cancellation/TimeoutCancellationTest.php @@ -3,6 +3,7 @@ namespace Amp\Cancellation; use Amp\CancelledException; +use Amp\DeferredFuture; use Amp\TestCase; use Amp\TimeoutCancellation; use Amp\TimeoutException; @@ -42,4 +43,20 @@ public function testWatcherCancellation(): void unset($cancellation); self::assertSame($identifiers, EventLoop::getIdentifiers()); } + + public function testWatcherUnreference(): void + { + $this->expectExceptionMessageMatches("/Event loop terminated without resuming the current suspension/"); + $deferred = new DeferredFuture; + $cancellation = new TimeoutCancellation(0.001); + $deferred->getFuture()->await($cancellation); + } + + public function testWatcherNoUnreference(): void + { + $this->expectException(CancelledException::class); + $cancellation = new TimeoutCancellation(0.001, reference: true); + $deferred = new DeferredFuture; + $deferred->getFuture()->await($cancellation); + } }