From 177c2dfe1ea463cc43a4856a466454bbea4ece2d Mon Sep 17 00:00:00 2001
From: Andrey Borodin <amborodin@acm.org>
Date: Tue, 6 Oct 2026 11:27:54 +0500
Subject: [PATCH v10 1/2] Reuse zstd decompression contexts when restoring WAL
 images

Restoring each zstd-compressed full-page image currently creates and
destroys a decompression context.  Avoid that overhead by keeping one
context per WAL reader, allocated on first use.

Free the context in XLogReaderFree().  In backend builds, also register
a reset callback with the reader's memory context so that errors do not
leak the allocation made by zstd.

Author: Andrey Borodin <amborodin@acm.org>
Reviewed-by: Michael Paquier <michael@paquier.xyz>
Discussion: https://www.postgresql.org/message-id/asRsOl73nf8GeBFd@paquier.xyz
---
 src/backend/access/transam/xlogreader.c | 51 +++++++++++++++++++++++--
 src/include/access/xlogreader.h         | 13 +++++++
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/src/backend/access/transam/xlogreader.c b/src/backend/access/transam/xlogreader.c
index 7db7c273b0c..4497eb52ec5 100644
--- a/src/backend/access/transam/xlogreader.c
+++ b/src/backend/access/transam/xlogreader.c
@@ -36,6 +36,7 @@
 #ifndef FRONTEND
 #include "pgstat.h"
 #include "storage/bufmgr.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 #else
 #include "common/logging.h"
@@ -55,6 +56,9 @@ static bool ValidXLogRecord(XLogReaderState *state, XLogRecord *record,
 static void ResetDecoder(XLogReaderState *state);
 static void WALOpenSegmentInit(WALOpenSegment *seg, WALSegmentContext *segcxt,
 							   int segsize, const char *waldir);
+#if defined(USE_ZSTD) && !defined(FRONTEND)
+static void XLogReaderFreeZstdContext(void *arg);
+#endif
 
 /* size of the buffer allocated for error message. */
 #define MAX_ERRORMSG_LEN 1000
@@ -171,10 +175,29 @@ XLogReaderFree(XLogReaderState *state)
 	pfree(state->errormsg_buf);
 	if (state->readRecordBuf)
 		pfree(state->readRecordBuf);
+#ifdef USE_ZSTD
+	if (state->zstd_dctx != NULL)
+	{
+#ifndef FRONTEND
+		MemoryContextUnregisterResetCallback(GetMemoryChunkContext(state),
+											 &state->zstd_dctx_cb);
+#endif
+		ZSTD_freeDCtx(state->zstd_dctx);
+	}
+#endif
 	pfree(state->readBuf);
 	pfree(state);
 }
 
+#if defined(USE_ZSTD) && !defined(FRONTEND)
+/* Release the zstd context even if an error bypasses XLogReaderFree(). */
+static void
+XLogReaderFreeZstdContext(void *arg)
+{
+	ZSTD_freeDCtx(arg);
+}
+#endif
+
 /*
  * Allocate readRecordBuf to fit a record of at least the given length.
  *
@@ -2177,9 +2200,31 @@ RestoreBlockImage(XLogReaderState *record, uint8 block_id, char *page)
 		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);
+			size_t		decomp_result;
+
+			if (record->zstd_dctx == NULL)
+			{
+				record->zstd_dctx = ZSTD_createDCtx();
+				if (record->zstd_dctx == NULL)
+				{
+					report_invalid_record(record, "out of memory while restoring image at %X/%08X, block %d",
+										  LSN_FORMAT_ARGS(record->ReadRecPtr),
+										  block_id);
+					return false;
+				}
+
+#ifndef FRONTEND
+				/* The reader may outlive the current memory context. */
+				record->zstd_dctx_cb.func = XLogReaderFreeZstdContext;
+				record->zstd_dctx_cb.arg = record->zstd_dctx;
+				MemoryContextRegisterResetCallback(GetMemoryChunkContext(record),
+												   &record->zstd_dctx_cb);
+#endif
+			}
+
+			decomp_result = ZSTD_decompressDCtx(record->zstd_dctx, tmp.data,
+												BLCKSZ - bkpb->hole_length,
+												ptr, bkpb->bimg_len);
 
 			if (ZSTD_isError(decomp_result))
 				decomp_success = false;
diff --git a/src/include/access/xlogreader.h b/src/include/access/xlogreader.h
index 4a9a687e879..f139a61d046 100644
--- a/src/include/access/xlogreader.h
+++ b/src/include/access/xlogreader.h
@@ -34,8 +34,13 @@
 #ifndef XLOGREADER_H
 #define XLOGREADER_H
 
+#ifdef USE_ZSTD
+#include <zstd.h>
+#endif
+
 #ifndef FRONTEND
 #include "access/transam.h"
+#include "utils/palloc.h"
 #endif
 
 #include "access/xlogrecord.h"
@@ -239,6 +244,14 @@ struct XLogReaderState
 	 * ----------------------------------------
 	 */
 
+#ifdef USE_ZSTD
+	/* Reused to decompress zstd-compressed full-page images. */
+	ZSTD_DCtx  *zstd_dctx;
+#ifndef FRONTEND
+	MemoryContextCallback zstd_dctx_cb;
+#endif
+#endif
+
 	/*
 	 * Buffer for decoded records.  This is a circular buffer, though
 	 * individual records can't be split in the middle, so some space is often
-- 
That's all, folks. May the source be with you.