Repository: zalando/patroni
Pull request #2710 (2023-06-09 09:17:35)

Fix a few concurrency bugs in Citus support
- the `_in_fligh` attribute is accessed from multiple threads and must be protected with mutex when it is changed
- allow adding tasks for `_in_fligh.group` from the `sync_pg_dist_node()` method when timeout is reached. Not doing so might result is indefinite transaction if REST API request from worker node failed.

coveralls: ## Pull Request Test Coverage Report for [Build 5220307179](https://coveralls.io/builds/60630688)

* **15** of **15**   **(100.0%)**  changed or added relevant lines in **1** file are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage remained the same at **99.858%**

---



|  Totals | [![Coverage Status](https://coveralls.io/builds/60630688/badge)](https://coveralls.io/builds/60630688) |
| :-- | --: |
| Change from base [Build 5186080925](https://coveralls.io/builds/60547136): |  0.0% |
| Covered Lines: | 12636 |
| Relevant Lines: | 12654 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/patroni/async_executor.py b/patroni/async_executor.py
index 39f11013..f7fac02a 100644
--- a/patroni/async_executor.py
+++ b/patroni/async_executor.py
@@ -1,3 +1,4 @@
+"""Implement facilities for executing asynchronous tasks."""
 import logging
 
 from threading import Event, Lock, RLock, Thread
@@ -13,14 +14,23 @@ class CriticalTask(object):
     """Represents a critical task in a background process that we either need to cancel or get the result of.
 
     Fields of this object may be accessed only when holding a lock on it. To perform the critical task the background
-    thread must, while holding lock on this object, check `is_cancelled` flag, run the task and mark the task as
-    complete using `complete()`.
+    thread must, while holding lock on this object, check ``is_cancelled`` flag, run the task and mark the task as
+    complete using :func:`complete`.
 
     The main thread must hold async lock to prevent the task from completing, hold lock on critical task object,
-    call cancel. If the task has completed `cancel()` will return False and `result` field will contain the result of
-    the task. When cancel returns True it is guaranteed that the background task will notice the `is_cancelled` flag.
+    call :func:`cancel`. If the task has completed :func:`cancel` will return ``False`` and ``result`` field will
+    contain the result of the task. When :func:`cancel` returns ``True`` it is guaranteed that the background task will
+    notice the ``is_cancelled`` flag.
+
+    :ivar is_cancelled: if the critical task has been cancelled.
+    :ivar result: contains the result of the task, if it has already been completed.
     """
+
     def __init__(self) -> None:
+        """Create a new instance of :class:`CriticalTask`.
+
+        Instantiate the lock and the task control attributes.
+        """
         self._lock = Lock()
         self.is_cancelled = False
         self.result = None
@@ -28,37 +38,59 @@ class CriticalTask(object):
     def reset(self) -> None:
         """Must be called every time the background task is finished.
 
-        Must be called from async thread. Caller must hold lock on async executor when calling."""
+        .. note::
+            Must be called from async thread. Caller must hold lock on async executor when calling.
+        """
         self.is_cancelled = False
         self.result = None
 
     def cancel(self) -> bool:
-        """Tries to cancel the task, returns True if the task has already run.
+        """Tries to cancel the task.
 
-        Caller must hold lock on async executor and the task when calling."""
+        .. note::
+            Caller must hold lock on async executor and the task when calling.
+
+        :returns: ``False`` if the task has already run, or ``True`` it has been cancelled.
+        """
         if self.result is not None:
             return False
         self.is_cancelled = True
         return True
 
     def complete(self, result: Any) -> None:
-        """Mark task as completed along with a result.
+        """Mark task as completed along with a *result*.
 
-        Must be called from async thread. Caller must hold lock on task when calling."""
+        .. note::
+            Must be called from async thread. Caller must hold lock on task when calling.
+        """
         self.result = result
 
     def __enter__(self) -> 'CriticalTask':
+        """Acquire the object lock when entering the context manager."""
         self._lock.acquire()
         return self
 
     def __exit__(self, exc_type: Optional[Type[BaseException]],
                  exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]) -> None:
+        """Release the object lock when exiting the context manager."""
         self._lock.release()
 
 
 class AsyncExecutor(object):
+    """Asynchronous executor of (long) tasks.
+
+    :ivar critical_task: a :class:`CriticalTask` instance to handle execution of critical background tasks.
+    """
 
     def __init__(self, cancellable: CancellableSubprocess, ha_wakeup: Callable[..., None]) -> None:
+        """Create a new instance of :class:`AsyncExecutor`.
+
+        Configure the given *cancellable* and *ha_wakeup*, initializes the control attributes, and instantiate the lock
+        and event objects that are used to access attributes and manage communication between threads.
+
+        :param cancellable: a subprocess that supports being cancelled.
+        :param ha_wakeup: function to wake up the HA loop.
+        """
         self._cancellable = cancellable
         self._ha_wakeup = ha_wakeup
         self._thread_lock = RLock()
@@ -70,9 +102,22 @@ class AsyncExecutor(object):
 
     @property
     def busy(self) -> bool:
+        """``True`` if there is an action scheduled to occur, else ``False``."""
         return self.scheduled_action is not None
 
     def schedule(self, action: str) -> Optional[str]:
+        """Schedule *action* to be executed.
+
+        .. note::
+            Must be called before executing a task.
+
+        .. note::
+            *action* can only be scheduled if there is no other action currently scheduled.
+
+        :param action: action to be executed.
+
+        :returns: ``None`` if *action* has been successfully scheduled, or the previously scheduled action, if any.
+        """
         with self._scheduled_action_lock:
             if self._scheduled_action is not None:
                 return self._scheduled_action
@@ -83,14 +128,32 @@ class AsyncExecutor(object):
 
     @property
     def scheduled_action(self) -> Optional[str]:
+        """The currently scheduled action, if any, else ``None``."""
         with self._scheduled_action_lock:
             return self._scheduled_action
 
     def reset_scheduled_action(self) -> None:
+        """Unschedule a previously scheduled action, if any.
+
+        .. note::
+            Must be called once the scheduled task finishes or is cancelled.
+        """
         with self._scheduled_action_lock:
             self._scheduled_action = None
 
-    def run(self, func: Callable[..., Any], args: Tuple[Any, ...] = ()) -> Optional[bool]:
+    def run(self, func: Callable[..., Any], args: Tuple[Any, ...] = ()) -> Optional[Any]:
+        """Run *func* with *args*.
+
+        .. note::
+            Expected to be executed through a thread.
+
+        :param func: function to be run. If it returns anything other than ``None``, HA loop will be woken up at the end
+            of :func:`run` execution.
+        :param args: arguments to be passed to *func*.
+
+        :returns: ``None`` if *func* execution has been cancelled or faced any exception, otherwise the result of
+            *func*.
+        """
         wakeup = False
         try:
             with self:
@@ -114,15 +177,36 @@ class AsyncExecutor(object):
                 self._ha_wakeup()
 
     def run_async(self, func: Callable[..., Any], args: Tuple[Any, ...] = ()) -> None:
+        """Start an async thread that runs *func* with *args*.
+
+        :param func: function to be run. Will be passed through args to :class:`~threading.Thread` with a target of
+            :func:`run`.
+        :param args: arguments to be passed along to :class:`~threading.Thread` with *func*.
+
+        """
         Thread(target=self.run, args=(func, args)).start()
 
     def try_run_async(self, action: str, func: Callable[..., Any], args: Tuple[Any, ...] = ()) -> Optional[str]:
+        """Try to run an async task, if none is currently being executed.
+
+        :param action: name of the task to be executed.
+        :param func: actual function that performs the task *action*.
+        :param args: arguments to be passed to *func*.
+
+        :returns: ``None`` if *func* was scheduled successfully, otherwise an error message informing of an already
+            ongoing task.
+        """
         prev = self.schedule(action)
         if prev is None:
             return self.run_async(func, args)
         return 'Failed to run {0}, {1} is already in progress'.format(action, prev)
 
     def cancel(self) -> None:
+        """Request cancellation of a scheduled async task, if any.
+
+        .. note::
+            Wait until task is cancelled before returning control to caller.
+        """
         with self:
             with self._scheduled_action_lock:
                 if self._scheduled_action is None:
@@ -137,9 +221,15 @@ class AsyncExecutor(object):
             self.reset_scheduled_action()
 
     def __enter__(self) -> 'AsyncExecutor':
+        """Acquire the thread lock when entering the context manager."""
         self._thread_lock.acquire()
         return self
 
     def __exit__(self, exc_type: Optional[Type[BaseException]],
                  exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]) -> None:
+        """Release the thread lock when exiting the context manager.
+
+        .. note::
+            The arguments are not used, but we need them to match the expected method signature.
+        """
         self._thread_lock.release()
diff --git a/patroni/postgresql/citus.py b/patroni/postgresql/citus.py
index 740369ef..e11c206c 100644
--- a/patroni/postgresql/citus.py
+++ b/patroni/postgresql/citus.py
@@ -76,8 +76,9 @@ class CitusHandler(Thread):
         self._connection = Connection()
         self._pg_dist_node: Dict[int, PgDistNode] = {}  # Cache of pg_dist_node: {groupid: PgDistNode()}
         self._tasks: List[PgDistNode] = []  # Requests to change pg_dist_node, every task is a `PgDistNode`
-        self._condition = Condition()  # protects _pg_dist_node, _tasks, and _schedule_load_pg_dist_node
         self._in_flight: Optional[PgDistNode] = None  # Reference to the `PgDistNode` being changed in a transaction
+        self._schedule_load_pg_dist_node = True  # Flag that "pg_dist_node" should be queried from the database
+        self._condition = Condition()  # protects _pg_dist_node, _tasks, _in_flight, and _schedule_load_pg_dist_node
         self.schedule_cache_rebuild()
 
     def is_enabled(self) -> bool:
@@ -117,7 +118,8 @@ class CitusHandler(Thread):
         except Exception as e:
             logger.error('Exception when executing query "%s", (%s): %r', sql, params, e)
             self._connection.close()
-            self._in_flight = None
+            with self._condition:
+                self._in_flight = None
             self.schedule_cache_rebuild()
             raise e
 
@@ -215,17 +217,20 @@ class CitusHandler(Thread):
     def process_task(self, task: PgDistNode) -> bool:
         """Updates a single row in `pg_dist_node` table, optionally in a transaction.
 
-        The transaction is started if we do a demote of the worker node
-        or before promoting the other worker if there is not transaction
-        in progress. And, the transaction it is committed when the
-        switchover/failover completed.
+        The transaction is started if we do a demote of the worker node or before promoting the other worker if
+        there is no transaction in progress. And, the transaction is committed when the switchover/failover completed.
 
-        This method returns `True` if node was updated (optionally,
-        transaction was committed) as an indicator that
-        the `self._pg_dist_node` cache should be updated.
+        .. note:
+            The maximum lifetime of the transaction in progress is controlled outside of this method.
 
-        The maximum lifetime of the transaction in progress
-        is controlled outside of this method."""
+        .. note:
+            Read access to `self._in_flight` isn't protected because we know it can't be changed outside of our thread.
+
+        :param task: reference to a :class:`PgDistNode` object that represents a row to be updated/created.
+        :returns: `True` if the row was succesfully created/updated or transaction in progress
+            was committed as an indicator that the `self._pg_dist_node` cache should be updated,
+            or, if the new transaction was opened, this method returns `False`.
+        """
 
         if task.event == 'after_promote':
             # The after_promote may happen without previous before_demote and/or
@@ -236,7 +241,6 @@ class CitusHandler(Thread):
                 self.update_node(task)
             if self._in_flight:
                 self.query('COMMIT')
-                self._in_flight = None
             return True
         else:  # before_demote, before_promote
             if task.timeout:
@@ -244,11 +248,11 @@ class CitusHandler(Thread):
             if not self._in_flight:
                 self.query('BEGIN')
             self.update_node(task)
-            self._in_flight = task
         return False
 
     def process_tasks(self) -> None:
         while True:
+            # Read access to `_in_flight` isn't protected because we know it can't be changed outside of our thread.
             if not self._in_flight and not self.load_pg_dist_node():
                 break
 
@@ -259,11 +263,17 @@ class CitusHandler(Thread):
                 update_cache = self.process_task(task)
             except Exception as e:
                 logger.error('Exception when working with pg_dist_node: %r', e)
-                update_cache = False
+                update_cache = None
             with self._condition:
                 if self._tasks:
                     if update_cache:
                         self._pg_dist_node[task.group] = task
+
+                    if update_cache is False:  # an indicator that process_tasks has started a transaction
+                        self._in_flight = task
+                    else:
+                        self._in_flight = None
+
                     if id(self._tasks[i]) == id(task):
                         self._tasks.pop(i)
             task.wakeup()
@@ -293,11 +303,19 @@ class CitusHandler(Thread):
         with self._condition:
             i = self.find_task_by_group(task.group)
 
-            # task.timeout is None is an indicator that it was scheduled
-            # from the sync_pg_dist_node() and we don't want to override
-            # already existing task created from REST API.
-            if task.timeout is None and (i is not None or self._in_flight and self._in_flight.group == task.group):
-                return False
+            # The `PgDistNode.timeout` == None is an indicator that it was scheduled from the sync_pg_dist_node().
+            if task.timeout is None:
+                # We don't want to override the already existing task created from REST API.
+                if i is not None and self._tasks[i].timeout is not None:
+                    return False
+
+                # There is a little race condition with tasks created from REST API - the call made "before" the member
+                # key is updated in DCS. Therefore it is possible that :func:`sync_pg_dist_node` will try to create a
+                # task based on the outdated values of "state"/"role". To solve it we introduce an artificial timeout.
+                # Only when the timeout is reached new tasks could be scheduled from sync_pg_dist_node()
+                if self._in_flight and self._in_flight.group == task.group and self._in_flight.timeout is not None\
+                        and self._in_flight.deadline > time.time():
+                    return False
 
             # Override already existing task for the same worker group
             if i is not None:


Test patch:
diff --git a/tests/test_citus.py b/tests/test_citus.py
index a2d096ba..40d4df8a 100644
--- a/tests/test_citus.py
+++ b/tests/test_citus.py
@@ -1,3 +1,4 @@
+import time
 from mock import Mock, patch
 from patroni.postgresql.citus import CitusHandler
 
@@ -16,7 +17,7 @@ class TestCitus(BaseTestPostgresql):
         self.cluster = get_cluster_initialized_with_leader()
         self.cluster.workers[1] = self.cluster
 
-    @patch('time.time', Mock(side_effect=[100, 130, 160, 190, 220, 250, 280, 310]))
+    @patch('time.time', Mock(side_effect=[100, 130, 160, 190, 220, 250, 280, 310, 340, 370]))
     @patch('patroni.postgresql.citus.logger.exception', Mock(side_effect=SleepException))
     @patch('patroni.postgresql.citus.logger.warning')
     @patch('patroni.postgresql.citus.PgDistNode.wait', Mock())
@@ -66,11 +67,14 @@ class TestCitus(BaseTestPostgresql):
             mock_logger.assert_called_once()
             self.assertTrue(mock_logger.call_args[0][0].startswith('Overriding existing task:'))
 
-        # add_task called from sync_pg_dist_node should not override already scheduled or in flight task
+        # add_task called from sync_pg_dist_node should not override already scheduled or in flight task until deadline
         self.assertIsNotNone(self.c.add_task('after_promote', 1, 'postgres://host:5432/postgres', 30))
         self.assertIsNone(self.c.add_task('after_promote', 1, 'postgres://host:5432/postgres'))
         self.c._in_flight = self.c._tasks.pop()
+        self.c._in_flight.deadline = self.c._in_flight.timeout + time.time()
         self.assertIsNone(self.c.add_task('after_promote', 1, 'postgres://host:5432/postgres'))
+        self.c._in_flight.deadline = 0
+        self.assertIsNotNone(self.c.add_task('after_promote', 1, 'postgres://host:5432/postgres'))
 
         # If there is no transaction in progress and cached pg_dist_node matching desired state task should not be added
         self.c._schedule_load_pg_dist_node = False