From f93bbd55bad7a3830e47563f313cbb8deb77731c Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Tue, 17 Jan 2023 20:06:25 +0100
Subject: [PATCH] Improve logging shutdown (#85812)

* Improve logging shutdown

* Update test
---
 homeassistant/util/logging.py | 27 +++++++++++++++------------
 tests/util/test_logging.py    | 12 +++++-------
 2 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/homeassistant/util/logging.py b/homeassistant/util/logging.py
index e493a3378fd..0595e4bb90a 100644
--- a/homeassistant/util/logging.py
+++ b/homeassistant/util/logging.py
@@ -11,7 +11,6 @@ import queue
 import traceback
 from typing import Any, TypeVar, cast, overload
 
-from homeassistant.const import EVENT_HOMEASSISTANT_CLOSE
 from homeassistant.core import HomeAssistant, callback, is_callback
 
 _T = TypeVar("_T")
@@ -35,6 +34,8 @@ class HideSensitiveDataFilter(logging.Filter):
 class HomeAssistantQueueHandler(logging.handlers.QueueHandler):
     """Process the log in another thread."""
 
+    listener: logging.handlers.QueueListener | None = None
+
     def prepare(self, record: logging.LogRecord) -> logging.LogRecord:
         """Prepare a record for queuing.
 
@@ -62,6 +63,18 @@ class HomeAssistantQueueHandler(logging.handlers.QueueHandler):
             self.emit(record)
         return return_value
 
+    def close(self) -> None:
+        """
+        Tidy up any resources used by the handler.
+
+        This adds shutdown of the QueueListener
+        """
+        super().close()
+        if not self.listener:
+            return
+        self.listener.stop()
+        self.listener = None
+
 
 @callback
 def async_activate_log_queue_handler(hass: HomeAssistant) -> None:
@@ -83,20 +96,10 @@ def async_activate_log_queue_handler(hass: HomeAssistant) -> None:
         migrated_handlers.append(handler)
 
     listener = logging.handlers.QueueListener(simple_queue, *migrated_handlers)
+    queue_handler.listener = listener
 
     listener.start()
 
-    @callback
-    def _async_stop_queue_handler(_: Any) -> None:
-        """Cleanup handler."""
-        # Ensure any messages that happen after close still get logged
-        for original_handler in migrated_handlers:
-            logging.root.addHandler(original_handler)
-        logging.root.removeHandler(queue_handler)
-        listener.stop()
-
-    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_CLOSE, _async_stop_queue_handler)
-
 
 def log_exception(format_err: Callable[..., Any], *args: Any) -> None:
     """Log an exception with additional context."""
diff --git a/tests/util/test_logging.py b/tests/util/test_logging.py
index a1fd8440971..7bb1613761a 100644
--- a/tests/util/test_logging.py
+++ b/tests/util/test_logging.py
@@ -7,7 +7,6 @@ from unittest.mock import patch
 
 import pytest
 
-from homeassistant.const import EVENT_HOMEASSISTANT_CLOSE
 from homeassistant.core import callback, is_callback
 import homeassistant.util.logging as logging_util
 
@@ -66,17 +65,16 @@ async def test_logging_with_queue_handler():
 async def test_migrate_log_handler(hass):
     """Test migrating log handlers."""
 
-    original_handlers = logging.root.handlers
-
     logging_util.async_activate_log_queue_handler(hass)
 
     assert len(logging.root.handlers) == 1
     assert isinstance(logging.root.handlers[0], logging_util.HomeAssistantQueueHandler)
 
-    hass.bus.async_fire(EVENT_HOMEASSISTANT_CLOSE)
-    await hass.async_block_till_done()
-
-    assert logging.root.handlers == original_handlers
+    # Test that the close hook shuts down the queue handler's thread
+    listener_thread = logging.root.handlers[0].listener._thread
+    assert listener_thread.is_alive()
+    logging.root.handlers[0].close()
+    assert not listener_thread.is_alive()
 
 
 @pytest.mark.no_fail_on_log_exception