fix: save data on SIGTERM #1307

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

6 files changed, 119 insertions, 17 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index e9ae1d6e..dbeaca2c 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -25,6 +25,12 @@ Unreleased
 - Feature: Added the `lcov` command to generate reports in LCOV format.
   Thanks, Bradley Burns. Closes `issue 587`_ and `issue 626`_.
 
+- Feature: coverage measurement data will now be written when a SIGTERM signal
+  is received by the process.  This includes
+  :meth:`Process.terminate <python:multiprocessing.Process.terminate>`,
+  and other ways to terminate a process.  Currently this is only on Linux and
+  Mac; Windows is not supported.  Fixes `issue 1307`_.
+
 - Dropped support for Python 3.6, which reached end-of-life on 2021-12-23.
 
 - Updated Python 3.11 support to 3.11.0a4, fixing `issue 1294`_.
@@ -45,6 +51,7 @@ Unreleased
 .. _issue 1288: https://github.com/nedbat/coveragepy/issues/1288
 .. _issue 1294: https://github.com/nedbat/coveragepy/issues/1294
 .. _issue 1303: https://github.com/nedbat/coveragepy/issues/1303
+.. _issue 1307: https://github.com/nedbat/coveragepy/issues/1307
 
 
 .. _changes_62:
diff --git a/coverage/control.py b/coverage/control.py
index bd51ffc5..913d6893 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,23 @@ 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")
+        # Statements after here won't be seen by metacov because we just wrote
+        # the data, and are about to kill the process.
+        signal.signal(signal.SIGTERM, self._old_sigterm)    # pragma: not covered
+        os.kill(os.getpid(), signal.SIGTERM)                # pragma: not covered
+
     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/doc/config.rst b/doc/config.rst
index e44390e5..b82df993 100644
--- a/doc/config.rst
+++ b/doc/config.rst
@@ -132,6 +132,8 @@ option, or coverage.py will produce very wrong results.
 .. _gevent: http://www.gevent.org/
 .. _eventlet: http://eventlet.net/
 
+See :ref:subprocess: for details of multi-process measurement.
+
 Before version 4.2, this option only accepted a single string.
 
 .. versionadded:: 4.0
diff --git a/doc/subprocess.rst b/doc/subprocess.rst
index b28cbc91..777ffbae 100644
--- a/doc/subprocess.rst
+++ b/doc/subprocess.rst
@@ -25,7 +25,7 @@ the name of the :ref:`configuration file <config>` to use.
 
 .. note::
 
-    If you have subprocesses because you are using :mod:`multiprocessing
+    If you have subprocesses created with :mod:`multiprocessing
     <python:multiprocessing>`, the ``--concurrency=multiprocessing``
     command-line option should take care of everything for you.  See
     :ref:`cmd_run` for details.
@@ -34,8 +34,8 @@ When using this technique, be sure to set the parallel option to true so that
 multiple coverage.py runs will each write their data to a distinct file.
 
 
-Configuring Python for sub-process coverage
--------------------------------------------
+Configuring Python for sub-process measurement
+----------------------------------------------
 
 Measuring coverage in sub-processes is a little tricky.  When you spawn a
 sub-process, you are invoking Python to run your program.  Usually, to get
@@ -84,18 +84,17 @@ start-up.  Be sure to remove the change when you uninstall coverage.py, or use
 a more defensive approach to importing it.
 
 
-Signal handlers and atexit
---------------------------
-
-.. hmm, this isn't specifically about subprocesses, is there a better place
-    where we could talk about this?
+Process termination
+-------------------
 
 To successfully write a coverage data file, the Python sub-process under
-analysis must shut down cleanly and have a chance for coverage.py to run the
-``atexit`` handler it registers.
+analysis must shut down cleanly and have a chance for coverage.py to run its
+termination code.  It will do that when the process ends naturally, or when a
+SIGTERM signal is received.
 
-For example if you send SIGTERM to end the sub-process, but your sub-process
-has never registered any SIGTERM handler, then a coverage file won't be
-written.  See the `atexit`_ docs for details of when the handler isn't run.
+Coverage.py uses :mod:`atexit <python:atexit>` to handle usual process ends,
+and a :mod:`signal <python:signal>` handler to catch SIGTERM signals.
 
-.. _atexit: https://docs.python.org/3/library/atexit.html
+Other ways of ending a process, like SIGKILL or :func:`os._exit
+<python:os._exit>`, will prevent coverage.py from writing its data file,
+leaving you with incomplete or non-existent coverage data.
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"