issue #147, pool: do not block when trying to terminate tasks

In order to stop cancelled or timing out tasks, the Pool Manager
loop needs to acquire the channel lock to prevent any worker from
accessing either the lock or the pipe during its termination.

The previous logic would try to acquire the lock indefinitely. This
led to possible deadlocks in case the Pool Manager was trying to stop
cancelled tasks with the results channel filled up while the pool was
terminated. This was due to the workers holding the channel lock while
being unable to push the data as the Message Manager loop was not
reading it anymore.

The new Pool Manager loop does not try to acquire the lock indefinitely but
rather polls it for availability. This allows the loop to continue
assessing its state and correctly terminate upon request.

The drawback of this implementation is that timing out or cancelled
tasks might be stopped a little later in case the pipe is busy
transferring lots of data.

Author: noxdafox

pebble/pool/channel.py

@@ -114,11 +114,16 @@ def windows_send(obj: Any):
 
         return unix_send if os.name != 'nt' else windows_send
 
-    @property
     @contextmanager
-    def lock(self):
-        with self.mutex:
-            yield self
+    def lock(self, block: bool = True, timeout: int = None) -> bool:
+        """Lock the channel, yields True if channel is locked."""
+        acquired = self.mutex.acquire(block=block, timeout=timeout)
+
+        try:
+            yield acquired
+        finally:
+            if acquired:
+                self.mutex.release()
 
     def initialize(self):
         """Close unused connections."""
@@ -151,32 +156,45 @@ def __exit__(self, *_):
         self.release()
 
     def _make_acquire_method(self) -> Callable:
-        def unix_acquire() -> bool:
-            return (
-                self.reader_mutex.acquire(timeout=CONSTS.channel_lock_timeout)
-                and
-                self.writer_mutex.acquire(timeout=CONSTS.channel_lock_timeout)
-            )
+        def unix_acquire(
+                block: bool = True, timeout: int = CONSTS.channel_lock_timeout
+        ) -> bool:
+            """Acquire both locks. Returns True if both locks where acquired.
+            Otherwise, handle the locks state.
+
+            """
+            if self.reader_mutex.acquire(block=block, timeout=timeout):
+                if self.writer_mutex.acquire(block=block, timeout=timeout):
+                    return True
+
+                self.reader_mutex.release()
+
+            return False
 
-        def windows_acquire() -> bool:
-            return self.reader_mutex.acquire(
-                timeout=CONSTS.channel_lock_timeout)
+        def windows_acquire(
+                block: bool = True, timeout: int = CONSTS.channel_lock_timeout
+        ) -> bool:
+            """Acquire the reader lock (on NT OS, writes are atomic)."""
+            return self.reader_mutex.acquire(block=block, timeout=timeout)
 
-        return unix_acquire if os.name != 'nt' else windows_acquire
+        return windows_acquire if os.name == 'nt' else unix_acquire
 
     def _make_release_method(self) -> Callable:
         def unix_release():
+            """Release both the locks."""
             self.reader_mutex.release()
             self.writer_mutex.release()
 
         def windows_release():
+            """Release the reader lock (on NT OS, writes are atomic)."""
             self.reader_mutex.release()
 
-        return unix_release if os.name != 'nt' else windows_release
+        return windows_release if os.name == 'nt' else unix_release
 
     @property
     @contextmanager
     def reader(self):
+        """Reader lock context manager."""
         if self.reader_mutex.acquire(timeout=CONSTS.channel_lock_timeout):
             try:
                 yield self
@@ -188,6 +206,7 @@ def reader(self):
     @property
     @contextmanager
     def writer(self):
+        """Writer lock context manager."""
         if self.writer_mutex.acquire(timeout=CONSTS.channel_lock_timeout):
             try:
                 yield self

pebble/pool/process.py

@@ -210,7 +210,7 @@ def start(self):
 
     def stop(self):
         self.worker_manager.close_channels()
-        self.worker_manager.stop_workers()
+        self.worker_manager.force_stop_workers()
 
     def schedule(self, task: Task):
         """Schedules a new Task in the PoolManager."""
@@ -238,16 +238,16 @@ def update_status(self):
     def update_tasks(self):
         """Handles timing out Tasks."""
         for task in self.task_manager.timeout_tasks():
-            self.worker_manager.stop_worker(task.worker_id)
-            self.task_manager.task_done(
-                task.id,
-                Result(ResultStatus.FAILURE,
-                       TimeoutError("Task timeout", task.timeout)))
+            if self.worker_manager.maybe_stop_worker(task.worker_id):
+                self.task_manager.task_done(
+                    task.id,
+                    Result(ResultStatus.FAILURE,
+                           TimeoutError("Task timeout", task.timeout)))
 
         for task in self.task_manager.cancelled_tasks():
-            self.worker_manager.stop_worker(task.worker_id)
-            self.task_manager.task_done(
-                task.id, Result(ResultStatus.FAILURE, CancelledError()))
+            if self.worker_manager.maybe_stop_worker(task.worker_id):
+                self.task_manager.task_done(
+                    task.id, Result(ResultStatus.FAILURE, CancelledError()))
 
     def update_workers(self):
         """Handles unexpected processes termination."""
@@ -390,9 +390,9 @@ def close_channels(self):
         self.pool_channel.close()
         self.workers_channel.close()
 
-    def stop_workers(self):
+    def force_stop_workers(self):
         for worker_id in tuple(self.workers.keys()):
-            self.stop_worker(worker_id, force=True)
+            stop_process(self.workers.pop(worker_id))
 
     def new_worker(self):
         try:
@@ -403,17 +403,18 @@ def new_worker(self):
         except OSError as error:
             raise BrokenProcessPool from error
 
-    def stop_worker(self, worker_id: int, force=False):
-        try:
-            if force:
-                stop_process(self.workers.pop(worker_id))
-            else:
-                with self.workers_channel.lock:
-                    stop_process(self.workers.pop(worker_id))
-        except ChannelError as error:
-            raise BrokenProcessPool from error
-        except KeyError:
-            return  # worker already expired
+    def maybe_stop_worker(self, worker_id: int) -> bool:
+        """Try to stop the assigned worker.
+        Returns True if the worker existed and could be stopped.
+
+        """
+        with self.workers_channel.lock(block=False) as locked:
+            if locked:
+                worker = self.workers.pop(worker_id, None)
+                if worker is not None:  # Worker have already ended
+                    stop_process(worker)
+
+            return locked
 
 
 def worker_process(params: Worker, channel: WorkerChannel):
@@ -465,7 +466,7 @@ def fetch_task(channel: WorkerChannel) -> Task:
 
 def task_transaction(channel: WorkerChannel) -> Task:
     """Ensures a task is fetched and acknowledged atomically."""
-    with channel.lock:
+    with channel.lock():
         if channel.poll(0):
             task = channel.recv()
             channel.send(Acknowledgement(os.getpid(), task.id))