Try to fetch missing WAL if pg_rewind complains about it (#1561)

It could happen that the WAL segment required for `pg_rewind` doesn't exist in the `pg_wal` anymore and therefore `pg_rewind` can't find the checkpoint location before the diverging point.
Starting from PostgreSQL 13 `pg_rewind` could use `restore_command` for fetching missing WALs, but we can do better than that.
On older PostgreSQL versions Patroni will parse the stdout and stderr of failed rewind attempt, try to fetch the missing WAL by calling the `restore_command`, and repeat an attempt.

patroni/postgresql/__init__.py

@@ -137,6 +137,10 @@ class Postgresql(object):
     def callback(self):
         return self.config.get('callbacks') or {}
 
+    @property
+    def wal_dir(self):
+        return os.path.join(self._data_dir, 'pg_' + self.wal_name)
+
     @property
     def wal_name(self):
         return 'wal' if self._major_version >= 100000 else 'xlog'
@@ -743,7 +747,7 @@ class Postgresql(object):
         return self._cluster_info_state_get('timeline')
 
     def get_history(self, timeline):
-        history_path = os.path.join(self._data_dir, 'pg_' + self.wal_name, '{0:08X}.history'.format(timeline))
+        history_path = os.path.join(self.wal_dir, '{0:08X}.history'.format(timeline))
         history_mtime = mtime(history_path)
         if history_mtime:
             try:

patroni/postgresql/cancellable.py

@@ -1,11 +1,9 @@
 import logging
-import os
 import psutil
 import subprocess
 
 from patroni.exceptions import PostgresException
 from patroni.utils import polling_loop
-from six import string_types
 from threading import Lock
 
 logger = logging.getLogger(__name__)
@@ -75,16 +73,16 @@ class CancellableSubprocess(CancellableExecutor):
         for s in ('stdin', 'stdout', 'stderr'):
             kwargs.pop(s, None)
 
-        communicate_input = 'communicate_input' in kwargs
-        if communicate_input:
-            input_data = kwargs.pop('communicate_input', None)
-            if not isinstance(input_data, string_types):
-                input_data = ''
-            if input_data and input_data[-1] != '\n':
-                input_data += '\n'
+        communicate = kwargs.pop('communicate', None)
+        if isinstance(communicate, dict):
+            input_data = communicate.get('input')
+            if input_data:
+                if input_data[-1] != '\n':
+                    input_data += '\n'
+                input_data = input_data.encode('utf-8')
             kwargs['stdin'] = subprocess.PIPE
-            kwargs['stdout'] = open(os.devnull, 'w')
-            kwargs['stderr'] = subprocess.STDOUT
+            kwargs['stdout'] = subprocess.PIPE
+            kwargs['stderr'] = subprocess.PIPE
 
         try:
             with self._lock:
@@ -95,10 +93,8 @@ class CancellableSubprocess(CancellableExecutor):
                 started = self._start_process(*args, **kwargs)
 
             if started:
-                if communicate_input:
-                    if input_data:
-                        self._process.communicate(input_data)
-                    self._process.stdin.close()
+                if isinstance(communicate, dict):
+                    communicate['stdout'], communicate['stderr'] = self._process.communicate(input_data)
                 return self._process.wait()
         finally:
             with self._lock:

patroni/postgresql/rewind.py

@@ -250,21 +250,82 @@ class Rewind(object):
     def checkpoint_after_promote(self):
         return self._state == REWIND_STATUS.CHECKPOINT
 
+    def _fetch_missing_wal(self, restore_command, wal_filename):
+        cmd = ''
+        length = len(restore_command)
+        i = 0
+        while i < length:
+            if restore_command[i] == '%' and i + 1 < length:
+                i += 1
+                if restore_command[i] == 'p':
+                    cmd += os.path.join(self._postgresql.wal_dir, wal_filename)
+                elif restore_command[i] == 'f':
+                    cmd += wal_filename
+                elif restore_command[i] == 'r':
+                    cmd += '000000010000000000000001'
+                elif restore_command[i] == '%':
+                    cmd += '%'
+                else:
+                    cmd += '%'
+                    i -= 1
+            else:
+                cmd += restore_command[i]
+            i += 1
+
+        logger.info('Trying to fetch the missing wal: %s', cmd)
+        return self._postgresql.cancellable.call(cmd, shell=True) == 0
+
+    def _find_missing_wal(self, data):
+        # could not open file "$PGDATA/pg_wal/0000000A00006AA100000068": No such file or directory
+        pattern = 'could not open file "'
+        for line in data.decode('utf-8').split('\n'):
+            b = line.find(pattern)
+            if b > -1:
+                b += len(pattern)
+                e = line.find('": ', b)
+                if e > -1:
+                    waldir, wal_filename = os.path.split(line[b:e])
+                    if waldir.endswith(os.path.sep + 'pg_' + self._postgresql.wal_name) and len(wal_filename) == 24:
+                        return wal_filename
+
     def pg_rewind(self, r):
         # prepare pg_rewind connection
         env = self._postgresql.config.write_pgpass(r)
-        env['PGOPTIONS'] = '-c statement_timeout=0'
+        env.update(LANG='C', LC_ALL='C', PGOPTIONS='-c statement_timeout=0')
         dsn = self._postgresql.config.format_dsn(r, True)
         logger.info('running pg_rewind from %s', dsn)
 
+        restore_command = self._postgresql.config.get('recovery_conf', {}).get('restore_command') \
+            if self._postgresql.major_version < 120000 else self._postgresql.get_guc_value('restore_command')
+
         cmd = [self._postgresql.pgcommand('pg_rewind')]
-        if self._postgresql.major_version >= 130000 and self._postgresql.get_guc_value('restore_command'):
+        if self._postgresql.major_version >= 130000 and restore_command:
             cmd.append('--restore-target-wal')
         cmd.extend(['-D', self._postgresql.data_dir, '--source-server', dsn])
-        try:
-            return self._postgresql.cancellable.call(cmd, env=env) == 0
-        except OSError:
-            return False
+
+        while True:
+            results = {}
+            ret = self._postgresql.cancellable.call(cmd, env=env, communicate=results)
+
+            logger.info('pg_rewind exit code=%s', ret)
+            if ret is None:
+                return False
+
+            logger.info(' stdout=%s', results['stdout'].decode('utf-8'))
+            logger.info(' stderr=%s', results['stderr'].decode('utf-8'))
+            if ret == 0:
+                return True
+
+            if not restore_command or self._postgresql.major_version >= 130000:
+                return False
+
+            missing_wal = self._find_missing_wal(results['stderr']) or self._find_missing_wal(results['stdout'])
+            if not missing_wal:
+                return False
+
+            if not self._fetch_missing_wal(restore_command, missing_wal):
+                logger.info('Failed to fetch WAL segment %s required for pg_rewind', missing_wal)
+                return False
 
     def execute(self, leader):
         if self._postgresql.is_running() and not self._postgresql.stop(checkpoint=False):
@@ -335,17 +396,17 @@ class Rewind(object):
             logger.exception('Error when reading postmaster.opts')
         return result
 
-    def single_user_mode(self, command=None, options=None):
+    def single_user_mode(self, communicate=None, options=None):
         """run a given command in a single-user mode. If the command is empty - then just start and stop"""
         cmd = [self._postgresql.pgcommand('postgres'), '--single', '-D', self._postgresql.data_dir]
         for opt, val in sorted((options or {}).items()):
             cmd.extend(['-c', '{0}={1}'.format(opt, val)])
         # need a database name to connect
         cmd.append('template1')
-        return self._postgresql.cancellable.call(cmd, communicate_input=command)
+        return self._postgresql.cancellable.call(cmd, communicate=communicate)
 
     def cleanup_archive_status(self):
-        status_dir = os.path.join(self._postgresql.data_dir, 'pg_' + self._postgresql.wal_name, 'archive_status')
+        status_dir = os.path.join(self._postgresql.wal_dir, 'archive_status')
         try:
             for f in os.listdir(status_dir):
                 path = os.path.join(status_dir, f)

tests/test_cancellable.py

@@ -13,7 +13,7 @@ class TestCancellableSubprocess(unittest.TestCase):
 
     def test_call(self):
         self.c.cancel()
-        self.assertRaises(PostgresException, self.c.call, communicate_input=None)
+        self.assertRaises(PostgresException, self.c.call)
 
     def test__kill_children(self):
         self.c._process_children = [Mock()]

tests/test_rewind.py

@@ -18,6 +18,28 @@ class MockThread(object):
         self._target(*self._args)
 
 
+def mock_cancellable_call(*args, **kwargs):
+    communicate = kwargs.pop('communicate', None)
+    if isinstance(communicate, dict):
+        communicate.update(stdout=b'', stderr=b'pg_rewind: error: could not open file ' +
+                                              b'"data/postgresql0/pg_xlog/000000010000000000000003": No such file')
+    return 1
+
+
+def mock_cancellable_call0(*args, **kwargs):
+    communicate = kwargs.pop('communicate', None)
+    if isinstance(communicate, dict):
+        communicate.update(stdout=b'', stderr=b'')
+    return 0
+
+
+def mock_cancellable_call1(*args, **kwargs):
+    communicate = kwargs.pop('communicate', None)
+    if isinstance(communicate, dict):
+        communicate.update(stdout=b'', stderr=b'')
+    return 1
+
+
 @patch('subprocess.call', Mock(return_value=0))
 @patch('psycopg2.connect', psycopg2_connect)
 class TestRewind(BaseTestPostgresql):
@@ -36,16 +58,23 @@ class TestRewind(BaseTestPostgresql):
         self.p.config._config['use_pg_rewind'] = False
         self.assertFalse(self.r.can_rewind)
 
-    @patch.object(Postgresql, 'major_version', PropertyMock(return_value=130000))
-    @patch.object(CancellableSubprocess, 'call')
-    def test_pg_rewind(self, mock_cancellable_subprocess_call):
+    def test_pg_rewind(self):
         r = {'user': '', 'host': '', 'port': '', 'database': '', 'password': ''}
-        mock_cancellable_subprocess_call.return_value = 0
-        with patch('subprocess.check_output', Mock(return_value=b'foo')):
-            self.assertTrue(self.r.pg_rewind(r))
-        mock_cancellable_subprocess_call.side_effect = OSError
-        with patch('subprocess.check_output', Mock(side_effect=Exception)):
-            self.assertFalse(self.r.pg_rewind(r))
+        with patch.object(Postgresql, 'major_version', PropertyMock(return_value=130000)),\
+                patch.object(CancellableSubprocess, 'call', Mock(return_value=None)):
+            with patch('subprocess.check_output', Mock(return_value=b'boo')):
+                self.assertFalse(self.r.pg_rewind(r))
+            with patch('subprocess.check_output', Mock(side_effect=Exception)):
+                self.assertFalse(self.r.pg_rewind(r))
+
+        with patch.object(Postgresql, 'major_version', PropertyMock(return_value=120000)),\
+                patch('subprocess.check_output', Mock(return_value=b'foo %f %p %r %% % %')):
+            with patch.object(CancellableSubprocess, 'call', mock_cancellable_call):
+                self.assertFalse(self.r.pg_rewind(r))
+            with patch.object(CancellableSubprocess, 'call', mock_cancellable_call0):
+                self.assertTrue(self.r.pg_rewind(r))
+            with patch.object(CancellableSubprocess, 'call', mock_cancellable_call1):
+                self.assertFalse(self.r.pg_rewind(r))
 
     @patch.object(Rewind, 'can_rewind', PropertyMock(return_value=True))
     def test__get_local_timeline_lsn(self):
@@ -61,7 +90,7 @@ class TestRewind(BaseTestPostgresql):
             with patch.object(MockCursor, 'fetchone', Mock(side_effect=[(0, 0, 1, 1,), Exception])):
                 self.r.rewind_or_reinitialize_needed_and_possible(self.leader)
 
-    @patch.object(CancellableSubprocess, 'call', Mock(return_value=0))
+    @patch.object(CancellableSubprocess, 'call', mock_cancellable_call)
     @patch.object(Postgresql, 'checkpoint', side_effect=['', '1'],)
     @patch.object(Postgresql, 'stop', Mock(return_value=False))
     @patch.object(Postgresql, 'start', Mock())
@@ -161,7 +190,8 @@ class TestRewind(BaseTestPostgresql):
     @patch('psutil.Popen')
     def test_single_user_mode(self, subprocess_popen_mock):
         subprocess_popen_mock.return_value.wait.return_value = 0
-        self.assertEqual(self.r.single_user_mode('CHECKPOINT', {'archive_mode': 'on'}), 0)
+        subprocess_popen_mock.return_value.communicate.return_value = ('', '')
+        self.assertEqual(self.r.single_user_mode({'input': 'CHECKPOINT'}, {'archive_mode': 'on'}), 0)
 
     @patch('os.listdir', Mock(side_effect=[OSError, ['a', 'b']]))
     @patch('os.unlink', Mock(side_effect=OSError))