feat: allow configuring worker shutdown timeout

src/ClusterWatcher.php

@@ -36,6 +36,9 @@ final class ClusterWatcher
     use ForbidCloning;
     use ForbidSerialization;
 
+    /**
+     * The default worker shutdown timeout in seconds.
+     */
     public const WORKER_TIMEOUT = 5;
 
     private readonly ContextFactory $contextFactory;
@@ -120,8 +123,10 @@ public function getMessageIterator(): ConcurrentIterator
 
     /**
      * @param int $count Number of cluster workers to spawn.
+     * @param null|int|float $workerShutdownTimeout The maximum time to wait for a worker to shut down, in seconds,
+     *  or null to wait indefinitely.
      */
-    public function start(int $count): void
+    public function start(int $count, null|int|float $workerShutdownTimeout = ClusterWatcher::WORKER_TIMEOUT): void
     {
         if ($this->running || $this->queue->isComplete()) {
             throw new \Error("The cluster watcher is already running or has already run");
@@ -137,7 +142,7 @@ public function start(int $count): void
         try {
             for ($i = 0; $i < $count; ++$i) {
                 $id = $this->nextId++;
-                $this->workers[$id] = $this->startWorker($id);
+                $this->workers[$id] = $this->startWorker($id, $workerShutdownTimeout);
             }
         } catch (\Throwable $exception) {
             $this->stop();
@@ -146,9 +151,11 @@ public function start(int $count): void
     }
 
     /**
-     * @param positive-int $id
+     * @param positive-int $id The worker ID.
+     * @param null|int|float $shutdownTimeout The maximum time to wait for the worker to shut down, in seconds,
+     *  or null to wait indefinitely.
      */
-    private function startWorker(int $id): ContextClusterWorker
+    private function startWorker(int $id, null|int|float $shutdownTimeout = ClusterWatcher::WORKER_TIMEOUT): ContextClusterWorker
     {
         $context = $this->contextFactory->start($this->script);
 
@@ -203,12 +210,13 @@ private function startWorker(int $id): ContextClusterWorker
             $socket,
             $deferredCancellation,
             $id,
+            $shutdownTimeout,
         ): void {
             async($this->provider->provideFor(...), $socket, $deferredCancellation->getCancellation())->ignore();
 
             try {
                 try {
-                    $worker->run();
+                    $worker->run($shutdownTimeout);
 
                     $worker->info("Worker {$id} terminated cleanly" .
                         ($this->running ? ", restarting..." : ""));

src/Internal/ContextClusterWorker.php

@@ -70,7 +70,13 @@ public function send(mixed $data): void
         $this->context->send(new WatcherMessage(WatcherMessageType::Data, $data));
     }
 
-    public function run(): void
+    /**
+     * Run the worker.
+     *
+     * @param null|int|float $shutdownTimeout The maximum time to wait for the worker to shut down, in seconds,
+     *    or null to wait indefinitely.
+     */
+    public function run(null|int|float $shutdownTimeout = ClusterWatcher::WORKER_TIMEOUT): void
     {
         $watcher = EventLoop::repeat(self::PING_TIMEOUT / 2, weakClosure(function (): void {
             if ($this->lastActivity < \time() - self::PING_TIMEOUT) {
@@ -109,7 +115,11 @@ public function run(): void
                 };
             }
 
-            $this->joinFuture->await(new TimeoutCancellation(ClusterWatcher::WORKER_TIMEOUT));
+            if ($shutdownTimeout === null) {
+                $this->joinFuture->await();
+            } else {
+                $this->joinFuture->await(new TimeoutCancellation($shutdownTimeout));
+            }
         } catch (\Throwable $exception) {
             $this->joinFuture->ignore();
             throw $exception;