From 9253013c789ffb121272bfeeaa9dcdebbef79ced Mon Sep 17 00:00:00 2001
From: Justin Pryzby <pryzbyj@telsasoft.com>
Date: Fri, 18 Feb 2022 22:54:18 -0600
Subject: [PATCH 1/6] add wal_compression=zstd

---
 doc/src/sgml/config.sgml                      |  9 +++---
 doc/src/sgml/installation.sgml                | 10 +++++++
 src/backend/access/transam/xloginsert.c       | 30 ++++++++++++++++++-
 src/backend/access/transam/xlogreader.c       | 20 +++++++++++++
 src/backend/utils/misc/guc.c                  |  3 ++
 src/backend/utils/misc/postgresql.conf.sample |  2 +-
 src/bin/pg_waldump/pg_waldump.c               |  2 ++
 src/include/access/xlog.h                     |  3 +-
 src/include/access/xlogrecord.h               |  5 +++-
 9 files changed, 76 insertions(+), 8 deletions(-)

diff --git a/doc/src/sgml/config.sgml b/doc/src/sgml/config.sgml
index 7ed8c82a9d..97740c7b66 100644
--- a/doc/src/sgml/config.sgml
+++ b/doc/src/sgml/config.sgml
@@ -3154,10 +3154,11 @@ include_dir 'conf.d'
         server compresses full page images written to WAL when
         <xref linkend="guc-full-page-writes"/> is on or during a base backup.
         A compressed page image will be decompressed during WAL replay.
-        The supported methods are <literal>pglz</literal> and
-        <literal>lz4</literal> (if <productname>PostgreSQL</productname> was
-        compiled with <option>--with-lz4</option>). The default value is
-        <literal>off</literal>. Only superusers can change this setting.
+        The supported methods are <literal>pglz</literal>, and
+        (if configured when <productname>PostgreSQL</productname> was built) 
+        <literal>lz4</literal> and zstd.
+        The default value is <literal>off</literal>.
+        Only superusers can change this setting.
        </para>
 
        <para>
diff --git a/doc/src/sgml/installation.sgml b/doc/src/sgml/installation.sgml
index 0f74252590..d32767b229 100644
--- a/doc/src/sgml/installation.sgml
+++ b/doc/src/sgml/installation.sgml
@@ -271,6 +271,14 @@ su - postgres
      </para>
     </listitem>
 
+    <listitem>
+     <para>
+      The <productname>ZSTD</productname> library can be used to enable
+      compression using that method; see
+      <xref linkend="guc-wal-compression"/>.
+     </para>
+    </listitem>
+
     <listitem>
      <para>
       To build the <productname>PostgreSQL</productname> documentation,
@@ -988,6 +996,8 @@ build-postgresql:
        <listitem>
         <para>
          Build with <productname>ZSTD</productname> compression support.
+         This enables use of <productname>ZSTD</productname> for
+         compression of WAL data.
         </para>
        </listitem>
       </varlistentry>
diff --git a/src/backend/access/transam/xloginsert.c b/src/backend/access/transam/xloginsert.c
index c260310c4c..847ce0c3fc 100644
--- a/src/backend/access/transam/xloginsert.c
+++ b/src/backend/access/transam/xloginsert.c
@@ -44,9 +44,17 @@
 #define LZ4_MAX_BLCKSZ		0
 #endif
 
+#ifdef USE_ZSTD
+#include <zstd.h>
+#define ZSTD_MAX_BLCKSZ		ZSTD_COMPRESSBOUND(BLCKSZ)
+#else
+#define ZSTD_MAX_BLCKSZ		0
+#endif
+
+/* Buffer size required to store a compressed version of backup block image */
 #define PGLZ_MAX_BLCKSZ		PGLZ_MAX_OUTPUT(BLCKSZ)
 
-#define COMPRESS_BUFSIZE	Max(PGLZ_MAX_BLCKSZ, LZ4_MAX_BLCKSZ)
+#define COMPRESS_BUFSIZE	Max(Max(PGLZ_MAX_BLCKSZ, LZ4_MAX_BLCKSZ), ZSTD_MAX_BLCKSZ)
 
 /*
  * For each block reference registered with XLogRegisterBuffer, we fill in
@@ -695,6 +703,14 @@ XLogRecordAssemble(RmgrId rmid, uint8 info,
 #endif
 						break;
 
+					case WAL_COMPRESSION_ZSTD:
+#ifdef USE_ZSTD
+						bimg.bimg_info |= BKPIMAGE_COMPRESS_ZSTD;
+#else
+						elog(ERROR, "ZSTD is not supported by this build");
+#endif
+						break;
+
 					case WAL_COMPRESSION_NONE:
 						Assert(false);	/* cannot happen */
 						break;
@@ -903,6 +919,18 @@ XLogCompressBackupBlock(char *page, uint16 hole_offset, uint16 hole_length,
 #endif
 			break;
 
+		case WAL_COMPRESSION_ZSTD:
+#ifdef USE_ZSTD
+			/* Uses level=1, not ZSTD_CLEVEL_DEFAULT */
+			len = ZSTD_compress(dest, COMPRESS_BUFSIZE, source, orig_len,
+								1);
+			if (ZSTD_isError(len))
+				len = -1;
+#else
+			elog(ERROR, "ZSTD is not supported by this build");
+#endif
+			break;
+
 		case WAL_COMPRESSION_NONE:
 			Assert(false);		/* cannot happen */
 			break;
diff --git a/src/backend/access/transam/xlogreader.c b/src/backend/access/transam/xlogreader.c
index 35029cf97d..d60e4cbaf6 100644
--- a/src/backend/access/transam/xlogreader.c
+++ b/src/backend/access/transam/xlogreader.c
@@ -21,6 +21,9 @@
 #ifdef USE_LZ4
 #include <lz4.h>
 #endif
+#ifdef USE_ZSTD
+#include <zstd.h>
+#endif
 
 #include "access/transam.h"
 #include "access/xlog_internal.h"
@@ -1618,6 +1621,23 @@ RestoreBlockImage(XLogReaderState *record, uint8 block_id, char *page)
 								  "LZ4",
 								  block_id);
 			return false;
+#endif
+		}
+		else if ((bkpb->bimg_info & BKPIMAGE_COMPRESS_ZSTD) != 0)
+		{
+#ifdef USE_ZSTD
+			size_t decomp_result = ZSTD_decompress(tmp.data,
+												   BLCKSZ-bkpb->hole_length,
+												   ptr, bkpb->bimg_len);
+
+			if (ZSTD_isError(decomp_result))
+				decomp_success = false;
+#else
+			report_invalid_record(record, "image at %X/%X compressed with %s not supported by build, block %d",
+								  LSN_FORMAT_ARGS(record->ReadRecPtr),
+								  "zstd",
+								  block_id);
+			return false;
 #endif
 		}
 		else
diff --git a/src/backend/utils/misc/guc.c b/src/backend/utils/misc/guc.c
index e4afd07bfe..66e6d664b0 100644
--- a/src/backend/utils/misc/guc.c
+++ b/src/backend/utils/misc/guc.c
@@ -550,6 +550,9 @@ static const struct config_enum_entry wal_compression_options[] = {
 	{"pglz", WAL_COMPRESSION_PGLZ, false},
 #ifdef USE_LZ4
 	{"lz4", WAL_COMPRESSION_LZ4, false},
+#endif
+#ifdef USE_ZSTD
+	{"zstd", WAL_COMPRESSION_ZSTD, false},
 #endif
 	{"on", WAL_COMPRESSION_PGLZ, false},
 	{"off", WAL_COMPRESSION_NONE, false},
diff --git a/src/backend/utils/misc/postgresql.conf.sample b/src/backend/utils/misc/postgresql.conf.sample
index 4a094bb38b..4cf5b26a36 100644
--- a/src/backend/utils/misc/postgresql.conf.sample
+++ b/src/backend/utils/misc/postgresql.conf.sample
@@ -220,7 +220,7 @@
 #wal_log_hints = off			# also do full page writes of non-critical updates
 					# (change requires restart)
 #wal_compression = off			# enables compression of full-page writes;
-					# off, pglz, lz4, or on
+					# off, pglz, lz4, zstd, or on
 #wal_init_zero = on			# zero-fill new WAL files
 #wal_recycle = on			# recycle WAL files
 #wal_buffers = -1			# min 32kB, -1 sets based on shared_buffers
diff --git a/src/bin/pg_waldump/pg_waldump.c b/src/bin/pg_waldump/pg_waldump.c
index a6251e1a96..c448658ea2 100644
--- a/src/bin/pg_waldump/pg_waldump.c
+++ b/src/bin/pg_waldump/pg_waldump.c
@@ -565,6 +565,8 @@ XLogDumpDisplayRecord(XLogDumpConfig *config, XLogReaderState *record)
 						method = "pglz";
 					else if ((bimg_info & BKPIMAGE_COMPRESS_LZ4) != 0)
 						method = "lz4";
+					else if ((bimg_info & BKPIMAGE_COMPRESS_ZSTD) != 0)
+						method = "zstd";
 					else
 						method = "unknown";
 
diff --git a/src/include/access/xlog.h b/src/include/access/xlog.h
index 4b45ac64db..09f6464331 100644
--- a/src/include/access/xlog.h
+++ b/src/include/access/xlog.h
@@ -75,7 +75,8 @@ typedef enum WalCompression
 {
 	WAL_COMPRESSION_NONE = 0,
 	WAL_COMPRESSION_PGLZ,
-	WAL_COMPRESSION_LZ4
+	WAL_COMPRESSION_LZ4,
+	WAL_COMPRESSION_ZSTD
 } WalCompression;
 
 /* Recovery states */
diff --git a/src/include/access/xlogrecord.h b/src/include/access/xlogrecord.h
index c1b1137aa7..052ac6817a 100644
--- a/src/include/access/xlogrecord.h
+++ b/src/include/access/xlogrecord.h
@@ -149,8 +149,11 @@ typedef struct XLogRecordBlockImageHeader
 /* compression methods supported */
 #define BKPIMAGE_COMPRESS_PGLZ	0x04
 #define BKPIMAGE_COMPRESS_LZ4	0x08
+#define BKPIMAGE_COMPRESS_ZSTD	0x10
+
 #define	BKPIMAGE_COMPRESSED(info) \
-	((info & (BKPIMAGE_COMPRESS_PGLZ | BKPIMAGE_COMPRESS_LZ4)) != 0)
+	((info & (BKPIMAGE_COMPRESS_PGLZ | BKPIMAGE_COMPRESS_LZ4 | \
+			  BKPIMAGE_COMPRESS_ZSTD)) != 0)
 
 /*
  * Extra header information used when page image has "hole" and
-- 
2.17.1

