From a0e259a6c7fba833e9709e077bc146627acebcaf Mon Sep 17 00:00:00 2001
From: Deeka Wong <huangdijia@gmail.com>
Date: Sun, 4 Feb 2024 12:16:27 +0800
Subject: [PATCH] Improve command-signals, adds unregister support (#547)

* Improve command-signals, adds unregister support

* Refactor signal handling logic in SignalRegistry class

* Add pushSignalHandler method to SignalRegistry class

* Refactor untrap method to accept optional signo parameter

---------

Co-authored-by: Deeka Wong <8337659+huangdijia@users.noreply.github.com>
---
 .vscode/cspell.json                           |  1 +
 src/command-signals/src/SignalRegistry.php    | 60 ++++++++++++-------
 .../src/Traits/InteractsWithSignals.php       | 24 ++++----
 3 files changed, 50 insertions(+), 35 deletions(-)

diff --git a/.vscode/cspell.json b/.vscode/cspell.json
index c130a4d82..0fb129d74 100644
--- a/.vscode/cspell.json
+++ b/.vscode/cspell.json
@@ -51,6 +51,7 @@
         "Traceparent",
         "ucsplit",
         "undot",
+        "untrap",
         "Wireable"
     ]
 }
diff --git a/src/command-signals/src/SignalRegistry.php b/src/command-signals/src/SignalRegistry.php
index bced0b71a..6eba98902 100644
--- a/src/command-signals/src/SignalRegistry.php
+++ b/src/command-signals/src/SignalRegistry.php
@@ -19,6 +19,9 @@
 
 class SignalRegistry
 {
+    /**
+     * @var array<int, callable[]>
+     */
     protected array $signalHandlers = [];
 
     /**
@@ -26,55 +29,66 @@ class SignalRegistry
      */
     protected array $handling = [];
 
-    protected bool $unregistered = false;
-
-    public function __construct(protected int $timeout = 1, protected int $concurrent = 0)
+    public function __construct(protected int $timeout = 1, protected int $concurrentLimit = 0)
     {
     }
 
+    /**
+     * @param int|int[] $signo
+     * @param callable(int $signo): void $signalHandler
+     */
     public function register(int|array $signo, callable $signalHandler): void
     {
         if (is_array($signo)) {
-            array_map(fn ($s) => $this->register((int) $s, $signalHandler), $signo);
+            array_map(fn ($s) => $this->pushSignalHandler((int) $s, $signalHandler), $signo);
             return;
         }
 
-        $this->signalHandlers[$signo] ??= [];
-        $this->signalHandlers[$signo][] = $signalHandler;
-
-        if ($this->isSignalHandling($signo)) {
-            return;
-        }
+        $this->pushSignalHandler($signo, $signalHandler);
 
         $this->handleSignal($signo);
     }
 
-    public function unregister(): void
+    /**
+     * @param int|int[]|null $signo
+     */
+    public function unregister(int|array|null $signo = null): void
     {
-        $this->unregistered = true;
+        match (true) {
+            // Unregister all signals
+            is_null($signo) => $this->signalHandlers = [],
+            // Unregister multiple signals
+            is_array($signo) => array_map(fn ($s) => isset($this->signalHandlers[$s]) && $this->signalHandlers[$s] = [], $signo),
+            // Unregister single signal
+            default => $this->signalHandlers[$signo] = [],
+        };
+    }
+
+    /**
+     * @param callable(int $signo): void $signalHandler
+     */
+    protected function pushSignalHandler(int $signo, callable $signalHandler): void
+    {
+        $this->signalHandlers[$signo] ??= [];
+        $this->signalHandlers[$signo][] = $signalHandler;
     }
 
     protected function handleSignal(int $signo): void
     {
+        if (isset($this->handling[$signo])) {
+            return;
+        }
+
         $this->handling[$signo] = Coroutine::create(function () use ($signo) {
             defer(fn () => posix_kill(posix_getpid(), $signo));
 
             while (true) {
                 if (Signal::wait($signo, $this->timeout)) {
-                    $callbacks = array_map(fn ($callback) => fn () => $callback($signo), $this->signalHandlers[$signo]);
-
-                    return parallel($callbacks, $this->concurrent);
-                }
+                    $callbacks = array_map(fn ($callback) => fn () => $callback($signo), $this->signalHandlers[$signo] ?? []);
 
-                if ($this->unregistered) {
-                    break;
+                    return parallel($callbacks, $this->concurrentLimit);
                 }
             }
         });
     }
-
-    protected function isSignalHandling(int $signo): bool
-    {
-        return isset($this->handling[$signo]);
-    }
 }
diff --git a/src/command-signals/src/Traits/InteractsWithSignals.php b/src/command-signals/src/Traits/InteractsWithSignals.php
index 3023b1e09..af0a37555 100644
--- a/src/command-signals/src/Traits/InteractsWithSignals.php
+++ b/src/command-signals/src/Traits/InteractsWithSignals.php
@@ -12,32 +12,32 @@
 namespace FriendsOfHyperf\CommandSignals\Traits;
 
 use FriendsOfHyperf\CommandSignals\SignalRegistry;
-use Psr\Container\ContainerExceptionInterface;
-use Psr\Container\NotFoundExceptionInterface;
-use TypeError;
 
 use function Hyperf\Coroutine\defer;
 use function Hyperf\Support\make;
 
 trait InteractsWithSignals
 {
-    protected ?SignalRegistry $SignalRegistry = null;
+    protected ?SignalRegistry $signalRegistry = null;
 
     /**
      * Define a callback to be run when the given signal(s) occurs.
      *
-     * @param callable(int $signal): void $callback
-     * @throws TypeError
-     * @throws NotFoundExceptionInterface
-     * @throws ContainerExceptionInterface
+     * @param int|int[] $signo
+     * @param callable(int $signo): void $callback
      */
     protected function trap(array|int $signo, callable $callback): void
     {
-        if (! $this->SignalRegistry) {
-            $this->SignalRegistry = make(SignalRegistry::class);
-            defer(fn () => $this->SignalRegistry->unregister());
+        if (! $this->signalRegistry) {
+            $this->signalRegistry = make(SignalRegistry::class);
+            defer(fn () => $this->signalRegistry->unregister());
         }
 
-        $this->SignalRegistry->register($signo, $callback);
+        $this->signalRegistry->register($signo, $callback);
+    }
+
+    protected function untrap(array|int|null $signo = null): void
+    {
+        $this->signalRegistry?->unregister($signo);
     }
 }