v5-0003-Compress-ReorderBuffer-spill-files-using-ZSTD.patch

From 27d33e09549d7f96b750750256ed37ec2eb9a04d Mon Sep 17 00:00:00 2001
From: Julien Tachoires <julmon@gmail.com>
Date: Tue, 22 Oct 2024 22:31:13 +0200
Subject: [PATCH 3/6] Compress ReorderBuffer spill files using ZSTD

---
 src/backend/replication/logical/Makefile      |   1 +
 src/backend/replication/logical/meson.build   |   1 +
 .../replication/logical/reorderbuffer.c       |  58 ++-
 .../replication/logical/reorderbuffer_zstd.c  | 361 ++++++++++++++++++
 .../replication/reorderbuffer_compression.h   |  54 +++
 5 files changed, 474 insertions(+), 1 deletion(-)
 create mode 100644 src/backend/replication/logical/reorderbuffer_zstd.c

diff --git a/src/backend/replication/logical/Makefile b/src/backend/replication/logical/Makefile
index 58dce86258..c71f06b729 100644
--- a/src/backend/replication/logical/Makefile
+++ b/src/backend/replication/logical/Makefile
@@ -28,6 +28,7 @@ OBJS = \
 	reorderbuffer.o \
 	reorderbuffer_pglz.o \
 	reorderbuffer_lz4.o \
+	reorderbuffer_zstd.o \
 	slotsync.o \
 	snapbuild.o \
 	tablesync.o \
diff --git a/src/backend/replication/logical/meson.build b/src/backend/replication/logical/meson.build
index dc2f55d0fa..70b8777290 100644
--- a/src/backend/replication/logical/meson.build
+++ b/src/backend/replication/logical/meson.build
@@ -14,6 +14,7 @@ backend_sources += files(
   'reorderbuffer.c',
   'reorderbuffer_pglz.c',
   'reorderbuffer_lz4.c',
+  'reorderbuffer_zstd.c',
   'slotsync.c',
   'snapbuild.c',
   'tablesync.c',
diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index 5d6f4bbfca..17f8208000 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -219,7 +219,7 @@ static const Size max_changes_in_memory = 4096; /* XXX for restore only */
 int			debug_logical_replication_streaming = DEBUG_LOGICAL_REP_STREAMING_BUFFERED;
 
 /* Compression strategy for spilled data. */
-int			logical_decoding_spill_compression = REORDER_BUFFER_LZ4_COMPRESSION;
+int			logical_decoding_spill_compression = REORDER_BUFFER_ZSTD_COMPRESSION;
 
 /* ---------------------------------------
  * primary reorderbuffer support routines
@@ -5411,6 +5411,9 @@ ReorderBufferNewCompressorState(MemoryContext context, int compression_method)
 		case REORDER_BUFFER_LZ4_COMPRESSION:
 			return lz4_NewCompressorState(context);
 			break;
+		case REORDER_BUFFER_ZSTD_COMPRESSION:
+			return zstd_NewCompressorState(context);
+			break;
 		case REORDER_BUFFER_NO_COMPRESSION:
 		default:
 			return NULL;
@@ -5434,6 +5437,9 @@ ReorderBufferFreeCompressorState(MemoryContext context, int compression_method,
 		case REORDER_BUFFER_LZ4_COMPRESSION:
 			return lz4_FreeCompressorState(context, compressor_state);
 			break;
+		case REORDER_BUFFER_ZSTD_COMPRESSION:
+			return zstd_FreeCompressorState(context, compressor_state);
+			break;
 		case REORDER_BUFFER_NO_COMPRESSION:
 		default:
 			break;
@@ -5561,6 +5567,37 @@ ReorderBufferCompress(ReorderBuffer *rb, ReorderBufferDiskChange **ondisk,
 
 				break;
 			}
+			/* ZSTD Compression */
+		case REORDER_BUFFER_ZSTD_COMPRESSION:
+			{
+				Size		dst_size = 0;
+				char	   *src = (char *) rb->outbuf + sizeof(ReorderBufferDiskChange);
+				Size		src_size = data_size - sizeof(ReorderBufferDiskChange);
+				StringInfo	buf = zstd_GetStringInfoBuffer(compressor_state);
+
+				dst_size = zstd_CompressBound(src_size);
+				enlargeStringInfo(buf, dst_size);
+
+				/* Use ZSTD streaming compression */
+				zstd_StreamingCompressData(rb->context, src, src_size,
+										   buf->data, &dst_size,
+										   compressor_state);
+				buf->len = dst_size;
+
+				ReorderBufferSerializeReserve(rb, (dst_size + sizeof(ReorderBufferDiskChange)));
+
+				hdr = (ReorderBufferDiskChange *) rb->outbuf;
+				hdr->comp_strat = REORDER_BUFFER_STRAT_ZSTD_STREAMING;
+				hdr->size = dst_size + sizeof(ReorderBufferDiskChange);
+				hdr->raw_size = src_size;
+
+				*ondisk = hdr;
+
+				memcpy((char *) rb->outbuf + sizeof(ReorderBufferDiskChange),
+					   buf->data, buf->len);
+
+				break;
+			}
 		default:
 			/* Other compression methods not yet supported */
 			break;
@@ -5636,6 +5673,25 @@ ReorderBufferDecompress(ReorderBuffer *rb, char *data,
 				 */
 				break;
 			}
+			/* ZSTD streaming decompression */
+		case REORDER_BUFFER_STRAT_ZSTD_STREAMING:
+			{
+				StringInfo	buf = zstd_GetStringInfoBuffer(compressor_state);
+				Size		src_size = ondisk->size - sizeof(ReorderBufferDiskChange);
+				Size		buf_size = ondisk->raw_size;
+
+				enlargeStringInfo(buf, buf_size);
+
+				zstd_StreamingDecompressData(rb->context, data, src_size,
+											 buf->data, buf_size,
+											 compressor_state);
+				buf->len = buf_size;
+
+				/* Copy decompressed data into the ReorderBuffer */
+				memcpy((char *) rb->outbuf + sizeof(ReorderBufferDiskChange),
+					   buf->data, buf->len);
+				break;
+			}
 		default:
 			/* Other compression methods not yet supported */
 			break;
diff --git a/src/backend/replication/logical/reorderbuffer_zstd.c b/src/backend/replication/logical/reorderbuffer_zstd.c
new file mode 100644
index 0000000000..83455a9c8b
--- /dev/null
+++ b/src/backend/replication/logical/reorderbuffer_zstd.c
@@ -0,0 +1,361 @@
+/*-------------------------------------------------------------------------
+ *
+ * reorderbuffer_zstd.c
+ *	  Functions for ReorderBuffer compression using ZSTD.
+ *
+ * Copyright (c) 2024-2024, PostgreSQL Global Development Group
+ *
+ *
+ * IDENTIFICATION
+ *	  src/backend/access/common/reorderbuffer_zstd.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#ifdef USE_ZSTD
+#include <zstd.h>
+#endif
+
+#include "replication/reorderbuffer_compression.h"
+
+#define NO_ZSTD_SUPPORT() \
+	ereport(ERROR, \
+			(errcode(ERRCODE_FEATURE_NOT_SUPPORTED), \
+			 errmsg("compression method zstd not supported"), \
+			 errdetail("This functionality requires the server to be built with zstd support.")))
+
+/*
+ * Allocate a new ZSTDStreamingCompressorState.
+ */
+void *
+zstd_NewCompressorState(MemoryContext context)
+{
+#ifndef USE_ZSTD
+	NO_ZSTD_SUPPORT();
+	return NULL;				/* keep compiler quiet */
+#else
+	ZSTDStreamingCompressorState *cstate;
+	MemoryContext oldcontext = MemoryContextSwitchTo(context);
+
+	cstate = (ZSTDStreamingCompressorState *)
+		MemoryContextAlloc(context, sizeof(ZSTDStreamingCompressorState));
+
+	cstate->buf = makeStringInfo();
+
+	/*
+	 * We do not allocate ZSTD buffers and contexts at this point because we
+	 * have no guarantee that we will need them later. Let's allocate only
+	 * when we are about to use them.
+	 */
+	cstate->zstd_c_ctx = NULL;
+	cstate->zstd_c_in_buf = NULL;
+	cstate->zstd_c_in_buf_size = 0;
+	cstate->zstd_c_out_buf = NULL;
+	cstate->zstd_c_out_buf_size = 0;
+	cstate->zstd_frame_size = 0;
+	cstate->zstd_d_ctx = NULL;
+	cstate->zstd_d_in_buf = NULL;
+	cstate->zstd_d_in_buf_size = 0;
+	cstate->zstd_d_out_buf = NULL;
+	cstate->zstd_d_out_buf_size = 0;
+
+	MemoryContextSwitchTo(oldcontext);
+
+	return (void *) cstate;
+#endif
+}
+
+/*
+ * Free ZSTD memory resources and the compressor state.
+ */
+void
+zstd_FreeCompressorState(MemoryContext context, void *compressor_state)
+{
+#ifndef USE_ZSTD
+	NO_ZSTD_SUPPORT();
+#else
+	ZSTDStreamingCompressorState *cstate;
+	MemoryContext oldcontext;
+
+	if (compressor_state == NULL)
+		return;
+
+	oldcontext = MemoryContextSwitchTo(context);
+
+	cstate = (ZSTDStreamingCompressorState *) compressor_state;
+
+	destroyStringInfo(cstate->buf);
+
+	if (cstate->zstd_c_ctx != NULL)
+	{
+		/* Compressor state was used for compression */
+		pfree(cstate->zstd_c_in_buf);
+		pfree(cstate->zstd_c_out_buf);
+		ZSTD_freeCCtx(cstate->zstd_c_ctx);
+	}
+	if (cstate->zstd_d_ctx != NULL)
+	{
+		/* Compressor state was used for decompression */
+		pfree(cstate->zstd_d_in_buf);
+		pfree(cstate->zstd_d_out_buf);
+		ZSTD_freeDCtx(cstate->zstd_d_ctx);
+	}
+
+	pfree(compressor_state);
+
+	MemoryContextSwitchTo(oldcontext);
+#endif
+}
+
+#ifdef USE_ZSTD
+/*
+ * Allocate ZSTD compression buffers and create the ZSTD compression context.
+ */
+static void
+zstd_CreateStreamCompressorState(MemoryContext context, void *compressor_state)
+{
+	ZSTDStreamingCompressorState *cstate;
+	MemoryContext oldcontext = MemoryContextSwitchTo(context);
+
+	cstate = (ZSTDStreamingCompressorState *) compressor_state;
+	cstate->zstd_c_in_buf_size = ZSTD_CStreamInSize();
+	cstate->zstd_c_in_buf = (char *) palloc0(cstate->zstd_c_in_buf_size);
+	cstate->zstd_c_out_buf_size = ZSTD_CStreamOutSize();
+	cstate->zstd_c_out_buf = (char *) palloc0(cstate->zstd_c_out_buf_size);
+	cstate->zstd_c_ctx = ZSTD_createCCtx();
+
+	if (cstate->zstd_c_ctx == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("could not create ZSTD compression context")));
+
+	/* Set compression level */
+	ZSTD_CCtx_setParameter(cstate->zstd_c_ctx, ZSTD_c_compressionLevel,
+						   ZSTD_COMPRESSION_LEVEL);
+
+	MemoryContextSwitchTo(oldcontext);
+}
+#endif
+
+#ifdef USE_ZSTD
+/*
+ * Allocate ZSTD decompression buffers and create the ZSTD decompression
+ * context.
+ */
+static void
+zstd_CreateStreamDecodeCompressorState(MemoryContext context, void *compressor_state)
+{
+	ZSTDStreamingCompressorState *cstate;
+	MemoryContext oldcontext = MemoryContextSwitchTo(context);
+
+	cstate = (ZSTDStreamingCompressorState *) compressor_state;
+	cstate->zstd_d_in_buf_size = ZSTD_DStreamInSize();
+	cstate->zstd_d_in_buf = (char *) palloc0(cstate->zstd_d_in_buf_size);
+	cstate->zstd_d_out_buf_size = ZSTD_DStreamOutSize();
+	cstate->zstd_d_out_buf = (char *) palloc0(cstate->zstd_d_out_buf_size);
+	cstate->zstd_d_ctx = ZSTD_createDCtx();
+
+	if (cstate->zstd_d_ctx == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("could not create ZSTD decompression context")));
+
+	MemoryContextSwitchTo(oldcontext);
+}
+#endif
+
+/*
+ * Data compression using ZSTD streaming API.
+ */
+void
+zstd_StreamingCompressData(MemoryContext context, char *src, Size src_size,
+						   char *dst, Size *dst_size, void *compressor_state)
+{
+#ifndef USE_ZSTD
+	NO_ZSTD_SUPPORT();
+#else
+	ZSTDStreamingCompressorState *cstate;
+
+	/* Size of remaining data to be copied from src into ZSTD input buffer */
+	Size		toCpy = src_size;
+	char	   *dst_data;
+
+	cstate = (ZSTDStreamingCompressorState *) compressor_state;
+	/* Allocate ZSTD buffers and context */
+	if (cstate->zstd_c_ctx == NULL)
+		zstd_CreateStreamCompressorState(context, compressor_state);
+
+	dst_data = dst;
+	*dst_size = 0;
+
+	/*
+	 * ZSTD streaming compression works with chunks: the source data needs to
+	 * be splitted out in chunks, each of them is then copied into ZSTD input
+	 * buffer. For each chunk, we proceed with compression. Streaming
+	 * compression is not intended to compress the whole input chunk, so we
+	 * have the call ZSTD_compressStream2() multiple times until the entire
+	 * chunk is consumed.
+	 */
+	while (toCpy > 0)
+	{
+		/* Are we on the last chunk? */
+		bool		last_chunk = (toCpy < cstate->zstd_c_in_buf_size);
+
+		/* Size of the data copied into ZSTD input buffer */
+		Size		cpySize = last_chunk ? toCpy : cstate->zstd_c_in_buf_size;
+		bool		finished = false;
+		ZSTD_inBuffer input;
+		ZSTD_EndDirective mode = last_chunk ? ZSTD_e_flush : ZSTD_e_continue;
+
+		/* Copy data from src into ZSTD input buffer */
+		memcpy(cstate->zstd_c_in_buf, src, cpySize);
+
+		/*
+		 * Close the frame when we are on the last chunk and we've reached max
+		 * frame size.
+		 */
+		if (last_chunk && (cstate->zstd_frame_size > ZSTD_MAX_FRAME_SIZE))
+		{
+			mode = ZSTD_e_end;
+			cstate->zstd_frame_size = 0;
+		}
+
+		cstate->zstd_frame_size += cpySize;
+
+		input.src = cstate->zstd_c_in_buf;
+		input.size = cpySize;
+		input.pos = 0;
+
+		do
+		{
+			Size		remaining;
+			ZSTD_outBuffer output;
+
+			output.dst = cstate->zstd_c_out_buf;
+			output.size = cstate->zstd_c_out_buf_size;
+			output.pos = 0;
+
+			remaining = ZSTD_compressStream2(cstate->zstd_c_ctx, &output,
+											 &input, mode);
+
+			if (ZSTD_isError(remaining))
+				ereport(ERROR,
+						(errcode(ERRCODE_DATA_CORRUPTED),
+						 errmsg_internal("ZSTD compression failed")));
+
+			/* Copy back compressed data from ZSTD output buffer */
+			memcpy(dst_data, (char *) cstate->zstd_c_out_buf, output.pos);
+
+			dst_data += output.pos;
+			*dst_size += output.pos;
+
+			/*
+			 * Compression is done when we are working on the last chunk and
+			 * there is nothing left to compress, or, when we reach the end of
+			 * the chunk.
+			 */
+			finished = last_chunk ? (remaining == 0) : (input.pos == input.size);
+		} while (!finished);
+
+		src += cpySize;
+		toCpy -= cpySize;
+	}
+#endif
+}
+
+/*
+ * Data decompression using ZSTD streaming API.
+ */
+void
+zstd_StreamingDecompressData(MemoryContext context, char *src, Size src_size,
+							 char *dst, Size dst_size, void *compressor_state)
+{
+#ifndef USE_ZSTD
+	NO_ZSTD_SUPPORT();
+#else
+	ZSTDStreamingCompressorState *cstate;
+
+	/* Size of remaining data to be copied from src into ZSTD input buffer */
+	Size		toCpy = src_size;
+	char	   *dst_data;
+	Size		decBytes = 0;	/* Size of decompressed data */
+
+	cstate = (ZSTDStreamingCompressorState *) compressor_state;
+	/* Allocate ZSTD buffers and context */
+	if (cstate->zstd_d_ctx == NULL)
+		zstd_CreateStreamDecodeCompressorState(context, compressor_state);
+
+	dst_data = dst;
+
+	while (toCpy > 0)
+	{
+		ZSTD_inBuffer input;
+		Size		cpySize = (toCpy > cstate->zstd_d_in_buf_size) ? cstate->zstd_d_in_buf_size : toCpy;
+
+		/* Copy data from src into ZSTD input buffer */
+		memcpy(cstate->zstd_d_in_buf, src, cpySize);
+
+		input.src = cstate->zstd_d_in_buf;
+		input.size = cpySize;
+		input.pos = 0;
+
+		while (input.pos < input.size)
+		{
+			ZSTD_outBuffer output;
+			Size		ret;
+
+			output.dst = cstate->zstd_d_out_buf;
+			output.size = cstate->zstd_d_out_buf_size;
+			output.pos = 0;
+
+			ret = ZSTD_decompressStream(cstate->zstd_d_ctx, &output, &input);
+
+			if (ZSTD_isError(ret))
+				ereport(ERROR,
+						(errcode(ERRCODE_DATA_CORRUPTED),
+						 errmsg_internal("ZSTD decompression failed")));
+
+			/* Copy back compressed data from ZSTD output buffer */
+			memcpy(dst_data, (char *) cstate->zstd_d_out_buf, output.pos);
+
+			dst_data += output.pos;
+			decBytes += output.pos;
+		}
+
+		src += cpySize;
+		toCpy -= cpySize;
+	}
+
+	Assert(dst_size == decBytes);
+#endif
+}
+
+Size
+zstd_CompressBound(Size src_size)
+{
+#ifndef USE_ZSTD
+	NO_ZSTD_SUPPORT();
+	return -1;
+#else
+	return ZSTD_compressBound(src_size);
+#endif
+}
+
+/*
+ * Returns the StringInfo buffer we use to store compressed/decompressed data.
+ */
+StringInfo
+zstd_GetStringInfoBuffer(void *compressor_state)
+{
+#ifndef USE_ZSTD
+	NO_ZSTD_SUPPORT();
+	return NULL;
+#else
+	ZSTDStreamingCompressorState *cstate;
+
+	cstate = (ZSTDStreamingCompressorState *) compressor_state;
+
+	return cstate->buf;
+#endif
+}
diff --git a/src/include/replication/reorderbuffer_compression.h b/src/include/replication/reorderbuffer_compression.h
index 3dbf47e18e..240c188f00 100644
--- a/src/include/replication/reorderbuffer_compression.h
+++ b/src/include/replication/reorderbuffer_compression.h
@@ -17,12 +17,17 @@
 #include <lz4.h>
 #endif
 
+#ifdef USE_ZSTD
+#include <zstd.h>
+#endif
+
 /* ReorderBuffer on disk compression algorithms */
 typedef enum ReorderBufferCompressionMethod
 {
 	REORDER_BUFFER_NO_COMPRESSION,
 	REORDER_BUFFER_PGLZ_COMPRESSION,
 	REORDER_BUFFER_LZ4_COMPRESSION,
+	REORDER_BUFFER_ZSTD_COMPRESSION,
 }			ReorderBufferCompressionMethod;
 
 /*
@@ -33,6 +38,7 @@ typedef enum ReorderBufferCompressionStrategy
 	REORDER_BUFFER_STRAT_UNCOMPRESSED,
 	REORDER_BUFFER_STRAT_PGLZ,
 	REORDER_BUFFER_STRAT_LZ4_STREAMING,
+	REORDER_BUFFER_STRAT_ZSTD_STREAMING,
 }			ReorderBufferCompressionStrategy;
 
 typedef struct PGLZCompressorState
@@ -72,6 +78,42 @@ typedef struct LZ4StreamingCompressorState
 }			LZ4StreamingCompressorState;
 #endif
 
+#ifdef USE_ZSTD
+/*
+ * Low compression level provides high compression speed and decent compression
+ * rate. Minimum level is 1, maximum is 22.
+ */
+#define ZSTD_COMPRESSION_LEVEL 1
+
+/*
+ * Maximum volume of data encoded in the current ZSTD frame. When this
+ * threshold is reached then we close the current frame and start a new one.
+ */
+#define ZSTD_MAX_FRAME_SIZE (64 * 1024)
+
+/*
+ * ZSTD streaming compression/decompression handlers and buffers.
+ */
+typedef struct ZSTDStreamingCompressorState
+{
+	/* Compression */
+	ZSTD_CCtx  *zstd_c_ctx;
+	Size		zstd_c_in_buf_size;
+	char	   *zstd_c_in_buf;
+	Size		zstd_c_out_buf_size;
+	char	   *zstd_c_out_buf;
+	Size		zstd_frame_size;
+	/* Decompression */
+	ZSTD_DCtx  *zstd_d_ctx;
+	Size		zstd_d_in_buf_size;
+	char	   *zstd_d_in_buf;
+	Size		zstd_d_out_buf_size;
+	char	   *zstd_d_out_buf;
+	/* Buffer used to store compressed/decompressed data */
+	StringInfo	buf;
+}			ZSTDStreamingCompressorState;
+#endif
+
 extern void *lz4_NewCompressorState(MemoryContext context);
 extern void lz4_FreeCompressorState(MemoryContext context,
 									void *compressor_state);
@@ -88,4 +130,16 @@ extern void *pglz_NewCompressorState(MemoryContext context);
 extern void pglz_FreeCompressorState(MemoryContext context,
 									 void *compressor_state);
 
+extern void *zstd_NewCompressorState(MemoryContext context);
+extern void zstd_FreeCompressorState(MemoryContext context,
+									 void *compressor_state);
+extern void zstd_StreamingCompressData(MemoryContext context, char *src,
+									   Size src_size, char *dst, Size *dst_size,
+									   void *compressor_state);
+extern void zstd_StreamingDecompressData(MemoryContext context, char *src,
+										 Size src_size, char *dst,
+										 Size dst_size, void *compressor_state);
+extern Size zstd_CompressBound(Size src_size);
+extern StringInfo zstd_GetStringInfoBuffer(void *compressor_state);
+
 #endif							/* REORDERBUFFER_COMPRESSION_H */
-- 
2.43.0