Take advantage of written_lsn and latest_end_lsn from pg_stat_wal_receiver (#3268)

The first one if available starting from PostgreSQL v13 and contains the
real write LSN. We will prefer it over value returned by
pg_last_wal_receive_lsn(), which is in fact flush LSN.

The second one is available starting from PostgreSQL v9.6 and  points to
WAL flush on the source host. In case of primary it will allow to better
calculate the replay lag, because values stored in DCS are updated only
every loop_wait seconds.

patroni/api.py

@@ -310,12 +310,13 @@ class RestApiHandler(BaseHTTPRequestHandler):
         """
         path = '/primary' if self.path == '/' else self.path
         response = self.get_postgresql_status()
+        latest_end_lsn = response.pop('latest_end_lsn', 0)
 
         patroni = self.server.patroni
         cluster = patroni.dcs.cluster
         config = global_config.from_cluster(cluster)
 
-        leader_optime = cluster and cluster.status.last_lsn
+        leader_optime = max(cluster and cluster.status.last_lsn or 0, latest_end_lsn)
         replayed_location = response.get('xlog', {}).get('replayed_location', 0)
         max_replica_lag = parse_int(self.path_query.get('lag', [sys.maxsize])[0], 'B')
         if max_replica_lag is None:
@@ -474,6 +475,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
         Postgres.
         """
         response = self.get_postgresql_status(True)
+        response.pop('latest_end_lsn', None)
         self._write_status_response(200, response)
 
     def do_GET_cluster(self) -> None:
@@ -1268,6 +1270,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
             * ``postmaster_start_time``: ``pg_postmaster_start_time()``;
             * ``role``: ``replica`` or ``primary`` based on ``pg_is_in_recovery()`` output;
             * ``server_version``: Postgres version without periods, e.g. ``150002`` for Postgres ``15.2``;
+            * ``latest_end_lsn``: latest_end_lsn value from ``pg_stat_get_wal_receiver()``, only on replica nodes;
             * ``xlog``: dictionary. Its structure depends on ``role``:
 
                 * If ``primary``:
@@ -1307,15 +1310,18 @@ class RestApiHandler(BaseHTTPRequestHandler):
 
             if postgresql.state not in ('running', 'restarting', 'starting'):
                 raise RetryFailedError('')
-            replication_state = ('(pg_catalog.pg_stat_get_wal_receiver()).status'
+            replication_state = ("pg_catalog.pg_{0}_{1}_diff(wr.latest_end_lsn, '0/0')::bigint, wr.status"
-                                 if postgresql.major_version >= 90600 else 'NULL') + ", " +\
+                                 if postgresql.major_version >= 90600 else "NULL, NULL") + ", " +\
-                ("pg_catalog.current_setting('restore_command')" if postgresql.major_version >= 120000 else "NULL")
+                ("pg_catalog.current_setting('restore_command')" if postgresql.major_version >= 120000 else "NULL") +\
+                ", " + ("pg_catalog.pg_wal_lsn_diff(wr.written_lsn, '0/0')::bigint"
+                        if postgresql.major_version >= 130000 else "NULL")
             stmt = ("SELECT " + postgresql.POSTMASTER_START_TIME + ", " + postgresql.TL_LSN + ","
                     " pg_catalog.pg_last_xact_replay_timestamp(), " + replication_state + ","
-                    " pg_catalog.array_to_json(pg_catalog.array_agg(pg_catalog.row_to_json(ri))) "
+                    " (SELECT pg_catalog.array_to_json(pg_catalog.array_agg(pg_catalog.row_to_json(ri))) "
                     "FROM (SELECT (SELECT rolname FROM pg_catalog.pg_authid WHERE oid = usesysid) AS usename,"
                     " application_name, client_addr, w.state, sync_state, sync_priority"
-                    " FROM pg_catalog.pg_stat_get_wal_senders() w, pg_catalog.pg_stat_get_activity(pid)) AS ri")
+                    " FROM pg_catalog.pg_stat_get_wal_senders() w, pg_catalog.pg_stat_get_activity(pid)) AS ri)") +\
+                (" FROM pg_catalog.pg_stat_get_wal_receiver() AS wr" if postgresql.major_version >= 90600 else "")
 
             row = self.query(stmt.format(postgresql.wal_name, postgresql.lsn_name,
                                          postgresql.wal_flush), retry=retry)[0]
@@ -1325,7 +1331,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
                 'role': 'replica' if row[1] == 0 else 'primary',
                 'server_version': postgresql.server_version,
                 'xlog': ({
-                    'received_location': row[4] or row[3],
+                    'received_location': row[10] or row[4] or row[3],
                     'replayed_location': row[3],
                     'replayed_timestamp': row[6],
                     'paused': row[5]} if row[1] == 0 else {
@@ -1347,12 +1353,15 @@ class RestApiHandler(BaseHTTPRequestHandler):
                     if not cluster or cluster.is_unlocked() or not cluster.leader else cluster.leader.timeline
                 result['timeline'] = postgresql.replica_cached_timeline(leader_timeline)
 
-            replication_state = postgresql.replication_state_from_parameters(row[1] > 0, row[7], row[8])
+            if row[7]:
+                result['latest_end_lsn'] = row[7]
+
+            replication_state = postgresql.replication_state_from_parameters(row[1] > 0, row[8], row[9])
             if replication_state:
                 result['replication_state'] = replication_state
 
-            if row[9]:
+            if row[11]:
-                result['replication'] = row[9]
+                result['replication'] = row[11]
 
         except (psycopg.Error, RetryFailedError, PostgresConnectionException):
             state = postgresql.state

patroni/postgresql/__init__.py

@@ -235,14 +235,17 @@ class Postgresql(object):
                         " AS confirmed_flush_lsn, pg_catalog.pg_wal_lsn_diff(restart_lsn, '0/0')::bigint"
                         f" AS restart_lsn, xmin FROM pg_catalog.pg_get_replication_slots(){filter_failover}) AS s)"
                         if self._should_query_slots and self.can_advance_slots else "NULL") + extra
-            extra = (", CASE WHEN latest_end_lsn IS NULL THEN NULL ELSE received_tli END,"
+
-                     " slot_name, conninfo, status, {0} FROM pg_catalog.pg_stat_get_wal_receiver()").format(extra)
+            written_lsn = ("pg_catalog.pg_wal_lsn_diff(written_lsn, '0/0')::bigint"
+                           if self._major_version >= 130000 else "NULL")
+            extra = (", CASE WHEN latest_end_lsn IS NULL THEN NULL ELSE received_tli END, {0}, slot_name, "
+                     "conninfo, status, {1} FROM pg_catalog.pg_stat_get_wal_receiver()").format(written_lsn, extra)
             if self.role == 'standby_leader':
                 extra = "timeline_id" + extra + ", pg_catalog.pg_control_checkpoint()"
             else:
                 extra = "0" + extra
         else:
-            extra = "0, NULL, NULL, NULL, NULL, NULL, NULL" + extra
+            extra = "0, NULL, NULL, NULL, NULL, NULL, NULL, NULL" + extra
 
         return ("SELECT " + self.TL_LSN + ", {3}").format(self.wal_name, self.lsn_name, self.wal_flush, extra)
 
@@ -477,8 +480,8 @@ class Postgresql(object):
                 result = self._is_leader_retry(self._query, self.cluster_info_query)[0]
                 cluster_info_state = dict(zip(['timeline', 'wal_position', 'replayed_location',
                                                'received_location', 'replay_paused', 'pg_control_timeline',
-                                               'received_tli', 'slot_name', 'conninfo', 'receiver_state',
+                                               'received_tli', 'write_location', 'slot_name', 'conninfo',
-                                               'restore_command', 'slots', 'synchronous_commit',
+                                               'receiver_state', 'restore_command', 'slots', 'synchronous_commit',
                                                'synchronous_standby_names', 'pg_stat_replication'], result))
                 if self._should_query_slots and self.can_advance_slots:
                     cluster_info_state['slots'] =\
@@ -498,7 +501,9 @@ class Postgresql(object):
         return self._cluster_info_state_get('replayed_location')
 
     def received_location(self) -> Optional[int]:
-        return self._cluster_info_state_get('received_location')
+        write = self._cluster_info_state_get('write_location')
+        received = self._cluster_info_state_get('received_location')
+        return max(received, write) if received and write else write or received
 
     def slots(self) -> Dict[str, int]:
         """Get replication slots state.
@@ -1251,8 +1256,9 @@ class Postgresql(object):
             received_location = self.received_location()
             pg_control_timeline = self._cluster_info_state_get('pg_control_timeline')
         else:
-            timeline, wal_position, replayed_location, received_location, _, pg_control_timeline = \
+            timeline, wal_position, replayed_location, received_location, _, pg_control_timeline, _, write_location = \
-                self._query(self.cluster_info_query)[0][:6]
+                self._query(self.cluster_info_query)[0][:8]
+            received_location = max(received_location or 0, write_location or 0)
 
         wal_position = self._wal_position(bool(timeline), wal_position, received_location, replayed_location)
         return timeline, wal_position, pg_control_timeline

tests/__init__.py

@@ -144,7 +144,7 @@ class MockCursor(object):
         elif sql.startswith('WITH slots AS (SELECT slot_name, active'):
             self.results = [(False, True)] if self.rowcount == 1 else []
         elif sql.startswith('SELECT CASE WHEN pg_catalog.pg_is_in_recovery()'):
-            self.results = [(1, 2, 1, 0, False, 1, 1, None, None, 'streaming', '',
+            self.results = [(1, 2, 1, 0, False, 1, 1, 1, None, None, 'streaming', '',
                              [{"slot_name": "ls", "confirmed_flush_lsn": 12345, "restart_lsn": 12344}],
                              'on', 'n1', None)]
         elif sql.startswith('SELECT pg_catalog.pg_is_in_recovery()'):

tests/test_api.py

@@ -32,7 +32,7 @@ class MockConnection:
 
     @staticmethod
     def query(sql, *params):
-        return [(postmaster_start_time, 0, '', 0, '', False, postmaster_start_time, 'streaming', None,
+        return [(postmaster_start_time, 0, '', 0, '', False, postmaster_start_time, 1, 'streaming', None, 0,
                  '[{"application_name":"walreceiver","client_addr":"1.2.3.4",'
                  + '"state":"streaming","sync_state":"async","sync_priority":0}]')]
 
@@ -239,7 +239,8 @@ class TestRestApiHandler(unittest.TestCase):
         with patch.object(MockHa, 'restart_scheduled', Mock(return_value=True)):
             MockRestApiServer(RestApiHandler, 'GET /primary')
         self.assertIsNotNone(MockRestApiServer(RestApiHandler, 'GET /primary'))
-        with patch.object(RestApiServer, 'query', Mock(return_value=[('', 1, '', '', '', '', False, None, None, '')])):
+        with patch.object(RestApiServer, 'query',
+                          Mock(return_value=[('', 1, '', '', '', '', False, None, 0, None, 0, '')])):
             self.assertIsNotNone(MockRestApiServer(RestApiHandler, 'GET /patroni'))
         with patch.object(global_config.__class__, 'is_standby_cluster', Mock(return_value=True)), \
                 patch.object(global_config.__class__, 'is_paused', Mock(return_value=True)):

tests/test_slots.py

@@ -109,7 +109,7 @@ class TestSlotsHandler(BaseTestPostgresql):
         with patch.object(Postgresql, '_query') as mock_query:
             self.p.reset_cluster_info_state(None)
             mock_query.return_value = [(
-                1, 0, 0, 0, 0, 0, 0, 0, 0, None, None,
+                1, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None,
                 [{"slot_name": "ls", "type": "logical", "datoid": 5, "plugin": "b", "xmin": 105,
                   "confirmed_flush_lsn": 12345, "catalog_xmin": 105, "restart_lsn": 12344},
                  {"slot_name": "blabla", "type": "physical", "datoid": None, "plugin": None, "xmin": 105,
@@ -118,7 +118,7 @@ class TestSlotsHandler(BaseTestPostgresql):
 
             self.p.reset_cluster_info_state(None)
             mock_query.return_value = [(
-                1, 0, 0, 0, 0, 0, 0, 0, 0, None, None,
+                1, 0, 0, 0, 0, 0, 0, 0, 0, 0, None, None,
                 [{"slot_name": "ls", "type": "logical", "datoid": 6, "plugin": "b", "xmin": 105,
                   "confirmed_flush_lsn": 12345, "catalog_xmin": 105}])]
             self.assertEqual(self.p.slots(), {'postgresql0': 0})