From 2dac70f61b9511099d885d950ec1c3d2b20e382f Mon Sep 17 00:00:00 2001
From: "Sami Imseih (AWS)" <simseih@amazon.com>
Date: Thu, 6 Aug 2026 21:04:43 +0000
Subject: [PATCH v1 1/1] tests

---
 .../test_custom_stats/t/001_custom_stats.pl   |  60 +++++
 .../test_custom_var_stats--1.0.sql            |  15 ++
 .../test_custom_stats/test_custom_var_stats.c | 206 ++++++++++++++++++
 src/tools/pgindent/typedefs.list              |   2 +
 4 files changed, 283 insertions(+)

diff --git a/src/test/modules/test_custom_stats/t/001_custom_stats.pl b/src/test/modules/test_custom_stats/t/001_custom_stats.pl
index 69f2284229e..ee5ff546ff7 100644
--- a/src/test/modules/test_custom_stats/t/001_custom_stats.pl
+++ b/src/test/modules/test_custom_stats/t/001_custom_stats.pl
@@ -168,5 +168,65 @@ $result = $node->safe_psql('postgres',
 );
 is($result, "0", "report of fixed-sized after manual reset");
 
+# Test cascade flush (A -> B -> C).  A's flush accumulates into B, B's flush
+# accumulates into C.  Seeding in the order C, B, A forces B to accumulate into
+# C only after C was already visited in the flush pass, so C must be re-queued
+# (moved to the list tail) to receive it.  Results are checked before COMMIT so
+# the transaction-end flush cannot mask a failed in-transaction re-queue.
+$node->safe_psql('postgres',
+	q(select test_custom_stats_var_create('cascade_test', 'cascade test')));
+
+my $b_before = $node->safe_psql('postgres',
+	q(SELECT coalesce(test_cascade_get_count('B'), 0)));
+my $c_before = $node->safe_psql('postgres',
+	q(SELECT coalesce(test_cascade_get_count('C'), 0)));
+
+my $bgpsql = $node->background_psql('postgres', on_error_stop => 1);
+$bgpsql->query_safe(qq(
+	SET stats_fetch_consistency = none;
+	BEGIN;
+	SELECT test_cascade_seed('C');
+	SELECT test_cascade_seed('B');
+	SELECT test_custom_stats_var_update('cascade_test');
+	SELECT test_custom_stats_var_update('cascade_test');
+	SELECT test_custom_stats_var_update('cascade_test');
+	SELECT pg_stat_force_next_flush();
+));
+
+my $b_after = $bgpsql->query_safe(
+	q(SELECT test_cascade_get_count('B')));
+my $c_after = $bgpsql->query_safe(
+	q(SELECT test_cascade_get_count('C')));
+
+is($b_after - $b_before, 3,
+	"cascade B receives 3 from A");
+is($c_after - $c_before, 3,
+	"cascade C receives 3 from B");
+
+$bgpsql->query_safe(qq(COMMIT));
+$bgpsql->quit;
+
+# Test that a callback which re-queues its OWN entry mid-flush does not abort
+# the flush scan.  test_selfrep_arm() arms object D; D's flush callback preps a
+# terminal entry E (which follows D on the pending list) and then re-preps D
+# itself.  If the self-reprep moved D (the scan cursor) to the list tail, the
+# scan would exit early and E would never be flushed.  E must end up at 1.
+my $selfrep = $node->background_psql('postgres', on_error_stop => 1);
+$selfrep->query_safe(qq(
+	SET stats_fetch_consistency = none;
+	BEGIN;
+	SELECT test_selfrep_arm();
+	SELECT pg_stat_force_next_flush();
+));
+
+my $d_count = $selfrep->query_safe(q(SELECT test_cascade_get_count('D')));
+my $e_count = $selfrep->query_safe(q(SELECT test_cascade_get_count('E')));
+
+is($d_count, "1", "self-reprep D flushed");
+is($e_count, "1", "self-reprep does not skip the following entry E");
+
+$selfrep->query_safe(qq(COMMIT));
+$selfrep->quit;
+
 # Test completed successfully
 done_testing();
diff --git a/src/test/modules/test_custom_stats/test_custom_var_stats--1.0.sql b/src/test/modules/test_custom_stats/test_custom_var_stats--1.0.sql
index 5ed8cfc2dcf..2de5b9aee4a 100644
--- a/src/test/modules/test_custom_stats/test_custom_var_stats--1.0.sql
+++ b/src/test/modules/test_custom_stats/test_custom_var_stats--1.0.sql
@@ -24,3 +24,18 @@ CREATE FUNCTION test_custom_stats_var_report(INOUT name TEXT,
 RETURNS SETOF record
 AS 'MODULE_PATHNAME', 'test_custom_stats_var_report'
 LANGUAGE C STRICT PARALLEL UNSAFE;
+
+CREATE FUNCTION test_cascade_get_count(IN which TEXT)
+RETURNS BIGINT
+AS 'MODULE_PATHNAME', 'test_cascade_get_count'
+LANGUAGE C STRICT PARALLEL UNSAFE;
+
+CREATE FUNCTION test_cascade_seed(IN which TEXT)
+RETURNS void
+AS 'MODULE_PATHNAME', 'test_cascade_seed'
+LANGUAGE C STRICT PARALLEL UNSAFE;
+
+CREATE FUNCTION test_selfrep_arm()
+RETURNS void
+AS 'MODULE_PATHNAME', 'test_selfrep_arm'
+LANGUAGE C STRICT PARALLEL UNSAFE;
diff --git a/src/test/modules/test_custom_stats/test_custom_var_stats.c b/src/test/modules/test_custom_stats/test_custom_var_stats.c
index 7ce44c7cbaf..25010b64246 100644
--- a/src/test/modules/test_custom_stats/test_custom_var_stats.c
+++ b/src/test/modules/test_custom_stats/test_custom_var_stats.c
@@ -43,6 +43,31 @@ PG_MODULE_MAGIC_EXT(
  */
 #define PGSTAT_KIND_TEST_CUSTOM_VAR_STATS 25
 
+/*
+ * Kind ID for the flush re-queue tests.  A single kind is used with several
+ * objects (objids); representing the dependent entries as distinct objects of
+ * one kind rather than distinct kinds keeps this to a single reserved ID while
+ * exercising the same per-entry re-queue mechanics.
+ */
+#define PGSTAT_KIND_CASCADE 27
+
+/*
+ * Object IDs within PGSTAT_KIND_CASCADE.
+ *
+ * B and C model a cascade: the test_custom_var_stats flush (A) accumulates
+ * into B, and B's flush accumulates into C.  Seeding C, B, A forces B to feed
+ * C only after C was already visited, so C must be re-queued.
+ *
+ * D and E model a self-reprep: D's flush accumulates into a following terminal
+ * entry E and also re-preps D itself.  If the re-queue relocated D (the scan
+ * cursor) to the tail, the scan would exit early and E would never be flushed;
+ * the flushed_this_pass gate must prevent that.
+ */
+#define CASCADE_OBJ_B	1
+#define CASCADE_OBJ_C	2
+#define CASCADE_OBJ_D	3
+#define CASCADE_OBJ_E	4
+
 /* File paths for auxiliary data serialization */
 #define TEST_CUSTOM_AUX_DATA_DESC "pg_stat/test_custom_var_stats_desc.stats"
 
@@ -70,6 +95,18 @@ typedef struct PgStatShared_CustomVarEntry
 	dsa_pointer description;	/* pointer to description string in DSA */
 } PgStatShared_CustomVarEntry;
 
+/* Pending and shared types for the PGSTAT_KIND_CASCADE objects */
+typedef struct PgStat_CascadeEntry
+{
+	PgStat_Counter count;
+} PgStat_CascadeEntry;
+
+typedef struct PgStatShared_CascadeEntry
+{
+	PgStatShared_Common header;
+	PgStat_CascadeEntry stats;
+} PgStatShared_CascadeEntry;
+
 /*--------------------------------------------------------------------------
  * Global Variables
  *--------------------------------------------------------------------------
@@ -128,6 +165,83 @@ static const PgStat_KindInfo custom_stats = {
 	.finish = test_custom_stats_var_finish,
 };
 
+/*
+ * cascade_flush_cb
+ *
+ * Single flush callback for PGSTAT_KIND_CASCADE, dispatching on the object id.
+ * All objects share the same pending/shared layout (PgStat_CascadeEntry) and
+ * flush their own count to shared memory.
+ *
+ * B (CASCADE_OBJ_B) accumulates into C; C is terminal.  D (CASCADE_OBJ_D)
+ * accumulates into terminal E and re-preps itself once.  The self-reprep must
+ * not relocate D (the scan cursor); the flushed_this_pass gate prevents that
+ * because D's flag is cleared just before this callback runs.
+ */
+static PgStat_FlushResult
+cascade_flush_cb(PgStat_EntryRef *entry_ref, bool nowait, bool xact_boundary)
+{
+	uint64		objid = entry_ref->shared_entry->key.objid;
+	PgStat_CascadeEntry *pending;
+	PgStatShared_CascadeEntry *shared;
+	PgStat_Counter shared_after;
+
+	pending = (PgStat_CascadeEntry *) entry_ref->pending;
+	shared = (PgStatShared_CascadeEntry *) entry_ref->shared_stats;
+
+	if (pending->count == 0)
+		return PGSTAT_FLUSH_DONE;
+
+	if (!pgstat_lock_entry(entry_ref, nowait))
+		return PGSTAT_FLUSH_LOCK_CONFLICT;
+
+	shared->stats.count += pending->count;
+	shared_after = shared->stats.count;
+	pgstat_unlock_entry(entry_ref);
+
+	if (objid == CASCADE_OBJ_B)
+	{
+		/* Accumulate into C. */
+		PgStat_EntryRef *c_ref =
+			pgstat_prep_pending_entry(PGSTAT_KIND_CASCADE, InvalidOid,
+									  CASCADE_OBJ_C, NULL);
+
+		((PgStat_CascadeEntry *) c_ref->pending)->count += pending->count;
+	}
+	else if (objid == CASCADE_OBJ_D)
+	{
+		/* Accumulate into terminal E (which follows D on the pending list). */
+		PgStat_EntryRef *e_ref =
+			pgstat_prep_pending_entry(PGSTAT_KIND_CASCADE, InvalidOid,
+									  CASCADE_OBJ_E, NULL);
+
+		((PgStat_CascadeEntry *) e_ref->pending)->count += 1;
+
+		/* Re-prep our own entry once, bounded so it cannot loop. */
+		if (shared_after < 2)
+		{
+			PgStat_EntryRef *d_ref =
+				pgstat_prep_pending_entry(PGSTAT_KIND_CASCADE, InvalidOid,
+										  CASCADE_OBJ_D, NULL);
+
+			((PgStat_CascadeEntry *) d_ref->pending)->count += 1;
+		}
+	}
+
+	pending->count = 0;
+	return PGSTAT_FLUSH_DONE;
+}
+
+static const PgStat_KindInfo cascade_stats = {
+	.name = "cascade",
+	.fixed_amount = false,
+	.accessed_across_databases = true,
+	.shared_size = sizeof(PgStatShared_CascadeEntry),
+	.shared_data_off = offsetof(PgStatShared_CascadeEntry, stats),
+	.shared_data_len = sizeof(((PgStatShared_CascadeEntry *) 0)->stats),
+	.pending_size = sizeof(PgStat_CascadeEntry),
+	.flush_pending_cb = cascade_flush_cb,
+};
+
 /*--------------------------------------------------------------------------
  * Module initialization
  *--------------------------------------------------------------------------
@@ -138,6 +252,7 @@ _PG_init(void)
 {
 	/* Register custom statistics kind */
 	pgstat_register_kind(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS, &custom_stats);
+	pgstat_register_kind(PGSTAT_KIND_CASCADE, &cascade_stats);
 }
 
 /*--------------------------------------------------------------------------
@@ -162,10 +277,15 @@ test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref, bool nowait,
 {
 	PgStat_StatCustomVarEntry *pending_entry;
 	PgStatShared_CustomVarEntry *shared_entry;
+	PgStat_EntryRef *b_ref;
+	PgStat_CascadeEntry *b_pending;
 
 	pending_entry = (PgStat_StatCustomVarEntry *) entry_ref->pending;
 	shared_entry = (PgStatShared_CustomVarEntry *) entry_ref->shared_stats;
 
+	if (pending_entry->numcalls == 0)
+		return PGSTAT_FLUSH_DONE;
+
 	if (!pgstat_lock_entry(entry_ref, nowait))
 		return PGSTAT_FLUSH_LOCK_CONFLICT;
 
@@ -174,6 +294,14 @@ test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref, bool nowait,
 
 	pgstat_unlock_entry(entry_ref);
 
+	/* Accumulate into cascade object B */
+	b_ref = pgstat_prep_pending_entry(PGSTAT_KIND_CASCADE, InvalidOid,
+									  CASCADE_OBJ_B, NULL);
+	b_pending = (PgStat_CascadeEntry *) b_ref->pending;
+	b_pending->count += pending_entry->numcalls;
+
+	memset(pending_entry, 0, sizeof(*pending_entry));
+
 	return PGSTAT_FLUSH_DONE;
 }
 
@@ -705,3 +833,81 @@ test_custom_stats_var_report(PG_FUNCTION_ARGS)
 
 	SRF_RETURN_DONE(funcctx);
 }
+
+/*
+ * Map a cascade object name ('B', 'C', 'D', 'E') to its object id within
+ * PGSTAT_KIND_CASCADE.
+ */
+static uint64
+cascade_objid(const char *which)
+{
+	if (strcmp(which, "B") == 0)
+		return CASCADE_OBJ_B;
+	else if (strcmp(which, "C") == 0)
+		return CASCADE_OBJ_C;
+	else if (strcmp(which, "D") == 0)
+		return CASCADE_OBJ_D;
+	else if (strcmp(which, "E") == 0)
+		return CASCADE_OBJ_E;
+
+	ereport(ERROR, (errmsg("argument must be 'B', 'C', 'D' or 'E'")));
+}
+
+/*
+ * test_cascade_get_count
+ *		Read the count from a cascade object's shared stats.
+ */
+PG_FUNCTION_INFO_V1(test_cascade_get_count);
+Datum
+test_cascade_get_count(PG_FUNCTION_ARGS)
+{
+	char	   *which = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	PgStat_CascadeEntry *entry;
+
+	entry = (PgStat_CascadeEntry *)
+		pgstat_fetch_entry(PGSTAT_KIND_CASCADE, InvalidOid,
+						   cascade_objid(which), NULL);
+
+	if (!entry)
+		PG_RETURN_INT64(0);
+
+	PG_RETURN_INT64(entry->count);
+}
+
+/*
+ * test_cascade_seed
+ *		Put a cascade object on the pending list with zero data.
+ */
+PG_FUNCTION_INFO_V1(test_cascade_seed);
+Datum
+test_cascade_seed(PG_FUNCTION_ARGS)
+{
+	char	   *which = text_to_cstring(PG_GETARG_TEXT_PP(0));
+
+	pgstat_prep_pending_entry(PGSTAT_KIND_CASCADE, InvalidOid,
+							  cascade_objid(which), NULL);
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * test_selfrep_arm
+ *		Put cascade object D on the pending list with a nonzero count so its
+ *		flush callback fires (which re-preps D and preps terminal E).  Used to
+ *		test that a self-reprepping callback does not abort the flush scan
+ *		early.
+ */
+PG_FUNCTION_INFO_V1(test_selfrep_arm);
+Datum
+test_selfrep_arm(PG_FUNCTION_ARGS)
+{
+	PgStat_EntryRef *ref;
+	PgStat_CascadeEntry *pending;
+
+	ref = pgstat_prep_pending_entry(PGSTAT_KIND_CASCADE, InvalidOid,
+									CASCADE_OBJ_D, NULL);
+	pending = (PgStat_CascadeEntry *) ref->pending;
+	pending->count += 1;
+
+	PG_RETURN_VOID();
+}
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 689e0164013..d6caa23836b 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2318,6 +2318,7 @@ PgIfAddrCallback
 PgStatShared_Archiver
 PgStatShared_Backend
 PgStatShared_BgWriter
+PgStatShared_CascadeEntry
 PgStatShared_Checkpointer
 PgStatShared_Common
 PgStatShared_CustomFixedEntry
@@ -2338,6 +2339,7 @@ PgStat_BackendPending
 PgStat_BackendSubEntry
 PgStat_BgWriterStats
 PgStat_BktypeIO
+PgStat_CascadeEntry
 PgStat_CheckpointerStats
 PgStat_Counter
 PgStat_EntryRef
-- 
2.50.1 (Apple Git-155)

