Issue CHEKPOINT explicitely after promote happened (#1498)

It is safe to call pg_rewind on the replica only when pg_control on the primary contains information about the latest timeline. Postgres is usually doing immediate checkpoint right after promote and in most cases it works just fine. Unfortunately we regularly receive complaints that it takes to long (minutes) until the checkpoint is done and replicas can't perform rewind. At the same time doing the checkpoint manually immediately helped. So Patroni starts doing the same. When the promotion happened and postgres is not running in recovery, we explicitly issue the checkpoint.

We are intentionally not using the AsyncExecutor here, because we want the HA loop continues doing its normal flow.

patroni/ha.py

@@ -945,7 +945,7 @@ class Ha(object):
                     return msg
 
                 # check if the node is ready to be used by pg_rewind
-                self._rewind.check_for_checkpoint_after_promote()
+                self._rewind.ensure_checkpoint_after_promote()
 
                 if self.is_standby_cluster():
                     # in case of standby cluster we don't really need to

patroni/postgresql/rewind.py

@@ -2,9 +2,12 @@ import logging
 import os
 import subprocess
 
-from patroni.dcs import Leader
-from patroni.postgresql.connection import get_connection_cursor
-from patroni.postgresql.misc import parse_history, parse_lsn
+from threading import Lock, Thread
+
+from .connection import get_connection_cursor
+from .misc import parse_history, parse_lsn
+from ..async_executor import CriticalTask
+from ..dcs import Leader
 
 logger = logging.getLogger(__name__)
 
@@ -16,6 +19,7 @@ class Rewind(object):
 
     def __init__(self, postgresql):
         self._postgresql = postgresql
+        self._checkpoint_task_lock = Lock()
         self.reset_state()
 
     @staticmethod
@@ -131,10 +135,32 @@ class Rewind(object):
             self._check_timeline_and_lsn(leader)
         return leader and leader.conn_url and self._state == REWIND_STATUS.NEED
 
-    def check_for_checkpoint_after_promote(self):
-        if self._state == REWIND_STATUS.INITIAL and self._postgresql.is_leader() and \
-                self._postgresql.get_master_timeline() == self._postgresql.pg_control_timeline():
-            self._state = REWIND_STATUS.CHECKPOINT
+    def __checkpoint(self, task):
+        try:
+            result = self._postgresql.checkpoint()
+        except Exception as e:
+            result = 'Exception: ' + str(e)
+        with task:
+            task.complete(not bool(result))
+
+    def ensure_checkpoint_after_promote(self):
+        """After promote issue a CHECKPOINT from a new thread and asynchronously check the result.
+        In case if CHECKPOINT failed, just check that timeline in pg_control was updated."""
+
+        if self._state == REWIND_STATUS.INITIAL and self._postgresql.is_leader():
+            with self._checkpoint_task_lock:
+                if self._checkpoint_task:
+                    with self._checkpoint_task:
+                        if self._checkpoint_task.result:
+                            self._state = REWIND_STATUS.CHECKPOINT
+                        if self._checkpoint_task.result is not False:
+                            return
+                else:
+                    self._checkpoint_task = CriticalTask()
+                    return Thread(target=self.__checkpoint, args=(self._checkpoint_task,)).start()
+
+            if self._postgresql.get_master_timeline() == self._postgresql.pg_control_timeline():
+                self._state = REWIND_STATUS.CHECKPOINT
 
     def checkpoint_after_promote(self):
         return self._state == REWIND_STATUS.CHECKPOINT
@@ -191,6 +217,8 @@ class Rewind(object):
 
     def reset_state(self):
         self._state = REWIND_STATUS.INITIAL
+        with self._checkpoint_task_lock:
+            self._checkpoint_task = None
 
     @property
     def is_needed(self):

tests/test_ha.py

@@ -172,6 +172,7 @@ def run_async(self, func, args=()):
 @patch('patroni.postgresql.polling_loop', Mock(return_value=range(1)))
 @patch('patroni.async_executor.AsyncExecutor.busy', PropertyMock(return_value=False))
 @patch('patroni.async_executor.AsyncExecutor.run_async', run_async)
+@patch('patroni.postgresql.rewind.Thread', Mock())
 @patch('subprocess.call', Mock(return_value=0))
 @patch('time.sleep', Mock())
 class TestHa(PostgresInit):

tests/test_rewind.py

@@ -7,6 +7,16 @@ from patroni.postgresql.rewind import Rewind
 from . import BaseTestPostgresql, MockCursor, psycopg2_connect
 
 
+class MockThread(object):
+
+    def __init__(self, target, args):
+        self._target = target
+        self._args = args
+
+    def start(self):
+        self._target(*self._args)
+
+
 @patch('subprocess.call', Mock(return_value=0))
 @patch('psycopg2.connect', psycopg2_connect)
 class TestRewind(BaseTestPostgresql):
@@ -102,6 +112,21 @@ class TestRewind(BaseTestPostgresql):
         self.r.check_leader_is_not_in_recovery()
         self.r.check_leader_is_not_in_recovery()
 
-    @patch.object(Postgresql, 'controldata', Mock(return_value={"Latest checkpoint's TimeLineID": 1}))
-    def test_check_for_checkpoint_after_promote(self):
-        self.r.check_for_checkpoint_after_promote()
+    @patch('patroni.postgresql.rewind.Thread', MockThread)
+    @patch.object(Postgresql, 'controldata')
+    @patch.object(Postgresql, 'checkpoint')
+    def test_ensure_checkpoint_after_promote(self, mock_checkpoint, mock_controldata):
+        mock_checkpoint.return_value = None
+        self.r.ensure_checkpoint_after_promote()
+        self.r.ensure_checkpoint_after_promote()
+
+        self.r.reset_state()
+        mock_controldata.return_value = {"Latest checkpoint's TimeLineID": 1}
+        mock_checkpoint.side_effect = Exception
+        self.r.ensure_checkpoint_after_promote()
+        self.r.ensure_checkpoint_after_promote()
+
+        self.r.reset_state()
+        mock_controldata.side_effect = TypeError
+        self.r.ensure_checkpoint_after_promote()
+        self.r.ensure_checkpoint_after_promote()