Update timeline on standby cluster (#1332)

Fixes https://github.com/zalando/patroni/issues/1031

patroni/ha.py

@@ -191,10 +191,20 @@ class Ha(object):
             if self._async_executor.scheduled_action in (None, 'promote') \
                     and data['state'] in ['running', 'restarting', 'starting']:
                 try:
-                    timeline, wal_position = self.state_handler.timeline_wal_position()
+                    timeline, wal_position, pg_control_timeline = self.state_handler.timeline_wal_position()
                     data['xlog_location'] = wal_position
                     if not timeline:
-                        timeline = self.state_handler.replica_cached_timeline(self._leader_timeline)
+                        # So far the only way to get the current timeline on the standby is from
+                        # the replication connection. In order to avoid opening the replication
+                        # connection on every iteration of HA loop we will do it only when noticed
+                        # that the timeline on the primary has changed.
+                        # Unfortunately such optimization isn't possible on the standby_leader,
+                        # therefore we will get the timeline from pg_control, either by calling
+                        # pg_control_checkpoint() on 9.6+ or by parsing the output of pg_controldata.
+                        if self.state_handler.role == 'standby_leader':
+                            timeline = pg_control_timeline or self.state_handler.pg_control_timeline()
+                        else:
+                            timeline = self.state_handler.replica_cached_timeline(timeline)
                     if timeline:
                         data['timeline'] = timeline
                 except Exception:
@@ -597,7 +607,7 @@ class Ha(object):
         """This method tries to determine whether I am healthy enough to became a new leader candidate or not."""
 
         # We don't call `last_operation()` here because it returns a string
-        _, my_wal_position = self.state_handler.timeline_wal_position()
+        _, my_wal_position, _ = self.state_handler.timeline_wal_position()
         if check_replication_lag and self.is_lagging(my_wal_position):
             logger.info('My wal position exceeds maximum replication lag')
             return False  # Too far behind last reported wal position on master

patroni/postgresql/__init__.py

@@ -134,6 +134,8 @@ class Postgresql(object):
 
     @property
     def cluster_info_query(self):
+        pg_control_timeline = 'timeline_id FROM pg_catalog.pg_control_checkpoint()' \
+                if self._major_version >= 90600 and self.role == 'standby_leader' else '0'
         return ("SELECT CASE WHEN pg_catalog.pg_is_in_recovery() THEN 0 "
                 "ELSE ('x' || pg_catalog.substr(pg_catalog.pg_{0}file_name("
                 "pg_catalog.pg_current_{0}_{1}()), 1, 8))::bit(32)::int END, "
@@ -142,7 +144,7 @@ class Postgresql(object):
                 "pg_catalog.pg_last_{0}_receive_{1}(), '0/0'), '0/0')::bigint,"
                 " pg_catalog.pg_{0}_{1}_diff(pg_catalog.pg_last_{0}_replay_{1}(), '0/0')::bigint)"
                 "ELSE pg_catalog.pg_{0}_{1}_diff(pg_catalog.pg_current_{0}_{1}(), '0/0')::bigint "
-                "END").format(self.wal_name, self.lsn_name)
+                "END, {2}").format(self.wal_name, self.lsn_name, pg_control_timeline)
 
     def _version_file_exists(self):
         return not self.data_directory_empty() and os.path.isfile(self._version_file)
@@ -289,7 +291,7 @@ class Postgresql(object):
         if not self._cluster_info_state:
             try:
                 result = self._is_leader_retry(self._query, self.cluster_info_query).fetchone()
-                self._cluster_info_state = dict(zip(['timeline', 'wal_position'], result))
+                self._cluster_info_state = dict(zip(['timeline', 'wal_position', 'pg_control_timeline'], result))
             except RetryFailedError as e:  # SELECT failed two times
                 self._cluster_info_state = {'error': str(e)}
                 if not self.is_starting() and self.pg_isready() == STATE_REJECT:
@@ -303,6 +305,12 @@ class Postgresql(object):
     def is_leader(self):
         return bool(self._cluster_info_state_get('timeline'))
 
+    def pg_control_timeline(self):
+        try:
+            return int(self.controldata().get("Latest checkpoint's TimeLineID"))
+        except (TypeError, ValueError):
+            logger.exception('Failed to parse timeline from pg_controldata output')
+
     def is_running(self):
         """Returns PostmasterProcess if one is running on the data directory or None. If most recently seen process
         is running updates the cached process based on pid file."""
@@ -735,11 +743,13 @@ class Postgresql(object):
         # This method could be called from different threads (simultaneously with some other `_query` calls).
         # If it is called not from main thread we will create a new cursor to execute statement.
         if current_thread().ident == self.__thread_ident:
-            return self._cluster_info_state_get('timeline'), self._cluster_info_state_get('wal_position')
+            return (self._cluster_info_state_get('timeline'),
+                    self._cluster_info_state_get('wal_position'),
+                    self._cluster_info_state_get('pg_control_timeline'))
 
         with self.connection().cursor() as cursor:
             cursor.execute(self.cluster_info_query)
-            return cursor.fetchone()[:2]
+            return cursor.fetchone()[:3]
 
     def postmaster_start_time(self):
         try:

patroni/postgresql/rewind.py

@@ -132,13 +132,9 @@ class Rewind(object):
         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():
-            try:
-                timeline = int(self._postgresql.controldata().get("Latest checkpoint's TimeLineID"))
-                if self._postgresql.get_master_timeline() == timeline:
-                    self._state = REWIND_STATUS.CHECKPOINT
-            except (TypeError, ValueError):
-                logger.exception('Failed to parse timeline from pg_controldata output')
+        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_after_promote(self):
         return self._state == REWIND_STATUS.CHECKPOINT

tests/__init__.py

@@ -86,7 +86,7 @@ class MockCursor(object):
         elif sql.startswith('SELECT slot_name'):
             self.results = [('blabla', 'physical'), ('foobar', 'physical'), ('ls', 'logical', 'a', 'b')]
         elif sql.startswith('SELECT CASE WHEN pg_catalog.pg_is_in_recovery()'):
-            self.results = [(1, 2)]
+            self.results = [(1, 2, 1)]
         elif sql.startswith('SELECT pg_catalog.pg_is_in_recovery()'):
             self.results = [(False, 2)]
         elif sql.startswith('SELECT pg_catalog.to_char'):

tests/test_ha.py

@@ -150,7 +150,7 @@ def run_async(self, func, args=()):
 
 @patch.object(Postgresql, 'is_running', Mock(return_value=MockPostmaster()))
 @patch.object(Postgresql, 'is_leader', Mock(return_value=True))
-@patch.object(Postgresql, 'timeline_wal_position', Mock(return_value=(1, 10)))
+@patch.object(Postgresql, 'timeline_wal_position', Mock(return_value=(1, 10, 1)))
 @patch.object(Postgresql, '_cluster_info_state_get', Mock(return_value=3))
 @patch.object(Postgresql, 'call_nowait', Mock(return_value=True))
 @patch.object(Postgresql, 'data_directory_empty', Mock(return_value=False))
@@ -199,9 +199,12 @@ class TestHa(PostgresInit):
         self.assertTrue(self.ha.update_lock(True))
 
     def test_touch_member(self):
-        self.p.timeline_wal_position = Mock(return_value=(0, 1))
+        self.p.timeline_wal_position = Mock(return_value=(0, 1, 0))
         self.p.replica_cached_timeline = Mock(side_effect=Exception)
         self.ha.touch_member()
+        self.p.timeline_wal_position = Mock(return_value=(0, 1, 1))
+        self.p.set_role('standby_leader')
+        self.ha.touch_member()
 
     def test_is_leader(self):
         self.assertFalse(self.ha.is_leader())
@@ -601,7 +604,7 @@ class TestHa(PostgresInit):
         # in synchronous_mode consider itself healthy if the former leader is accessible in read-only and ahead of us
         with patch.object(Ha, 'is_synchronous_mode', Mock(return_value=True)):
             self.assertTrue(self.ha._is_healthiest_node(self.ha.old_cluster.members))
-        with patch('patroni.postgresql.Postgresql.timeline_wal_position', return_value=(1, 1)):
+        with patch('patroni.postgresql.Postgresql.timeline_wal_position', return_value=(1, 1, 1)):
             self.assertFalse(self.ha._is_healthiest_node(self.ha.old_cluster.members))
         with patch('patroni.postgresql.Postgresql.replica_cached_timeline', return_value=1):
             self.assertFalse(self.ha._is_healthiest_node(self.ha.old_cluster.members))

tests/test_postgresql.py

@@ -331,7 +331,7 @@ class TestPostgresql(BaseTestPostgresql):
         self.assertTrue(self.p.promote(0))
 
     def test_timeline_wal_position(self):
-        self.assertEqual(self.p.timeline_wal_position(), (1, 2))
+        self.assertEqual(self.p.timeline_wal_position(), (1, 2, 1))
         Thread(target=self.p.timeline_wal_position).start()
 
     @patch.object(PostmasterProcess, 'from_pidfile')