From aa96c33cbc9b33d3638d2ee2cdec3df21e08dea1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=20M=C3=B6ller?= <n1ngu@riseup.net>
Date: Sun, 4 Feb 2024 20:32:09 +0100
Subject: [PATCH] Rewrite worker with concurrent.futures instead of
 multiprocessing

Unlike multiprocessing pools, the concurrent library pools will raise an
exception if child processes have died and can't accept further tasks. This
should cause the main worker to stop instead of stalling in an inifinite wait.

Introduce a dummy noop execution in the pool to probe if it is still operative
to mitigate the volume of lost tasks if the next submited execution would raise
a BrokenExecutor exception. This isn't 100% bulletproof because the pool might
still break between the probe and a new task being pulled.

The maxtasksperchild CLI option will only work for Python 3.11 onwards because
former concurrent.futures pools do not support this option.

Closes #11104
---
 trytond/CHANGELOG          |  1 +
 trytond/bin/trytond-worker |  8 +++++---
 trytond/trytond/worker.py  | 38 ++++++++++++++++++++++++++------------
 3 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/trytond/CHANGELOG b/trytond/CHANGELOG
index dd83141946..3165388392 100644
--- a/trytond/CHANGELOG
+++ b/trytond/CHANGELOG
@@ -1,3 +1,4 @@
+* Remove trytond-worker max tasks per child option for Python <= 3.10
 * Do not copy readonly relational fields
 * Support list of single value in sortable_values tools
 * Allow resetting password of user via RPC
diff --git a/trytond/bin/trytond-worker b/trytond/bin/trytond-worker
index f8de953109..c08efa95a7 100755
--- a/trytond/bin/trytond-worker
+++ b/trytond/bin/trytond-worker
@@ -29,6 +29,8 @@ import trytond.worker as worker
 # Import after application is configured
 from trytond.pool import Pool
 
-with commandline.pidfile(options):
-    Pool.start()
-    worker.work(options)
+# Ensure main module can be safely imported by a new interpreter
+if __name__ == '__main__':
+    with commandline.pidfile(options):
+        Pool.start()
+        worker.work(options)
diff --git a/trytond/trytond/worker.py b/trytond/trytond/worker.py
index 8bcdfe0003..5f24602b44 100644
--- a/trytond/trytond/worker.py
+++ b/trytond/trytond/worker.py
@@ -5,8 +5,9 @@ import logging
 import random
 import selectors
 import signal
+import sys
 import time
-from multiprocessing import Pool as MPool
+from concurrent import futures
 from multiprocessing import cpu_count
 
 from sql import Flavor
@@ -23,10 +24,10 @@ logger = logging.getLogger(__name__)
 
 
 class Queue(object):
-    def __init__(self, database_name, mpool):
+    def __init__(self, database_name, executor):
         self.database = backend.Database(database_name).connect()
         self.connection = self.database.get_connection(autocommit=True)
-        self.mpool = mpool
+        self.executor = executor
 
     def pull(self, name=None):
         database_list = Pool.database_list()
@@ -38,17 +39,21 @@ class Queue(object):
         return Queue.pull(self.database, self.connection, name=name)
 
     def run(self, task_id):
-        return self.mpool.apply_async(run_task, (self.database.name, task_id))
+        return self.executor.submit(run_task, self.database.name, task_id)
 
 
 class TaskList(list):
     def filter(self):
         for t in list(self):
-            if t.ready():
+            if not t.running():
                 self.remove(t)
         return self
 
 
+def _noop():
+    pass
+
+
 def work(options):
     Flavor.set(backend.Database.flavor)
     if not config.getboolean('queue', 'worker', default=False):
@@ -58,10 +63,16 @@ def work(options):
     except NotImplementedError:
         processes = 1
     logger.info("start %d workers", processes)
-    mpool = MPool(
-        processes, initializer, (options.database_names,),
-        options.maxtasksperchild)
-    queues = [Queue(name, mpool) for name in options.database_names]
+    executor_options = dict(
+        max_workers=processes,
+        mp_context=None,
+        initializer=initializer,
+        initargs=(options.database_names,),
+        max_tasks_per_child=options.maxtasksperchild,
+        )
+    if sys.version_info < (3, 11):
+        del executor_options["max_tasks_per_child"]
+
 
     tasks = TaskList()
     selector = selectors.DefaultSelector()
@@ -73,7 +84,11 @@ def work(options):
             # Add some randomness to avoid concurrent pulling
             time.sleep(0.1 * random.random())
             while len(tasks.filter()) >= processes:
-                time.sleep(0.1)
+                futures.wait(tasks, return_when=futures.FIRST_COMPLETED)
+
+            # Probe process pool is still operative before pulling a new task
+            executor.submit(_noop).result()
+
             for queue in queues:
                 try:
                     task_id, next_ = queue.pull(options.name)
@@ -90,9 +105,8 @@ def work(options):
                     connection.poll()
                     while connection.notifies:
                         connection.notifies.pop(0)
-    except KeyboardInterrupt:
-        mpool.close()
     finally:
+        executor.shutdown()
         selector.close()