v5-0002-Compress-ReorderBuffer-spill-files-using-LZ4.patch

From 3fa3afc0f0ae660b5ab7f6c7851f749adea9e421 Mon Sep 17 00:00:00 2001
From: Julien Tachoires <julmon@gmail.com>
Date: Tue, 22 Oct 2024 20:40:42 +0200
Subject: [PATCH 2/6] Compress ReorderBuffer spill files using LZ4

---
 src/backend/replication/logical/Makefile      |   1 +
 src/backend/replication/logical/meson.build   |   1 +
 .../replication/logical/reorderbuffer.c       |  89 +++++-
 .../replication/logical/reorderbuffer_lz4.c   | 268 ++++++++++++++++++
 .../replication/reorderbuffer_compression.h   |  49 ++++
 5 files changed, 407 insertions(+), 1 deletion(-)
 create mode 100644 src/backend/replication/logical/reorderbuffer_lz4.c

diff --git a/src/backend/replication/logical/Makefile b/src/backend/replication/logical/Makefile
index 96c733d009..58dce86258 100644
--- a/src/backend/replication/logical/Makefile
+++ b/src/backend/replication/logical/Makefile
@@ -27,6 +27,7 @@ OBJS = \
 	relation.o \
 	reorderbuffer.o \
 	reorderbuffer_pglz.o \
+	reorderbuffer_lz4.o \
 	slotsync.o \
 	snapbuild.o \
 	tablesync.o \
diff --git a/src/backend/replication/logical/meson.build b/src/backend/replication/logical/meson.build
index 48df907b35..dc2f55d0fa 100644
--- a/src/backend/replication/logical/meson.build
+++ b/src/backend/replication/logical/meson.build
@@ -13,6 +13,7 @@ backend_sources += files(
   'relation.c',
   'reorderbuffer.c',
   'reorderbuffer_pglz.c',
+  'reorderbuffer_lz4.c',
   'slotsync.c',
   'snapbuild.c',
   'tablesync.c',
diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index 649b196d5f..5d6f4bbfca 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_PGLZ_COMPRESSION;
+int			logical_decoding_spill_compression = REORDER_BUFFER_LZ4_COMPRESSION;
 
 /* ---------------------------------------
  * primary reorderbuffer support routines
@@ -5408,6 +5408,9 @@ ReorderBufferNewCompressorState(MemoryContext context, int compression_method)
 		case REORDER_BUFFER_PGLZ_COMPRESSION:
 			return pglz_NewCompressorState(context);
 			break;
+		case REORDER_BUFFER_LZ4_COMPRESSION:
+			return lz4_NewCompressorState(context);
+			break;
 		case REORDER_BUFFER_NO_COMPRESSION:
 		default:
 			return NULL;
@@ -5428,6 +5431,9 @@ ReorderBufferFreeCompressorState(MemoryContext context, int compression_method,
 		case REORDER_BUFFER_PGLZ_COMPRESSION:
 			return pglz_FreeCompressorState(context, compressor_state);
 			break;
+		case REORDER_BUFFER_LZ4_COMPRESSION:
+			return lz4_FreeCompressorState(context, compressor_state);
+			break;
 		case REORDER_BUFFER_NO_COMPRESSION:
 		default:
 			break;
@@ -5494,6 +5500,67 @@ ReorderBufferCompress(ReorderBuffer *rb, ReorderBufferDiskChange **ondisk,
 
 				break;
 			}
+			/* LZ4 Compression */
+		case REORDER_BUFFER_LZ4_COMPRESSION:
+			{
+				Size		dst_size = 0;
+				char	   *src = (char *) rb->outbuf + sizeof(ReorderBufferDiskChange);
+				Size		src_size = data_size - sizeof(ReorderBufferDiskChange);
+				StringInfo	buf = lz4_GetStringInfoBuffer(compressor_state);
+
+				/*
+				 * LZ4 streaming compression implies keeping a copy of the
+				 * "raw" data in LZ4 input ring buffer. If the "raw" data does
+				 * not fit in this buffer, then we should not try to compress
+				 * it. Let's store it uncompressed.
+				 *
+				 * Even if individual changes larger than 64kB shouldn't
+				 * exist, we still want to be sure that this case is covered
+				 * anyway.
+				 */
+				if (unlikely(src_size > LZ4_RING_BUFFER_SIZE))
+					return ReorderBufferCompress(rb, ondisk,
+												 REORDER_BUFFER_NO_COMPRESSION,
+												 data_size, compressor_state);
+
+				/*
+				 * Make sure the buffer we'll use to store compressed data has
+				 * enough space.
+				 */
+				dst_size = lz4_CompressBound(src_size);
+				enlargeStringInfo(buf, dst_size);
+
+				/* Use LZ4 streaming compression */
+				lz4_StreamingCompressData(rb->context, src, src_size, buf->data,
+										  &dst_size, compressor_state);
+				buf->len = dst_size;
+
+				/*
+				 * Make sure the ReorderBuffer has enough space to store
+				 * compressed data. Compressed data must be smaller than raw
+				 * data, so, the ReorderBuffer should already have room for
+				 * compressed data, but we do this to avoid buffer overflow
+				 * risks.
+				 */
+				ReorderBufferSerializeReserve(rb, (dst_size + sizeof(ReorderBufferDiskChange)));
+
+				hdr = (ReorderBufferDiskChange *) rb->outbuf;
+				hdr->comp_strat = REORDER_BUFFER_STRAT_LZ4_STREAMING;
+				hdr->size = dst_size + sizeof(ReorderBufferDiskChange);
+				hdr->raw_size = src_size;
+
+				/*
+				 * Update ondisk: hdr pointer has potentially changed due to
+				 * ReorderBufferSerializeReserve()
+				 */
+				*ondisk = hdr;
+
+				/* Copy back compressed data into the ReorderBuffer */
+				memcpy((char *) rb->outbuf + sizeof(ReorderBufferDiskChange),
+					   buf->data, buf->len);
+
+				break;
+			}
 		default:
 			/* Other compression methods not yet supported */
 			break;
@@ -5549,6 +5616,26 @@ ReorderBufferDecompress(ReorderBuffer *rb, char *data,
 							 errmsg_internal("compressed PGLZ data is corrupted")));
 				break;
 			}
+			/* LZ4 streaming decompression */
+		case REORDER_BUFFER_STRAT_LZ4_STREAMING:
+			{
+				char	   *buf = NULL;
+				Size		src_size = ondisk->size - sizeof(ReorderBufferDiskChange);
+				Size		buf_size = ondisk->raw_size;
+
+				lz4_StreamingDecompressData(rb->context, data, src_size, &buf,
+											buf_size, compressor_state);
+
+				/* Copy decompressed data into the ReorderBuffer */
+				memcpy((char *) rb->outbuf + sizeof(ReorderBufferDiskChange),
+					   buf, buf_size);
+
+				/*
+				 * Not necessary to free buf in this case: it points to the
+				 * decompressed data stored in LZ4 output ring buffer.
+				 */
+				break;
+			}
 		default:
 			/* Other compression methods not yet supported */
 			break;
diff --git a/src/backend/replication/logical/reorderbuffer_lz4.c b/src/backend/replication/logical/reorderbuffer_lz4.c
new file mode 100644
index 0000000000..b9a7e717aa
--- /dev/null
+++ b/src/backend/replication/logical/reorderbuffer_lz4.c
@@ -0,0 +1,268 @@
+/*-------------------------------------------------------------------------
+ *
+ * reorderbuffer_lz4.c
+ *	  Functions for ReorderBuffer compression using LZ4.
+ *
+ * Copyright (c) 2024-2024, PostgreSQL Global Development Group
+ *
+ *
+ * IDENTIFICATION
+ *	  src/backend/access/common/reorderbuffer_lz4.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#ifdef USE_LZ4
+#include <lz4.h>
+#endif
+
+#include "replication/reorderbuffer_compression.h"
+
+#define NO_LZ4_SUPPORT() \
+	ereport(ERROR, \
+			(errcode(ERRCODE_FEATURE_NOT_SUPPORTED), \
+			 errmsg("compression method lz4 not supported"), \
+			 errdetail("This functionality requires the server to be built with lz4 support.")))
+
+/*
+ * Allocate a new LZ4StreamingCompressorState.
+ */
+void *
+lz4_NewCompressorState(MemoryContext context)
+{
+#ifndef USE_LZ4
+	NO_LZ4_SUPPORT();
+	return NULL;				/* keep compiler quiet */
+#else
+	LZ4StreamingCompressorState *cstate;
+	MemoryContext oldcontext = MemoryContextSwitchTo(context);
+
+	cstate = (LZ4StreamingCompressorState *)
+		MemoryContextAlloc(context, sizeof(LZ4StreamingCompressorState));
+
+	cstate->buf = makeStringInfo();
+
+	/*
+	 * We do not allocate LZ4 ring buffers and streaming handlers 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->lz4_in_buf = NULL;
+	cstate->lz4_out_buf = NULL;
+	cstate->lz4_in_buf_offset = 0;
+	cstate->lz4_out_buf_offset = 0;
+	cstate->lz4_stream = NULL;
+	cstate->lz4_stream_decode = NULL;
+
+	MemoryContextSwitchTo(oldcontext);
+
+	return (void *) cstate;
+#endif
+}
+
+/*
+ * Free LZ4 memory resources and the compressor state.
+ */
+void
+lz4_FreeCompressorState(MemoryContext context, void *compressor_state)
+{
+#ifndef USE_LZ4
+	NO_LZ4_SUPPORT();
+#else
+	LZ4StreamingCompressorState *cstate;
+	MemoryContext oldcontext;
+
+	if (compressor_state == NULL)
+		return;
+
+	oldcontext = MemoryContextSwitchTo(context);
+
+	cstate = (LZ4StreamingCompressorState *) compressor_state;
+
+	destroyStringInfo(cstate->buf);
+
+	if (cstate->lz4_in_buf != NULL)
+	{
+		pfree(cstate->lz4_in_buf);
+		LZ4_freeStream(cstate->lz4_stream);
+	}
+	if (cstate->lz4_out_buf != NULL)
+	{
+		pfree(cstate->lz4_out_buf);
+		LZ4_freeStreamDecode(cstate->lz4_stream_decode);
+	}
+
+	pfree(compressor_state);
+
+	MemoryContextSwitchTo(oldcontext);
+#endif
+}
+
+#ifdef USE_LZ4
+/*
+ * Allocate LZ4 input ring buffer and create the streaming compression handler.
+ */
+static void
+lz4_CreateStreamCompressorState(MemoryContext context, void *compressor_state)
+{
+	LZ4StreamingCompressorState *cstate;
+	MemoryContext oldcontext = MemoryContextSwitchTo(context);
+
+	cstate = (LZ4StreamingCompressorState *) compressor_state;
+	cstate->lz4_in_buf = (char *) palloc0(LZ4_RING_BUFFER_SIZE);
+	cstate->lz4_stream = LZ4_createStream();
+
+	MemoryContextSwitchTo(oldcontext);
+}
+#endif
+
+#ifdef USE_LZ4
+/*
+ * Allocate LZ4 output ring buffer and create the streaming decompression
+ */
+static void
+lz4_CreateStreamDecodeCompressorState(MemoryContext context,
+									  void *compressor_state)
+{
+	LZ4StreamingCompressorState *cstate;
+	MemoryContext oldcontext = MemoryContextSwitchTo(context);
+
+	cstate = (LZ4StreamingCompressorState *) compressor_state;
+	cstate->lz4_out_buf = (char *) palloc0(LZ4_RING_BUFFER_SIZE);
+	cstate->lz4_stream_decode = LZ4_createStreamDecode();
+
+	MemoryContextSwitchTo(oldcontext);
+}
+#endif
+
+/*
+ * Data compression using LZ4 streaming API.
+ */
+void
+lz4_StreamingCompressData(MemoryContext context, char *src, Size src_size,
+						  char *dst, Size *dst_size, void *compressor_state)
+{
+#ifndef USE_LZ4
+	NO_LZ4_SUPPORT();
+#else
+	LZ4StreamingCompressorState *cstate;
+	int			lz4_cmp_size = 0;	/* compressed size */
+	char	   *lz4_in_bufPtr;	/* input ring buffer pointer */
+
+	cstate = (LZ4StreamingCompressorState *) compressor_state;
+
+	/* Allocate LZ4 input ring buffer and streaming compression handler */
+	if (cstate->lz4_in_buf == NULL)
+		lz4_CreateStreamCompressorState(context, compressor_state);
+
+	/* Ring buffer offset wraparound */
+	if ((cstate->lz4_in_buf_offset + src_size) > LZ4_RING_BUFFER_SIZE)
+		cstate->lz4_in_buf_offset = 0;
+
+	/* Get the pointer of the next entry in the ring buffer */
+	lz4_in_bufPtr = cstate->lz4_in_buf + cstate->lz4_in_buf_offset;
+
+	/* Copy data that should be compressed into LZ4 input ring buffer */
+	memcpy(lz4_in_bufPtr, src, src_size);
+
+	/* Use LZ4 streaming compression API */
+	lz4_cmp_size = LZ4_compress_fast_continue(cstate->lz4_stream,
+											  lz4_in_bufPtr, dst, src_size,
+											  *dst_size, 1);
+
+	if (lz4_cmp_size <= 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_DATA_CORRUPTED),
+				 errmsg_internal("LZ4 compression failed")));
+
+	/* Move the input ring buffer offset */
+	cstate->lz4_in_buf_offset += src_size;
+
+	*dst_size = lz4_cmp_size;
+#endif
+}
+
+/*
+ * Data decompression using LZ4 streaming API.
+ * LZ4 decompression uses the output ring buffer to store decompressed data,
+ * thus, we don't need to create a new buffer. We return the pointer to data
+ * location.
+ */
+void
+lz4_StreamingDecompressData(MemoryContext context, char *src, Size src_size,
+							char **dst, Size dst_size, void *compressor_state)
+{
+#ifndef USE_LZ4
+	NO_LZ4_SUPPORT();
+#else
+	LZ4StreamingCompressorState *cstate;
+	char	   *lz4_out_bufPtr; /* output ring buffer pointer */
+	int			lz4_dec_size;	/* decompressed data size */
+
+	cstate = (LZ4StreamingCompressorState *) compressor_state;
+
+	/* Allocate LZ4 output ring buffer and streaming decompression handler */
+	if (cstate->lz4_out_buf == NULL)
+		lz4_CreateStreamDecodeCompressorState(context, compressor_state);
+
+	/* Ring buffer offset wraparound */
+	if ((cstate->lz4_out_buf_offset + dst_size) > LZ4_RING_BUFFER_SIZE)
+		cstate->lz4_out_buf_offset = 0;
+
+	/* Get current entry pointer in the ring buffer */
+	lz4_out_bufPtr = cstate->lz4_out_buf + cstate->lz4_out_buf_offset;
+
+	lz4_dec_size = LZ4_decompress_safe_continue(cstate->lz4_stream_decode,
+												src,
+												lz4_out_bufPtr,
+												src_size,
+												dst_size);
+
+	Assert(lz4_dec_size == dst_size);
+
+	if (lz4_dec_size < 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_DATA_CORRUPTED),
+				 errmsg_internal("compressed LZ4 data is corrupted")));
+	else if (lz4_dec_size != dst_size)
+		ereport(ERROR,
+				(errcode(ERRCODE_DATA_CORRUPTED),
+				 errmsg_internal("decompressed LZ4 data size differs from original size")));
+
+	/* Move the output ring buffer offset */
+	cstate->lz4_out_buf_offset += lz4_dec_size;
+
+	/* Point to the decompressed data location */
+	*dst = lz4_out_bufPtr;
+#endif
+}
+
+Size
+lz4_CompressBound(Size src_size)
+{
+#ifndef USE_LZ4
+	NO_LZ4_SUPPORT();
+	return -1;
+#else
+	return LZ4_COMPRESSBOUND(src_size);
+#endif
+}
+
+/*
+ * Returns the StringInfo buffer we use to store compressed/decompressed data.
+ */
+StringInfo
+lz4_GetStringInfoBuffer(void *compressor_state)
+{
+#ifndef USE_LZ4
+	NO_LZ4_SUPPORT();
+	return NULL;
+#else
+	LZ4StreamingCompressorState *cstate;
+
+	cstate = (LZ4StreamingCompressorState *) compressor_state;
+
+	return cstate->buf;
+#endif
+}
diff --git a/src/include/replication/reorderbuffer_compression.h b/src/include/replication/reorderbuffer_compression.h
index 9e9565ca7f..3dbf47e18e 100644
--- a/src/include/replication/reorderbuffer_compression.h
+++ b/src/include/replication/reorderbuffer_compression.h
@@ -13,11 +13,16 @@
 #ifndef REORDERBUFFER_COMPRESSION_H
 #define REORDERBUFFER_COMPRESSION_H
 
+#ifdef USE_LZ4
+#include <lz4.h>
+#endif
+
 /* ReorderBuffer on disk compression algorithms */
 typedef enum ReorderBufferCompressionMethod
 {
 	REORDER_BUFFER_NO_COMPRESSION,
 	REORDER_BUFFER_PGLZ_COMPRESSION,
+	REORDER_BUFFER_LZ4_COMPRESSION,
 }			ReorderBufferCompressionMethod;
 
 /*
@@ -27,6 +32,7 @@ typedef enum ReorderBufferCompressionStrategy
 {
 	REORDER_BUFFER_STRAT_UNCOMPRESSED,
 	REORDER_BUFFER_STRAT_PGLZ,
+	REORDER_BUFFER_STRAT_LZ4_STREAMING,
 }			ReorderBufferCompressionStrategy;
 
 typedef struct PGLZCompressorState
@@ -35,6 +41,49 @@ typedef struct PGLZCompressorState
 	StringInfo	buf;
 }			PGLZCompressorState;
 
+#ifdef USE_LZ4
+/*
+ * We use a fairly small LZ4 ring buffer size (64kB). Using a larger buffer
+ * size provide better compression ratio, but as long as we have to allocate
+ * two LZ4 ring buffers per ReorderBufferTXN, we should keep it small.
+
+ * 64kB is also twice the maximum size of a block, which is enough to cover
+ * changes like UPDATE that will contain data of the old and new version of a
+ * tuple.
+ */
+#define LZ4_RING_BUFFER_SIZE (64 * 1024)
+
+/*
+ * LZ4 streaming compression/decompression contextes and buffers.
+ */
+typedef struct LZ4StreamingCompressorState
+{
+	/* Streaming compression handler */
+	LZ4_stream_t *lz4_stream;
+	/* Streaming decompression handler */
+	LZ4_streamDecode_t *lz4_stream_decode;
+	/* LZ4 in/out ring buffers used for streaming compression */
+	char	   *lz4_in_buf;
+	int			lz4_in_buf_offset;
+	char	   *lz4_out_buf;
+	int			lz4_out_buf_offset;
+	/* Buffer used to store compressed data */
+	StringInfo	buf;
+}			LZ4StreamingCompressorState;
+#endif
+
+extern void *lz4_NewCompressorState(MemoryContext context);
+extern void lz4_FreeCompressorState(MemoryContext context,
+									void *compressor_state);
+extern void lz4_StreamingCompressData(MemoryContext context, char *src,
+									  Size src_size, char *dst, Size *dst_size,
+									  void *compressor_state);
+extern void lz4_StreamingDecompressData(MemoryContext context, char *src,
+										Size src_size, char **dst,
+										Size dst_size, void *compressor_state);
+extern Size lz4_CompressBound(Size src_size);
+extern StringInfo lz4_GetStringInfoBuffer(void *compressor_state);
+
 extern void *pglz_NewCompressorState(MemoryContext context);
 extern void pglz_FreeCompressorState(MemoryContext context,
 									 void *compressor_state);
-- 
2.43.0