scheduler.py: Notifications for sched suspending

Add a notification for SUSPEND & UNSUSPEND. This splits the
jobs_suspended context manager into two methods, which are triggered
by suspend() context manager yielding in Stream().

2 files changed, 19 insertions, 7 deletions

diff --git a/src/buildstream/_scheduler/scheduler.py b/src/buildstream/_scheduler/scheduler.py
index 082580bcf..45f6a39d4 100644
--- a/src/buildstream/_scheduler/scheduler.py
+++ b/src/buildstream/_scheduler/scheduler.py
@@ -24,7 +24,6 @@ import asyncio
 from itertools import chain
 import signal
 import datetime
-from contextlib import contextmanager
 
 # Local imports
 from .resources import Resources
@@ -57,6 +56,8 @@ class NotificationType(FastEnum):
     SCHED_START_TIME = "sched_start_time"
     RUNNING = "running"
     TERMINATED = "terminated"
+    SUSPEND = "suspend"
+    UNSUSPEND = "unsuspend"
 
 
 # Notification()
@@ -242,15 +243,17 @@ class Scheduler():
 
     # jobs_suspended()
     #
-    # A context manager for running with jobs suspended
+    # Suspend jobs after being notified
     #
-    @contextmanager
     def jobs_suspended(self):
         self._disconnect_signals()
         self._suspend_jobs()
 
-        yield
-
+    # jobs_unsuspended()
+    #
+    # Unsuspend jobs after being notified
+    #
+    def jobs_unsuspended(self):
         self._resume_jobs()
         self._connect_signals()
 
@@ -500,6 +503,10 @@ class Scheduler():
             self.terminate_jobs()
         elif notification.notification_type == NotificationType.QUIT:
             self.stop_queueing()
+        elif notification.notification_type == NotificationType.SUSPEND:
+            self.jobs_suspended()
+        elif notification.notification_type == NotificationType.UNSUSPEND:
+            self.jobs_unsuspended()
         else:
             # Do not raise exception once scheduler process is separated
             # as we don't want to pickle exceptions between processes
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index d3e1d362c..77d746ab7 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -1129,8 +1129,13 @@ class Stream():
     #
     @contextmanager
     def suspend(self):
-        with self._scheduler.jobs_suspended():
-            yield
+        # Send the notification to suspend jobs
+        notification = Notification(NotificationType.SUSPEND)
+        self._notify(notification)
+        yield
+        # Unsuspend jobs on context exit
+        notification = Notification(NotificationType.UNSUSPEND)
+        self._notify(notification)
 
     #############################################################
     #                    Private Methods                        #