From 1cf0daec4fe038736a9e2618650bda860576d4c5 Mon Sep 17 00:00:00 2001
From: wunder957 <wakedwunder957@gmail.com>
Date: Thu, 21 Sep 2023 11:07:31 +0800
Subject: [PATCH] [BUG] Poll when shutdown (#74)

* Poll when shutdown

* Make it configuable
---
 duetector/static/config.toml |  2 ++
 duetector/tools/poller.py    | 11 +++++++++++
 tests/test_poller.py         | 25 +++++++++++++++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/duetector/static/config.toml b/duetector/static/config.toml
index 71c339b..de8bb51 100644
--- a/duetector/static/config.toml
+++ b/duetector/static/config.toml
@@ -100,6 +100,7 @@ max_workers = 10
 
 [monitor.bcc.poller]
 interval_ms = 500
+call_when_shutdown = true
 
 [monitor.sh]
 disabled = false
@@ -111,6 +112,7 @@ max_workers = 10
 
 [monitor.sh.poller]
 interval_ms = 500
+call_when_shutdown = true
 
 [server]
 token = ""
diff --git a/duetector/tools/poller.py b/duetector/tools/poller.py
index 43b213a..83eb608 100644
--- a/duetector/tools/poller.py
+++ b/duetector/tools/poller.py
@@ -20,6 +20,7 @@ class Poller(Configuable):
 
     default_config = {
         "interval_ms": 500,
+        "call_when_shutdown": True,
     }
     """
     Default config for this poller.
@@ -42,6 +43,13 @@ def interval_ms(self):
         """
         return self.config.interval_ms
 
+    @property
+    def call_when_shutdown(self):
+        """
+        Whether to call ``func`` when shutdown
+        """
+        return self.config.call_when_shutdown
+
     def start(self, func, *args, **kwargs):
         """
         Start a poller thread, until ``shutdown`` is called.
@@ -56,6 +64,9 @@ def _poll():
             while not self.shutdown_event.is_set():
                 func(*args, **kwargs)
                 self.shutdown_event.wait(timeout=self.interval_ms / 1000)
+            # call func one last time before exit
+            if self.call_when_shutdown:
+                func(*args, **kwargs)
 
         self._thread = threading.Thread(target=_poll)
         self.shutdown_event.clear()
diff --git a/tests/test_poller.py b/tests/test_poller.py
index d4ece31..eaf7469 100644
--- a/tests/test_poller.py
+++ b/tests/test_poller.py
@@ -10,6 +10,17 @@ def config():
     yield {
         "poller": {
             "interval_ms": 500,
+            "call_when_shutdown": False,
+        }
+    }
+
+
+@pytest.fixture
+def config_with_call_when_shutdown():
+    yield {
+        "poller": {
+            "interval_ms": 500,
+            "call_when_shutdown": True,
         }
     }
 
@@ -19,6 +30,11 @@ def poller(config):
     return Poller(config)
 
 
+@pytest.fixture
+def poller_with_call_when_shutdown(config_with_call_when_shutdown):
+    return Poller(config_with_call_when_shutdown)
+
+
 def test_poller(poller: Poller, capsys):
     poller.start(lambda: print("hello"))
     time.sleep(poller.interval_ms / 1000 * 1.5)
@@ -28,5 +44,14 @@ def test_poller(poller: Poller, capsys):
     assert captured.out == "hello\nhello\n"
 
 
+def test_poller_with_call_when_shutdown(poller_with_call_when_shutdown: Poller, capsys):
+    poller_with_call_when_shutdown.start(lambda: print("hello"))
+    time.sleep(poller_with_call_when_shutdown.interval_ms / 1000 * 1.5)
+    poller_with_call_when_shutdown.shutdown()
+    poller_with_call_when_shutdown.wait()
+    captured = capsys.readouterr()
+    assert captured.out == "hello\nhello\nhello\n"
+
+
 if __name__ == "__main__":
     pytest.main(["-vv", "-s", __file__])