[BUG] Poll when shutdown (#74)

* Poll when shutdown

* Make it configuable

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 = ""

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()

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__])