Fix mkdocs serve cleanup to remove temporary site directories on SIGTERM

docs/about/release-notes.md

@@ -29,6 +29,7 @@ The current members of the MkDocs-NG team.
 
 * Fix anchor link validation so `mkdocs build -v --strict` still fails when missing anchors are reported as warnings. #30
 * Fix `validation.links.not_found` is always reported as INFO for excluded pages. #32
+* Fix `mkdocs serve` cleanup so temporary site directories are removed when the process receives `SIGTERM`. #36
 
 ### Maintenance
 

mkdocs/commands/serve.py

@@ -2,6 +2,7 @@
 
 import logging
 import shutil
+import signal
 import tempfile
 from os.path import isdir, isfile, join
 from typing import TYPE_CHECKING
@@ -17,6 +18,10 @@
 log = logging.getLogger(__name__)
 
 
+class _ShutdownRequested(KeyboardInterrupt):
+    pass
+
+
 def serve(
     config_file: str | None = None,
     livereload: bool = True,
@@ -80,33 +85,42 @@ def error_handler(code) -> bytes | None:
 
     server.error_handler = error_handler
 
-    try:
-        # Perform the initial build
-        builder(config)
-
-        if livereload:
-            # Watch the documentation files, the config file and the theme files.
-            server.watch(config.docs_dir)
-            if config.config_file_path:
-                server.watch(config.config_file_path)
+    def handle_sigterm(_signum, _frame) -> None:
+        raise _ShutdownRequested
 
-            if watch_theme:
-                for d in config.theme.dirs:
-                    server.watch(d)
-
-            # Run `serve` plugin events.
-            server = config.plugins.on_serve(server, config=config, builder=builder)
-
-            for item in config.watch:
-                server.watch(item)
+    previous_sigterm_handler = signal.signal(signal.SIGTERM, handle_sigterm)
 
+    try:
         try:
-            server.serve(open_in_browser=open_in_browser)
-        except KeyboardInterrupt:
+            # Perform the initial build
+            builder(config)
+
+            if livereload:
+                # Watch the documentation files, the config file and the theme files.
+                server.watch(config.docs_dir)
+                if config.config_file_path:
+                    server.watch(config.config_file_path)
+
+                if watch_theme:
+                    for d in config.theme.dirs:
+                        server.watch(d)
+
+                # Run `serve` plugin events.
+                server = config.plugins.on_serve(server, config=config, builder=builder)
+
+                for item in config.watch:
+                    server.watch(item)
+
+            try:
+                server.serve(open_in_browser=open_in_browser)
+            except KeyboardInterrupt:
+                log.info("Shutting down...")
+            finally:
+                server.shutdown()
+        except _ShutdownRequested:
             log.info("Shutting down...")
-        finally:
-            server.shutdown()
     finally:
+        signal.signal(signal.SIGTERM, previous_sigterm_handler)
         config.plugins.on_shutdown()
         if isdir(site_dir):
             shutil.rmtree(site_dir)

mkdocs/tests/serve_tests.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python
+
+import contextlib
+import signal
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest import mock
+
+from mkdocs.commands import serve
+from mkdocs.tests.base import tempdir
+
+
+class ServeTests(unittest.TestCase):
+    @tempdir()
+    def test_sigterm_cleans_temporary_site_dir(self, temp_dir):
+        site_dir = Path(temp_dir, "site")
+        site_dir.mkdir()
+
+        config = SimpleNamespace(
+            config_file_path=None,
+            dev_addr=("127.0.0.1", 8000),
+            docs_dir=str(Path(temp_dir, "docs")),
+            plugins=mock.Mock(),
+            site_url=None,
+            theme=SimpleNamespace(dirs=[]),
+            watch=[],
+        )
+        config.plugins.on_serve.side_effect = lambda server, **kwargs: server
+
+        signal_handlers = {signal.SIGTERM: signal.SIG_DFL}
+
+        def signal_side_effect(signum, handler):
+            previous_handler = signal_handlers[signum]
+            signal_handlers[signum] = handler
+            return previous_handler
+
+        def serve_side_effect(**kwargs):
+            signal_handlers[signal.SIGTERM](signal.SIGTERM, None)
+
+        server = mock.Mock()
+        server.serve.side_effect = serve_side_effect
+
+        with contextlib.ExitStack() as stack:
+            stack.enter_context(
+                mock.patch(
+                    "mkdocs.commands.serve.tempfile.mkdtemp", return_value=site_dir
+                )
+            )
+            stack.enter_context(
+                mock.patch("mkdocs.commands.serve.load_config", return_value=config)
+            )
+            stack.enter_context(mock.patch("mkdocs.commands.serve.build"))
+            stack.enter_context(
+                mock.patch(
+                    "mkdocs.commands.serve.LiveReloadServer", return_value=server
+                )
+            )
+            mock_signal = stack.enter_context(
+                mock.patch(
+                    "mkdocs.commands.serve.signal.signal",
+                    side_effect=signal_side_effect,
+                )
+            )
+
+            serve.serve()
+
+        self.assertFalse(site_dir.exists())
+        server.shutdown.assert_called_once_with()
+        config.plugins.on_shutdown.assert_called_once_with()
+        self.assertEqual(signal_handlers[signal.SIGTERM], signal.SIG_DFL)
+        mock_signal.assert_has_calls(
+            [
+                mock.call(signal.SIGTERM, mock.ANY),
+                mock.call(signal.SIGTERM, signal.SIG_DFL),
+            ]
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()