Repository: zalando/patroni
Pull request #2384 (2022-08-15 10:00:27)

Archive possibly missing WALs before rewind
There is currently a risk to lose some WAL segments entirely in case
archive_mode was set to 'on' before a promote and there are some WALs
with .ready files on the former leader we are trying to rewind. It happens
because of the pg_rewind's modus operandi: it simply syncs the content of
pg_wal directory of the old leader with the new leader's one. Including
deletion of all WALs that are not present on the current leader, regardless
their archive status on the former one. Thus, in case the new leader has
already recycled such files, we just remove them entirely.

In case archive_mode was set to 'always' and the .ready WALs are actually
present in archive, it is for end user who writes the archive_command to
avoid overwriting and to properly test it. 
According to [pg docs](https://www.postgresql.org/docs/current/continuous-archiving.html):   

> The archive command should generally be designed to refuse to  overwrite any pre-existing archive file. This is an important safety feature to preserve the integrity of your archive in case of  administrator error (such as sending the output of two different servers to the same archive directory).    
It is advisable to test your proposed archive command to ensure that it indeed does not overwrite an existing file, and that it returns nonzero status in this case. The example command above for Unix ensures this by including a separate test step.

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

* **38** of **38**   **(100.0%)**  changed or added relevant lines in **1** file are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage increased (+**0.0003%**) to **99.882%**

---



|  Totals | [![Coverage Status](https://coveralls.io/builds/51658078/badge)](https://coveralls.io/builds/51658078) |
| :-- | --: |
| Change from base [Build 2797700302](https://coveralls.io/builds/51422778): |  0.0003% |
| Covered Lines: | 11039 |
| Relevant Lines: | 11052 |

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

CyberDem0n: :+1:
hughcapet: 👍 

Patch:
diff --git a/patroni/postgresql/rewind.py b/patroni/postgresql/rewind.py
index 481f0b7a..6a9b6057 100644
--- a/patroni/postgresql/rewind.py
+++ b/patroni/postgresql/rewind.py
@@ -1,6 +1,8 @@
 import logging
 import os
+import re
 import shlex
+import shutil
 import six
 import subprocess
 
@@ -277,28 +279,37 @@ class Rewind(object):
     def checkpoint_after_promote(self):
         return self._state == REWIND_STATUS.CHECKPOINT
 
-    def _fetch_missing_wal(self, restore_command, wal_filename):
+    def _buid_archiver_command(self, command, wal_filename):
+        """Replace placeholders in the given archiver command's template.
+        Applicable for archive_command and restore_command.
+        Can also be used for archive_cleanup_command and recovery_end_command,
+        however %r value is always set to 000000010000000000000001."""
         cmd = ''
-        length = len(restore_command)
+        length = len(command)
         i = 0
         while i < length:
-            if restore_command[i] == '%' and i + 1 < length:
+            if command[i] == '%' and i + 1 < length:
                 i += 1
-                if restore_command[i] == 'p':
+                if command[i] == 'p':
                     cmd += os.path.join(self._postgresql.wal_dir, wal_filename)
-                elif restore_command[i] == 'f':
+                elif command[i] == 'f':
                     cmd += wal_filename
-                elif restore_command[i] == 'r':
+                elif command[i] == 'r':
                     cmd += '000000010000000000000001'
-                elif restore_command[i] == '%':
+                elif command[i] == '%':
                     cmd += '%'
                 else:
                     cmd += '%'
                     i -= 1
             else:
-                cmd += restore_command[i]
+                cmd += command[i]
             i += 1
 
+        return cmd
+
+    def _fetch_missing_wal(self, restore_command, wal_filename):
+        cmd = self._buid_archiver_command(restore_command, wal_filename)
+
         logger.info('Trying to fetch the missing wal: %s', cmd)
         return self._postgresql.cancellable.call(shlex.split(cmd)) == 0
 
@@ -315,6 +326,41 @@ class Rewind(object):
                     if waldir.endswith('/pg_' + self._postgresql.wal_name) and len(wal_filename) == 24:
                         return wal_filename
 
+    def _archive_ready_wals(self):
+        """Try to archive WALs that have .ready files just in case
+        archive_mode was not set to 'always' before promote, while
+        after it the WALs were recycled on the promoted replica.
+        With this we prevent the entire loss of such WALs and the
+        consequent old leader's start failure."""
+        archive_mode = self._postgresql.get_guc_value('archive_mode')
+        archive_cmd = self._postgresql.get_guc_value('archive_command')
+        if archive_mode not in ('on', 'always') or not archive_cmd:
+            return
+
+        walseg_regex = re.compile(r'^[0-9A-F]{24}(\.partial){0,1}\.ready$')
+        status_dir = os.path.join(self._postgresql.wal_dir, 'archive_status')
+        try:
+            wals_to_archive = [f[:-6] for f in os.listdir(status_dir) if walseg_regex.match(f)]
+        except OSError as e:
+            return logger.error('Unable to list %s: %r', status_dir, e)
+
+        # skip fsync, as postgres --single or pg_rewind will anyway run it
+        for wal in sorted(wals_to_archive):
+            if os.path.isfile(os.path.join(self._postgresql.wal_dir, wal)):
+                cmd = self._buid_archiver_command(archive_cmd, wal)
+                # it is the author of archive_command, who is responsible
+                # for not overriding the WALs already present in archive
+                logger.info('Trying to archive %s: %s', wal, cmd)
+                if self._postgresql.cancellable.call(shlex.split(cmd)) == 0:
+                    old_name = os.path.join(status_dir, wal + '.ready')
+                    new_name = os.path.join(status_dir, wal + '.done')
+                    try:
+                        shutil.move(old_name, new_name)
+                    except Exception as e:
+                        logger.error('Unable to rename %s to %s: %r', old_name, new_name, e)
+                else:
+                    logger.info('Failed to archive WAL segment %s', wal)
+
     def pg_rewind(self, r):
         # prepare pg_rewind connection
         env = self._postgresql.config.write_pgpass(r)
@@ -367,6 +413,8 @@ class Rewind(object):
         if self._postgresql.is_running() and not self._postgresql.stop(checkpoint=False):
             return logger.warning('Can not run pg_rewind because postgres is still running')
 
+        self._archive_ready_wals()
+
         # prepare pg_rewind connection
         r = self._conn_kwargs(leader, self._postgresql.config.rewind_credentials)
 
@@ -465,6 +513,7 @@ class Rewind(object):
             logger.exception('Unable to list %s', status_dir)
 
     def ensure_clean_shutdown(self):
+        self._archive_ready_wals()
         self.cleanup_archive_status()
 
         # Start in a single user mode and stop to produce a clean shutdown


Test patch:
diff --git a/tests/test_rewind.py b/tests/test_rewind.py
index 2d510b33..b9398205 100644
--- a/tests/test_rewind.py
+++ b/tests/test_rewind.py
@@ -218,6 +218,62 @@ class TestRewind(BaseTestPostgresql):
         self.r.cleanup_archive_status()
         self.r.cleanup_archive_status()
 
+    @patch('os.path.isfile', Mock(return_value=True))
+    @patch('shutil.move', Mock(side_effect=OSError))
+    @patch('patroni.postgresql.rewind.logger.info')
+    def test_archive_ready_wals(self, mock_logger_info):
+        with patch('os.listdir', Mock(side_effect=OSError)), \
+              patch.object(Postgresql, 'get_guc_value', Mock(side_effect=['on', 'command %f'])):
+            self.r._archive_ready_wals()
+            mock_logger_info.assert_not_called()
+
+        # each assert_not_called() calls get_guc_value('archive_mode') + get_guc_value('archive_command')
+        get_guc_value_res = [
+            '', 'command %f',
+            'on', '',
+        ]
+        with patch.object(Postgresql, 'get_guc_value', Mock(side_effect=get_guc_value_res)):
+            for _ in range(len(get_guc_value_res)//2):
+                self.r._archive_ready_wals()
+                mock_logger_info.assert_not_called()
+
+        with patch('os.listdir', Mock(return_value=['000000000000000000000000.ready'])):
+            # successful archive_command call
+            with patch.object(CancellableSubprocess, 'call',  Mock(return_value=0)):
+                get_guc_value_res = [
+                    'on', 'command %f',
+                    'always', 'command %f',
+                ]
+                with patch.object(Postgresql, 'get_guc_value', Mock(side_effect=get_guc_value_res)):
+                    for _ in range(len(get_guc_value_res)//2):
+                        self.r._archive_ready_wals()
+                        mock_logger_info.assert_called_once()
+                        self.assertEqual(('Trying to archive %s: %s',
+                                          '000000000000000000000000', 'command 000000000000000000000000'),
+                                         mock_logger_info.call_args[0])
+                        mock_logger_info.reset_mock()
+
+            # failed archive_command call
+            with patch.object(CancellableSubprocess, 'call', Mock(return_value=1)):
+                with patch.object(Postgresql, 'get_guc_value', Mock(side_effect=['on', 'command %f'])):
+                    self.r._archive_ready_wals()
+                    self.assertEqual(('Trying to archive %s: %s',
+                                      '000000000000000000000000', 'command 000000000000000000000000'),
+                                     mock_logger_info.call_args_list[0][0])
+                    self.assertEqual(('Failed to archive WAL segment %s', '000000000000000000000000'),
+                                     mock_logger_info.call_args_list[1][0])
+                    mock_logger_info.reset_mock()
+
+        wal_files_to_skip = [
+            '000000000000000000000000.done',
+            '000000000000000000000001.partial.done',
+            '002.ready',
+            'U00000000000000000000001.ready',
+        ]
+        with patch('os.listdir', Mock(return_value=wal_files_to_skip)):
+            self.r._archive_ready_wals()
+            mock_logger_info.assert_not_called()
+
     @patch('os.unlink', Mock())
     @patch('os.listdir', Mock(return_value=[]))
     @patch('os.path.isfile', Mock(return_value=True))