Improve logging.

Author: ayjayt

src/py/kaleido/_sync_server.py

@@ -8,11 +8,15 @@
 from threading import Thread
 from typing import TYPE_CHECKING, NamedTuple
 
+import logistro
+
 from .kaleido import Kaleido
 
 if TYPE_CHECKING:
     from typing import Any, Callable
 
+_logger = logistro.getLogger(__name__)
+
 
 class Task(NamedTuple):
     fn: str
@@ -31,9 +35,10 @@ async def _server(self, *args, **kwargs):
         async with Kaleido(*args, **kwargs) as k:  # multiple processor? Enable GPU?
             while True:
                 task = self._task_queue.get()  # thread dies if main thread dies
+                _logger.debug(f"Got task for kaleido_sync_server: {task!s}")
                 if task is None:
-                    self._task_queue.task_done()
-                    return
+                    _logger.debug("Task was none.")
+                    break
                 if not hasattr(k, task.fn):
                     raise _BadFunctionName(f"Kaleido has no attribute {task.fn}")
                 try:
@@ -45,6 +50,9 @@ async def _server(self, *args, **kwargs):
 
                 self._task_queue.task_done()
 
+        self._task_queue.task_done()
+        return  # noqa: PLR1711 useless return, but readability
+
     def __new__(cls):
         # Create the singleton on first instantiation
         if cls._instance is None:
@@ -73,23 +81,38 @@ def open(self, *args: Any, silence_warnings=False, **kwargs: Any) -> None:
         )
         self._task_queue: Queue[Task | None] = Queue()
         self._return_queue: Queue[Any] = Queue()
+        _logger.debug("Starting kaleido_sync_server thread.")
         self._thread.start()
         self._initialized = True
-        close = partial(self.close, silence_warnings=True)
+        close = partial(self.close, silence_warnings=True, _atexit=True)
+        _logger.debug("Registering close with atexit.")
         atexit.register(close)
 
-    def close(self, *, silence_warnings=False):
+        # python bug
+        from time import sleep  # noqa: PLC0415 import at top, is hack
+
+        sleep(0.1)
+        # python seems to sometimes not like calling atext.register
+        # too close to the end of a program
+
+    def close(self, *, silence_warnings=False, _atexit=False):
         """Reset the singleton back to an uninitialized state."""
+        if _atexit:
+            _logger.debug("atexit trying to close kaleido_sync_server")
         if not self.is_running():
+            _logger.debug("Can't close kaleido_sync_server: not running.")
             if not silence_warnings:
                 warnings.warn(
                     "Server already closed.",
                     RuntimeWarning,
                     stacklevel=2,
                 )
             return
+        _logger.debug("Putting None to thread queue to end.")
         self._task_queue.put(None)
+        _logger.debug("Signaled thread to end, now going to join.")
         self._thread.join()
+        _logger.debug("Thread joined.")
         del self._thread
         del self._task_queue
         del self._return_queue