From 7e6ab1d0ce578486fb96b55a56162f71fe5a370c Mon Sep 17 00:00:00 2001
From: Bertrand Drouvot <bertranddrouvot.pg@gmail.com>
Date: Thu, 5 Feb 2026 05:54:34 +0000
Subject: [PATCH v11 2/5] Add anytime flush tests for custom stats

---
 .../test_custom_stats/t/001_custom_stats.pl   | 43 ++++++++++++++
 .../test_custom_fixed_stats--1.0.sql          |  5 ++
 .../test_custom_fixed_stats.c                 | 57 +++++++++++++++++++
 .../test_custom_var_stats--1.0.sql            |  5 ++
 .../test_custom_stats/test_custom_var_stats.c | 27 +++++++++
 5 files changed, 137 insertions(+)
  35.8% src/test/modules/test_custom_stats/t/
  64.1% src/test/modules/test_custom_stats/

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 9e6a7a38577..6ba4022418f 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
@@ -156,5 +156,48 @@ $result = $node->safe_psql('postgres',
 );
 is($result, "0", "report of fixed-sized after manual reset");
 
+# Test FLUSH_ANYTIME mechanism with custom fixed stats
+# This verifies that custom stats can be flushed during a transaction
+
+# Reset stats first
+$node->safe_psql('postgres', q(select test_custom_stats_fixed_reset()));
+$node->safe_psql('postgres', q(select pg_stat_force_next_flush()));
+
+my $anytime_test = q[
+    BEGIN;
+    SET LOCAL stats_fetch_consistency = none;
+    -- Accumulate stats
+    select test_custom_stats_fixed_anytime_update() from generate_series(1, 2);
+    -- Wait (has to be greater than PGSTAT_MIN_INTERVAL)
+    select pg_sleep(1.5);
+    -- Check
+    select 'fixed_anytime:'||numcalls from test_custom_stats_fixed_report();
+];
+
+$result = $node->safe_psql('postgres', $anytime_test);
+like($result, qr/^fixed_anytime:2/m,
+	"anytime fixed stats flushed during transaction");
+
+# Test FLUSH_ANYTIME mechanism with custom variable stats
+# This verifies that custom stats can be flushed during a transaction
+
+$node->safe_psql('postgres', q(select pg_stat_force_next_flush()));
+
+$anytime_test = q[
+    BEGIN;
+    SET LOCAL stats_fetch_consistency = none;
+    -- Accumulate stats
+    select test_custom_stats_var_anytime_update('entry2');
+    select test_custom_stats_var_anytime_update('entry2');
+    -- Wait (has to be greater than PGSTAT_MIN_INTERVAL)
+    select pg_sleep(1.5);
+    -- Check
+    select 'var_anytime:'||calls from test_custom_stats_var_report('entry2');
+];
+
+$result = $node->safe_psql('postgres', $anytime_test);
+like($result, qr/^var_anytime:2/m,
+	"anytime var stats flushed during transaction");
+
 # Test completed successfully
 done_testing();
diff --git a/src/test/modules/test_custom_stats/test_custom_fixed_stats--1.0.sql b/src/test/modules/test_custom_stats/test_custom_fixed_stats--1.0.sql
index 69a93b5241f..da3a798f289 100644
--- a/src/test/modules/test_custom_stats/test_custom_fixed_stats--1.0.sql
+++ b/src/test/modules/test_custom_stats/test_custom_fixed_stats--1.0.sql
@@ -18,3 +18,8 @@ CREATE FUNCTION test_custom_stats_fixed_reset()
 RETURNS void
 AS 'MODULE_PATHNAME', 'test_custom_stats_fixed_reset'
 LANGUAGE C STRICT PARALLEL UNSAFE;
+
+CREATE FUNCTION test_custom_stats_fixed_anytime_update()
+RETURNS void
+AS 'MODULE_PATHNAME'
+LANGUAGE C STRICT PARALLEL UNSAFE;
diff --git a/src/test/modules/test_custom_stats/test_custom_fixed_stats.c b/src/test/modules/test_custom_stats/test_custom_fixed_stats.c
index 485e08e5c19..e7fbb2737ef 100644
--- a/src/test/modules/test_custom_stats/test_custom_fixed_stats.c
+++ b/src/test/modules/test_custom_stats/test_custom_fixed_stats.c
@@ -18,6 +18,7 @@
 #include "pgstat.h"
 #include "utils/builtins.h"
 #include "utils/pgstat_internal.h"
+#include "utils/timeout.h"
 #include "utils/timestamp.h"
 
 PG_MODULE_MAGIC_EXT(
@@ -44,11 +45,13 @@ typedef struct PgStatShared_CustomFixedEntry
 static void test_custom_stats_fixed_init_shmem_cb(void *stats);
 static void test_custom_stats_fixed_reset_all_cb(TimestampTz ts);
 static void test_custom_stats_fixed_snapshot_cb(void);
+static bool test_custom_stats_fixed_flush_cb(bool nowait, bool anytime_only);
 
 static const PgStat_KindInfo custom_stats = {
 	.name = "test_custom_fixed_stats",
 	.fixed_amount = true,		/* exactly one entry */
 	.write_to_file = true,		/* persist to stats file */
+	.flush_mode = FLUSH_ANYTIME,	/* can be flushed anytime */
 
 	.shared_size = sizeof(PgStat_StatCustomFixedEntry),
 	.shared_data_off = offsetof(PgStatShared_CustomFixedEntry, stats),
@@ -57,8 +60,12 @@ static const PgStat_KindInfo custom_stats = {
 	.init_shmem_cb = test_custom_stats_fixed_init_shmem_cb,
 	.reset_all_cb = test_custom_stats_fixed_reset_all_cb,
 	.snapshot_cb = test_custom_stats_fixed_snapshot_cb,
+	.flush_static_cb = test_custom_stats_fixed_flush_cb,
 };
 
+/* Pending statistics */
+static PgStat_StatCustomFixedEntry PendingCustomStats = {0};
+
 /*
  * Kind ID for test_custom_fixed_stats.
  */
@@ -142,6 +149,38 @@ test_custom_stats_fixed_snapshot_cb(void)
 #undef FIXED_COMP
 }
 
+/*
+ * test_custom_stats_fixed_flush_cb
+ *		Flush pending stats to shared memory
+ */
+static bool
+test_custom_stats_fixed_flush_cb(bool nowait, bool anytime_only)
+{
+	PgStatShared_CustomFixedEntry *stats_shmem;
+
+	/* Nothing to flush if no calls were made */
+	if (PendingCustomStats.numcalls == 0)
+		return false;
+
+	stats_shmem = pgstat_get_custom_shmem_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
+
+	if (!nowait)
+		LWLockAcquire(&stats_shmem->lock, LW_EXCLUSIVE);
+	else if (!LWLockConditionalAcquire(&stats_shmem->lock, LW_EXCLUSIVE))
+		return true;
+
+	pgstat_begin_changecount_write(&stats_shmem->changecount);
+	stats_shmem->stats.numcalls += PendingCustomStats.numcalls;
+	pgstat_end_changecount_write(&stats_shmem->changecount);
+
+	LWLockRelease(&stats_shmem->lock);
+
+	/* Reset pending stats */
+	PendingCustomStats.numcalls = 0;
+
+	return false;				/* successfully flushed */
+}
+
 /*--------------------------------------------------------------------------
  * SQL-callable functions
  *--------------------------------------------------------------------------
@@ -223,3 +262,21 @@ test_custom_stats_fixed_report(PG_FUNCTION_ARGS)
 	/* Return as tuple */
 	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
 }
+
+/*
+ * test_custom_stats_fixed_anytime_update
+ *		Increment call counter and schedule anytime flush
+ */
+PG_FUNCTION_INFO_V1(test_custom_stats_fixed_anytime_update);
+Datum
+test_custom_stats_fixed_anytime_update(PG_FUNCTION_ARGS)
+{
+	/* Accumulate in pending stats */
+	PendingCustomStats.numcalls++;
+
+	/* Schedule anytime stats update */
+	pgstat_schedule_anytime_update();
+	pgstat_report_fixed = true;
+
+	PG_RETURN_VOID();
+}
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..ed66d38981e 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,8 @@ 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_custom_stats_var_anytime_update(IN name TEXT)
+RETURNS void
+AS 'MODULE_PATHNAME', 'test_custom_stats_var_anytime_update'
+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 4c207611236..e9f1bda6b32 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
@@ -18,6 +18,7 @@
 #include "storage/dsm_registry.h"
 #include "utils/builtins.h"
 #include "utils/pgstat_internal.h"
+#include "utils/timeout.h"
 
 PG_MODULE_MAGIC_EXT(
 					.name = "test_custom_var_stats",
@@ -108,6 +109,7 @@ static const PgStat_KindInfo custom_stats = {
 	.name = "test_custom_var_stats",
 	.fixed_amount = false,		/* variable number of entries */
 	.write_to_file = true,		/* persist across restarts */
+	.flush_mode = FLUSH_ANYTIME,	/* can be flushed anytime */
 	.track_entry_count = true,	/* count active entries */
 	.accessed_across_databases = true,	/* global statistics */
 	.shared_size = sizeof(PgStatShared_CustomVarEntry),
@@ -690,3 +692,28 @@ test_custom_stats_var_report(PG_FUNCTION_ARGS)
 
 	SRF_RETURN_DONE(funcctx);
 }
+
+/*
+ * test_custom_stats_var_anytime_update
+ *		Increment custom statistic counter and schedule anytime flush
+ */
+PG_FUNCTION_INFO_V1(test_custom_stats_var_anytime_update);
+Datum
+test_custom_stats_var_anytime_update(PG_FUNCTION_ARGS)
+{
+	char	   *stat_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	PgStat_EntryRef *entry_ref;
+	PgStat_StatCustomVarEntry *pending_entry;
+
+	/* Get pending entry in local memory */
+	entry_ref = pgstat_prep_pending_entry(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS, InvalidOid,
+										  PGSTAT_CUSTOM_VAR_STATS_IDX(stat_name), NULL);
+
+	pending_entry = (PgStat_StatCustomVarEntry *) entry_ref->pending;
+	pending_entry->numcalls++;
+
+	/* Schedule anytime stats update */
+	pgstat_schedule_anytime_update();
+
+	PG_RETURN_VOID();
+}
-- 
2.34.1

