From 27d33e09549d7f96b750750256ed37ec2eb9a04d Mon Sep 17 00:00:00 2001 From: Julien Tachoires 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 +#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 #endif +#ifdef USE_ZSTD +#include +#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