fix: save data on SIGTERM #1307nedbat/fix-1307

This covers multiprocessing.Process.terminate(), and maybe other cases also.

3 files changed, 95 insertions, 3 deletions

diff --git a/coverage/control.py b/coverage/control.py
index bd51ffc5..623ecf6c 100644
--- a/coverage/control.py
+++ b/coverage/control.py
@@ -9,6 +9,7 @@ import contextlib
 import os
 import os.path
 import platform
+import signal
 import sys
 import time
 import warnings
@@ -228,6 +229,7 @@ class Coverage:
         self._exclude_re = None
         self._debug = None
         self._file_mapper = None
+        self._old_sigterm = None
 
         # State machine variables:
         # Have we initialized everything?
@@ -526,6 +528,11 @@ class Coverage:
         self._should_write_debug = True
 
         atexit.register(self._atexit)
+        if not env.WINDOWS:
+            # The Python docs seem to imply that SIGTERM works uniformly even
+            # on Windows, but that's not my experience, and this agrees:
+            # https://stackoverflow.com/questions/35772001/x/35792192#35792192
+            self._old_sigterm = signal.signal(signal.SIGTERM, self._on_sigterm)
 
     def _init_data(self, suffix):
         """Create a data file if we don't have one yet."""
@@ -583,15 +590,21 @@ class Coverage:
             self._collector.stop()
         self._started = False
 
-    def _atexit(self):
+    def _atexit(self, event="atexit"):
         """Clean up on process shutdown."""
         if self._debug.should("process"):
-            self._debug.write(f"atexit: pid: {os.getpid()}, instance: {self!r}")
+            self._debug.write(f"{event}: pid: {os.getpid()}, instance: {self!r}")
         if self._started:
             self.stop()
         if self._auto_save:
             self.save()
 
+    def _on_sigterm(self, signum_unused, frame_unused):
+        """A handler for signal.SIGTERM."""
+        self._atexit("sigterm")
+        signal.signal(signal.SIGTERM, self._old_sigterm)
+        os.kill(os.getpid(), signal.SIGTERM)
+
     def erase(self):
         """Erase previously collected coverage data.
 
diff --git a/coverage/multiproc.py b/coverage/multiproc.py
index 1f9225f3..3a9bd633 100644
--- a/coverage/multiproc.py
+++ b/coverage/multiproc.py
@@ -27,7 +27,7 @@ class ProcessWithCoverage(OriginalProcess):         # pylint: disable=abstract-m
         """Wrapper around _bootstrap to start coverage."""
         try:
             from coverage import Coverage       # avoid circular import
-            cov = Coverage(data_suffix=True)
+            cov = Coverage(data_suffix=True, auto_data=True)
             cov._warn_preimported_source = False
             cov.start()
             debug = cov._debug
diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py
index 001455f2..14dfc5ca 100644
--- a/tests/test_concurrency.py
+++ b/tests/test_concurrency.py
@@ -693,3 +693,82 @@ def test_thread_safe_save_data(tmpdir):
     finally:
         os.chdir(old_dir)
         should_run[0] = False
+
+
+@pytest.mark.skipif(env.WINDOWS, reason="SIGTERM doesn't work the same on Windows")
+class SigtermTest(CoverageTest):
+    """Tests of our handling of SIGTERM."""
+
+    def test_sigterm_saves_data(self):
+        # A terminated process should save its coverage data.
+        self.make_file("clobbered.py", """\
+            import multiprocessing
+            import time
+
+            def subproc(x):
+                if x.value == 3:
+                    print("THREE", flush=True)  # line 6, missed
+                else:
+                    print("NOT THREE", flush=True)
+                x.value = 0
+                time.sleep(60)
+
+            if __name__ == "__main__":
+                print("START", flush=True)
+                x = multiprocessing.Value("L", 1)
+                proc = multiprocessing.Process(target=subproc, args=(x,))
+                proc.start()
+                while x.value != 0:
+                    time.sleep(.05)
+                proc.terminate()
+                print("END", flush=True)
+            """)
+        self.make_file(".coveragerc", """\
+            [run]
+            parallel = True
+            concurrency = multiprocessing
+            """)
+        out = self.run_command("coverage run clobbered.py")
+        # Under the Python tracer on Linux, we get the "Trace function changed"
+        # message. Does that matter?
+        if "Trace function changed" in out:
+            lines = out.splitlines(True)
+            assert len(lines) == 5  # "trace function changed" and "self.warn("
+            out = "".join(lines[:3])
+        assert out == "START\nNOT THREE\nEND\n"
+        self.run_command("coverage combine")
+        out = self.run_command("coverage report -m")
+        assert self.squeezed_lines(out)[2] == "clobbered.py 17 1 94% 6"
+
+    def test_sigterm_still_runs(self):
+        # A terminated process still runs its own SIGTERM handler.
+        self.make_file("handler.py", """\
+            import multiprocessing
+            import signal
+            import time
+
+            def subproc(x):
+                print("START", flush=True)
+                def on_sigterm(signum, frame):
+                    print("SIGTERM", flush=True)
+
+                signal.signal(signal.SIGTERM, on_sigterm)
+                x.value = 0
+                time.sleep(.1)
+                print("END", flush=True)
+
+            if __name__ == "__main__":
+                x = multiprocessing.Value("L", 1)
+                proc = multiprocessing.Process(target=subproc, args=(x,))
+                proc.start()
+                while x.value != 0:
+                    time.sleep(.02)
+                proc.terminate()
+            """)
+        self.make_file(".coveragerc", """\
+            [run]
+            parallel = True
+            concurrency = multiprocessing
+            """)
+        out = self.run_command("coverage run handler.py")
+        assert out == "START\nSIGTERM\nEND\n"