From aec7c54ea0473e8ddc57f03793fb1aefe7e868e8 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Tue, 1 Apr 2025 21:01:07 +0300
Subject: [PATCH] TEST: add consume_multixids function

---
 src/test/modules/xid_wraparound/Makefile      |  1 +
 src/test/modules/xid_wraparound/meson.build   |  1 +
 .../xid_wraparound/multixid_wraparound.c      | 96 +++++++++++++++++++
 .../xid_wraparound/xid_wraparound--1.0.sql    |  4 +
 4 files changed, 102 insertions(+)
 create mode 100644 src/test/modules/xid_wraparound/multixid_wraparound.c

diff --git a/src/test/modules/xid_wraparound/Makefile b/src/test/modules/xid_wraparound/Makefile
index 7a6e0f667629..ebb3d8fcb3ec 100644
--- a/src/test/modules/xid_wraparound/Makefile
+++ b/src/test/modules/xid_wraparound/Makefile
@@ -3,6 +3,7 @@
 MODULE_big = xid_wraparound
 OBJS = \
 	$(WIN32RES) \
+	multixid_wraparound.o \
 	xid_wraparound.o
 PGFILEDESC = "xid_wraparound - tests for XID wraparound"
 
diff --git a/src/test/modules/xid_wraparound/meson.build b/src/test/modules/xid_wraparound/meson.build
index f7dada67f678..98ad381614cf 100644
--- a/src/test/modules/xid_wraparound/meson.build
+++ b/src/test/modules/xid_wraparound/meson.build
@@ -1,6 +1,7 @@
 # Copyright (c) 2023-2025, PostgreSQL Global Development Group
 
 xid_wraparound_sources = files(
+  'multixid_wraparound.c',
   'xid_wraparound.c',
 )
 
diff --git a/src/test/modules/xid_wraparound/multixid_wraparound.c b/src/test/modules/xid_wraparound/multixid_wraparound.c
new file mode 100644
index 000000000000..af567c6e541a
--- /dev/null
+++ b/src/test/modules/xid_wraparound/multixid_wraparound.c
@@ -0,0 +1,96 @@
+/*--------------------------------------------------------------------------
+ *
+ * multixid_wraparound.c
+ *		Utilities for testing multixids
+ *
+ *
+ * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *		src/test/modules/xid_wraparound/multixid_wraparound.c
+ *
+ * -------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "access/multixact.h"
+#include "access/xact.h"
+#include "miscadmin.h"
+#include "storage/proc.h"
+#include "utils/xid8.h"
+
+static int mxactMemberComparator(const void *arg1, const void *arg2);
+
+/*
+ * Consume the specified number of multi-XIDs, with specified number of
+ * members each.
+ */
+PG_FUNCTION_INFO_V1(consume_multixids);
+Datum
+consume_multixids(PG_FUNCTION_ARGS)
+{
+	int64		nmultis = PG_GETARG_INT64(0);
+	int32		nmembers = PG_GETARG_INT32(1);
+	MultiXactMember *members;
+	MultiXactId	lastmxid = InvalidMultiXactId;
+
+	if (nmultis < 0)
+		elog(ERROR, "invalid nxids argument: %" PRId64, nmultis);
+	if (nmembers < 1)
+		elog(ERROR, "invalid nmembers argument: %d", nmembers);
+
+	/*
+	 * We consume XIDs by calling GetNewTransactionId(true), which marks the
+	 * consumed XIDs as subtransactions of the current top-level transaction.
+	 * For that to work, this transaction must have a top-level XID.
+	 *
+	 * GetNewTransactionId registers them in the subxid cache in PGPROC, until
+	 * the cache overflows, but beyond that, we don't keep track of the
+	 * consumed XIDs.
+	 */
+	(void) GetTopTransactionId();
+
+	members = palloc((nmultis + nmembers) * sizeof(MultiXactMember));
+	for (int32 i = 0; i < nmultis + nmembers; i++)
+	{
+		FullTransactionId xid;
+
+		xid = GetNewTransactionId(true);
+		members[i].xid = XidFromFullTransactionId(xid);
+		members[i].status = MultiXactStatusForKeyShare;
+	}
+	/*
+	 * pre-sort the array like mXactCacheGetBySet does, so that the qsort call
+	 * in mXactCacheGetBySet() is cheaper.
+	 */
+	qsort(members, nmultis + nmembers, sizeof(MultiXactMember), mxactMemberComparator);
+
+	for (int64 i = 0; i < nmultis; i++)
+	{
+		lastmxid = MultiXactIdCreateFromMembers(nmembers, &members[i]);
+		CHECK_FOR_INTERRUPTS();
+	}
+
+	pfree(members);
+
+	PG_RETURN_TRANSACTIONID(lastmxid);
+}
+
+/* copied from multixact.c */
+static int
+mxactMemberComparator(const void *arg1, const void *arg2)
+{
+	MultiXactMember member1 = *(const MultiXactMember *) arg1;
+	MultiXactMember member2 = *(const MultiXactMember *) arg2;
+
+	if (member1.xid > member2.xid)
+		return 1;
+	if (member1.xid < member2.xid)
+		return -1;
+	if (member1.status > member2.status)
+		return 1;
+	if (member1.status < member2.status)
+		return -1;
+	return 0;
+}
diff --git a/src/test/modules/xid_wraparound/xid_wraparound--1.0.sql b/src/test/modules/xid_wraparound/xid_wraparound--1.0.sql
index 96356b4b9745..ed7520c3d86c 100644
--- a/src/test/modules/xid_wraparound/xid_wraparound--1.0.sql
+++ b/src/test/modules/xid_wraparound/xid_wraparound--1.0.sql
@@ -10,3 +10,7 @@ AS 'MODULE_PATHNAME' LANGUAGE C;
 CREATE FUNCTION consume_xids_until(targetxid xid8)
 RETURNS xid8 VOLATILE PARALLEL UNSAFE STRICT
 AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION consume_multixids(nmultis bigint, nmembers int4)
+RETURNS bigint VOLATILE PARALLEL UNSAFE STRICT
+AS 'MODULE_PATHNAME' LANGUAGE C;