Reuse zstd decompression contexts when restoring WAL images

Restoring each zstd-compressed full-page image currently creates and
destroys a decompression context.  One decompression context gets now
saved in a XLogReaderState on first use, reused for the decompression of
any follow-up zstd-compressed pages.

The decompression context is freed in XLogReaderFree().  In non-FRONTEND
builds, a memory context reset callback is registered to make sure that
the decompression context does not leak.  Note that zstd has APIs to
either initialize a pre-allocated decompression context or assign custom
allocation and free routines, but none of these are available for
builds that dynamically link to the zstd library (ZSTDLIB_STATIC_API).

Tomas Vondra has done some benchmarking showing that for a FPI-heavy set
of WAL records, reusing the decompression context can nicely improve the
recovery time, bringing it close to lz4.  In my own measurements, I have
seen up to a 4% faster recovery time when replaying records made mostly
of zstd-compressed full-page images between HEAD and this improvement.
Andrey has seen better numbers than that.

Author: Andrey Borodin <x4mmm@yandex-team.ru>
Reviewed-by: Michael Paquier <michael@paquier.xyz>
Discussion: https://www.postgresql.org/message-id/asRsOl73nf8GeBFd@paquier.xyz

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 decomoression context. */
+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;

src/include/access/xlogreader.h

@@ -239,6 +239,15 @@ struct XLogReaderState
 	 * ----------------------------------------
 	 */
 
+#ifdef USE_ZSTD
+	/* Decompression context reused for zstd-compressed full-page images. */
+	void	   *zstd_dctx;
+#ifndef FRONTEND
+	/* Reset callback for zstd_dctx */
+	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