Allow referencing a TimeoutCancellation

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   $unreference Whether to unreference the timer.
      */
-    public function __construct(float $timeout, string $message = "Operation timed out")
+    public function __construct(float $timeout, string $message = "Operation timed out", bool $unreference = true)
     {
         $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 ($unreference) {
+            EventLoop::unreference($this->watcher);
+        }
     }
 
     /**

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, unreference: false);
+        $deferred = new DeferredFuture;
+        $deferred->getFuture()->await($cancellation);
+    }
 }