From 391010039326fd016de766a93fddbdb43507d4d7 Mon Sep 17 00:00:00 2001
From: Ned Batchelder <ned@nedbatchelder.com>
Date: Sat, 22 Jan 2022 21:58:11 -0500
Subject: [PATCH] fix: save data on SIGTERM #1307

This covers multiprocessing.Process.terminate(), and maybe other cases also.
---
 coverage/control.py       | 13 +++++--
 coverage/multiproc.py     |  2 +-
 tests/test_concurrency.py | 72 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 84 insertions(+), 3 deletions(-)

diff --git a/coverage/control.py b/coverage/control.py
index bd51ffc56..6e609ae3e 100644
--- a/coverage/control.py
+++ b/coverage/control.py
@@ -9,6 +9,7 @@
 import os
 import os.path
 import platform
+import signal
 import sys
 import time
 import warnings
@@ -228,6 +229,7 @@ def __init__(
         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,7 @@ def _init_for_start(self):
         self._should_write_debug = True
 
         atexit.register(self._atexit)
+        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 +586,21 @@ def stop(self):
             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 1f9225f3e..3a9bd6339 100644
--- a/coverage/multiproc.py
+++ b/coverage/multiproc.py
@@ -27,7 +27,7 @@ def _bootstrap(self, *args, **kwargs):
         """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 001455f26..ad3c4fa89 100644
--- a/tests/test_concurrency.py
+++ b/tests/test_concurrency.py
@@ -693,3 +693,75 @@ def random_load():                              # pragma: nested
     finally:
         os.chdir(old_dir)
         should_run[0] = False
+
+
+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")
+        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"