[buildstream] 01/13: _signals.py: allow calling signal handler from non-main threads

[gi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

github-bot pushed a commit to branch bschubert/no-multiprocessing-full
in repository https://gitbox.apache.org/repos/asf/buildstream.git

commit 8d91b9739417b00de6525e7c0c49f745880c5c63
Author: Benjamin Schubert <co...@benschubert.me>
AuthorDate: Wed Jun 17 20:57:37 2020 +0000

    _signals.py: allow calling signal handler from non-main threads
    
    This modifies the signal terminator so that it can be called from any
    thread.
    
    This checks that either:
    
    - The signal handler is already in place
    - Or the caller is in the main thread, allowing to set the signal
      handler.
    
    This also removes the exact callback that was added instead of removing
    the last one, and fixes the `suspend_handler` to do the same.
    
    This is required, as we don't know which interleaving of calls will be
    done, and we can't guarantee that the last one is the right one to
    remove
---
 src/buildstream/_cas/casserver.py |  7 ++++++-
 src/buildstream/_signals.py       | 13 ++++++-------
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/buildstream/_cas/casserver.py b/src/buildstream/_cas/casserver.py
index 013fb07..04c5eb8 100644
--- a/src/buildstream/_cas/casserver.py
+++ b/src/buildstream/_cas/casserver.py
@@ -30,6 +30,7 @@ import grpc
 import click
 
 from .._protos.build.bazel.remote.asset.v1 import remote_asset_pb2_grpc
+from .. import _signals
 from .._protos.build.bazel.remote.execution.v2 import (
     remote_execution_pb2,
     remote_execution_pb2_grpc,
@@ -137,7 +138,11 @@ def create_server(repo, *, enable_push, quota, index_only, log_level=LogLevel.Le
             _ReferenceStorageServicer(casd_channel, root, enable_push=enable_push), server
         )
 
-        yield server
+        # Ensure we have the signal handler set for SIGTERM
+        # This allows threads from GRPC to call our methods that do register
+        # handlers at exit.
+        with _signals.terminator(lambda: None):
+            yield server
 
     finally:
         casd_channel.close()
diff --git a/src/buildstream/_signals.py b/src/buildstream/_signals.py
index 03b55b0..1edd445 100644
--- a/src/buildstream/_signals.py
+++ b/src/buildstream/_signals.py
@@ -80,13 +80,10 @@ def terminator_handler(signal_, frame):
 def terminator(terminate_func):
     global terminator_stack  # pylint: disable=global-statement
 
-    # Signal handling only works in the main thread
-    if threading.current_thread() != threading.main_thread():
-        yield
-        return
-
     outermost = bool(not terminator_stack)
 
+    assert threading.current_thread() == threading.main_thread() or not outermost
+
     terminator_stack.append(terminate_func)
     if outermost:
         original_handler = signal.signal(signal.SIGTERM, terminator_handler)
@@ -96,7 +93,7 @@ def terminator(terminate_func):
     finally:
         if outermost:
             signal.signal(signal.SIGTERM, original_handler)
-        terminator_stack.pop()
+        terminator_stack.remove(terminate_func)
 
 
 # Just a simple object for holding on to two callbacks
@@ -146,6 +143,8 @@ def suspendable(suspend_callback, resume_callback):
     global suspendable_stack  # pylint: disable=global-statement
 
     outermost = bool(not suspendable_stack)
+    assert threading.current_thread() == threading.main_thread() or not outermost
+
     suspender = Suspender(suspend_callback, resume_callback)
     suspendable_stack.append(suspender)
 
@@ -158,7 +157,7 @@ def suspendable(suspend_callback, resume_callback):
         if outermost:
             signal.signal(signal.SIGTSTP, original_stop)
 
-        suspendable_stack.pop()
+        suspendable_stack.remove(suspender)
 
 
 # blocked()