From a5b49223544baede253938b71c63122869fbfdc1 Mon Sep 17 00:00:00 2001
From: Sami Imseih <samimseih@gmail.com>
Date: Sun, 10 May 2026 07:06:04 -0500
Subject: [PATCH v3 1/1] pgstat: Allow pg_stat_force_next_flush() to work
 in-transaction

Previously, pg_stat_force_next_flush() deferred the actual flush until
after the transaction ended.  Extend it to also flush immediately when
called in-transaction.

Non-transactional counters (numscans, tuples_returned, tuples_fetched,
blocks_fetched, blocks_hit) are flushed right away since they reflect
completed work that does not depend on transaction outcome.
Transactional counters (tuples_inserted/updated/deleted and the derived
live/dead tuple counts) are deferred until transaction end, since their
final values depend on commit/abort.

Also remove a test query that checked last_seq_scan/last_idx_scan inside a
transaction; it only previously appeared to work because pg_stat_force_next_flush()
previously did not flush in-transaction, and would now be unstable.
---
 doc/src/sgml/monitoring.sgml                  |   9 +-
 src/backend/utils/activity/pgstat.c           |  38 +-
 src/backend/utils/activity/pgstat_database.c  |  13 +-
 src/backend/utils/activity/pgstat_function.c  |  23 +-
 src/backend/utils/activity/pgstat_relation.c  | 132 ++++--
 .../utils/activity/pgstat_subscription.c      |  16 +-
 src/include/utils/pgstat_internal.h           |  53 ++-
 .../test_custom_stats/test_custom_var_stats.c |  20 +-
 src/test/regress/expected/stats.out           | 395 +++++++++++++++++-
 src/test/regress/sql/stats.sql                | 237 ++++++++++-
 10 files changed, 848 insertions(+), 88 deletions(-)

diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index f1422826d29..f2e4d21c52f 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -4538,7 +4538,8 @@ description | Waiting for a newly initialized WAL file to reach durable storage
       </para>
       <para>
        The time of the last sequential scan on this table, based on the
-       most recent transaction stop time
+       most recent transaction stop time, or the statement start time
+       when flushed mid-transaction
       </para></entry>
      </row>
 
@@ -4566,7 +4567,8 @@ description | Waiting for a newly initialized WAL file to reach durable storage
       </para>
       <para>
        The time of the last index scan on this table, based on the
-       most recent transaction stop time
+       most recent transaction stop time, or the statement start time
+       when flushed mid-transaction
       </para></entry>
      </row>
 
@@ -5067,7 +5069,8 @@ description | Waiting for a newly initialized WAL file to reach durable storage
       </para>
       <para>
        The time of the last scan on this index, based on the
-       most recent transaction stop time
+       most recent transaction stop time, or the statement start time
+       when flushed mid-transaction
       </para></entry>
      </row>
 
diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index 50cd07822b4..d18190b31ca 100644
--- a/src/backend/utils/activity/pgstat.c
+++ b/src/backend/utils/activity/pgstat.c
@@ -716,8 +716,8 @@ pgstat_initialize(void)
  * a timeout after which to call pgstat_report_stat(true), but are not
  * required to do so.
  *
- * Note that this is called only when not within a transaction, so it is fair
- * to use transaction stop time as an approximation of current time.
+ * A non-forced flush is only ever called outside of a transaction, while a
+ * forced flush may also happen within one (e.g. pg_stat_force_next_flush()).
  */
 long
 pgstat_report_stat(bool force)
@@ -729,7 +729,7 @@ pgstat_report_stat(bool force)
 	bool		nowait;
 
 	pgstat_assert_is_up();
-	Assert(!IsTransactionOrTransactionBlock());
+	Assert(force || !IsTransactionOrTransactionBlock());
 
 	/* "absorb" the forced flush even if there's nothing to flush */
 	if (pgStatForceNextFlush)
@@ -815,12 +815,13 @@ pgstat_report_stat(bool force)
 
 	/*
 	 * If some of the pending stats could not be flushed due to lock
-	 * contention, let the caller know when to retry.
+	 * contention, or only partially flushed due to in-transaction counters
+	 * being deferred, let the caller know when to retry.
 	 */
 	if (partial_flush)
 	{
-		/* force should have prevented us from getting here */
-		Assert(!force);
+		/* with force, only active transaction state can cause a partial flush */
+		Assert(!force || IsTransactionOrTransactionBlock());
 
 		/* remember since when stats have been pending */
 		if (pending_since == 0)
@@ -842,6 +843,9 @@ pgstat_report_stat(bool force)
 void
 pgstat_force_next_flush(void)
 {
+	if (IsTransactionOrTransactionBlock())
+		pgstat_report_stat(true);
+
 	pgStatForceNextFlush = true;
 }
 
@@ -1420,16 +1424,25 @@ pgstat_flush_pending_entries(bool nowait)
 		PgStat_HashKey key = entry_ref->shared_entry->key;
 		PgStat_Kind kind = key.kind;
 		const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
-		bool		did_flush;
+		PgStat_FlushResult result;
+		bool		xact_boundary = !IsTransactionOrTransactionBlock();
 		dlist_node *next;
 
 		Assert(!kind_info->fixed_amount);
 		Assert(kind_info->flush_pending_cb != NULL);
 
 		/* flush the stats, if possible */
-		did_flush = kind_info->flush_pending_cb(entry_ref, nowait);
+		result = kind_info->flush_pending_cb(entry_ref, nowait, xact_boundary);
 
-		Assert(did_flush || nowait);
+		/*
+		 * A lock conflict can only happen when we allowed the callback to
+		 * give up without waiting, and a partial flush can only happen when
+		 * we are inside a transaction and the callback had to retain
+		 * transactional state.
+		 */
+		Assert(result == PGSTAT_FLUSH_DONE ||
+			   (result == PGSTAT_FLUSH_LOCK_CONFLICT && nowait) ||
+			   (result == PGSTAT_FLUSH_PARTIAL && !xact_boundary));
 
 		/* determine next entry, before deleting the pending entry */
 		if (dlist_has_next(&pgStatPending, cur))
@@ -1437,8 +1450,11 @@ pgstat_flush_pending_entries(bool nowait)
 		else
 			next = NULL;
 
-		/* if successfully flushed, remove entry */
-		if (did_flush)
+		/*
+		 * Never free pending entries mid-transaction; callers may hold
+		 * pointers.
+		 */
+		if (result == PGSTAT_FLUSH_DONE && xact_boundary)
 			pgstat_delete_pending_entry(entry_ref);
 		else
 			have_pending = true;
diff --git a/src/backend/utils/activity/pgstat_database.c b/src/backend/utils/activity/pgstat_database.c
index 7f3bc016593..54b17f148d7 100644
--- a/src/backend/utils/activity/pgstat_database.c
+++ b/src/backend/utils/activity/pgstat_database.c
@@ -432,10 +432,13 @@ pgstat_reset_database_timestamp(Oid dboid, TimestampTz ts)
  * Flush out pending stats for the entry
  *
  * If nowait is true and the lock could not be immediately acquired, returns
- * false without flushing the entry.  Otherwise returns true.
+ * PGSTAT_FLUSH_LOCK_CONFLICT without flushing the entry.  Database stats are
+ * not transactional, so xact_boundary is unused and this always returns
+ * PGSTAT_FLUSH_DONE once flushed.
  */
-bool
-pgstat_database_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
+PgStat_FlushResult
+pgstat_database_flush_cb(PgStat_EntryRef *entry_ref, bool nowait,
+						 bool xact_boundary)
 {
 	PgStatShared_Database *sharedent;
 	PgStat_StatDBEntry *pendingent;
@@ -444,7 +447,7 @@ pgstat_database_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	sharedent = (PgStatShared_Database *) entry_ref->shared_stats;
 
 	if (!pgstat_lock_entry(entry_ref, nowait))
-		return false;
+		return PGSTAT_FLUSH_LOCK_CONFLICT;
 
 #define PGSTAT_ACCUM_DBCOUNT(item)		\
 	(sharedent)->stats.item += (pendingent)->item
@@ -496,7 +499,7 @@ pgstat_database_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 
 	memset(pendingent, 0, sizeof(*pendingent));
 
-	return true;
+	return PGSTAT_FLUSH_DONE;
 }
 
 void
diff --git a/src/backend/utils/activity/pgstat_function.c b/src/backend/utils/activity/pgstat_function.c
index f0366e13990..5ce6b1e9523 100644
--- a/src/backend/utils/activity/pgstat_function.c
+++ b/src/backend/utils/activity/pgstat_function.c
@@ -187,10 +187,16 @@ pgstat_end_function_usage(PgStat_FunctionCallUsage *fcu, bool finalize)
  * Flush out pending stats for the entry
  *
  * If nowait is true and the lock could not be immediately acquired, returns
- * false without flushing the entry.  Otherwise returns true.
+ * PGSTAT_FLUSH_LOCK_CONFLICT without flushing the entry.
+ *
+ * Function stats are not transactional, so this always flushes everything and
+ * returns PGSTAT_FLUSH_DONE.  The pending entry is zeroed after flushing to
+ * avoid double-counting if the caller keeps it alive (as happens
+ * mid-transaction).
  */
-bool
-pgstat_function_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
+PgStat_FlushResult
+pgstat_function_flush_cb(PgStat_EntryRef *entry_ref, bool nowait,
+						 bool xact_boundary)
 {
 	PgStat_FunctionCounts *localent;
 	PgStatShared_Function *shfuncent;
@@ -201,7 +207,7 @@ pgstat_function_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	/* localent always has non-zero content */
 
 	if (!pgstat_lock_entry(entry_ref, nowait))
-		return false;
+		return PGSTAT_FLUSH_LOCK_CONFLICT;
 
 	shfuncent->stats.numcalls += localent->numcalls;
 	shfuncent->stats.total_time +=
@@ -211,7 +217,14 @@ pgstat_function_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 
 	pgstat_unlock_entry(entry_ref);
 
-	return true;
+	/*
+	 * Zeroing total_time can cause double-counting if a recursive function
+	 * calls pg_stat_force_next_flush().  This isn't ideal, but not worth
+	 * adding complexity to handle that case.
+	 */
+	memset(localent, 0, sizeof(*localent));
+
+	return PGSTAT_FLUSH_DONE;
 }
 
 void
diff --git a/src/backend/utils/activity/pgstat_relation.c b/src/backend/utils/activity/pgstat_relation.c
index 04f2eb21d0b..7edea33346c 100644
--- a/src/backend/utils/activity/pgstat_relation.c
+++ b/src/backend/utils/activity/pgstat_relation.c
@@ -808,84 +808,120 @@ pgstat_twophase_postabort(FullTransactionId fxid, uint16 info,
  * Flush out pending stats for the entry
  *
  * If nowait is true and the lock could not be immediately acquired, returns
- * false without flushing the entry.  Otherwise returns true.
+ * PGSTAT_FLUSH_LOCK_CONFLICT without flushing the entry.
+ *
+ * Transactional counters (tuples_inserted/updated/deleted and the derived
+ * live/dead tuple counts) can only be flushed at a transaction boundary, since
+ * their final values depend on the transaction outcome.  When called
+ * mid-transaction (xact_boundary is false) for a relation that still has
+ * active transaction state, those counters are retained, only the
+ * non-transactional counters are flushed, and PGSTAT_FLUSH_PARTIAL is
+ * returned.  Otherwise everything is flushed and PGSTAT_FLUSH_DONE is returned.
  *
  * Some of the stats are copied to the corresponding pending database stats
  * entry when successfully flushing.
  */
-bool
-pgstat_relation_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
+PgStat_FlushResult
+pgstat_relation_flush_cb(PgStat_EntryRef *entry_ref, bool nowait,
+						 bool xact_boundary)
 {
 	Oid			dboid;
 	PgStat_TableStatus *lstats; /* pending stats entry  */
 	PgStatShared_Relation *shtabstats;
 	PgStat_StatTabEntry *tabentry;	/* table entry of shared stats */
 	PgStat_StatDBEntry *dbentry;	/* pending database entry */
+	bool		flush_txn;
 
 	dboid = entry_ref->shared_entry->key.dboid;
 	lstats = (PgStat_TableStatus *) entry_ref->pending;
 	shtabstats = (PgStatShared_Relation *) entry_ref->shared_stats;
 
+	/*
+	 * The transactional counters can be flushed once we reach a transaction
+	 * boundary, or when this relation has no active transaction state (i.e.
+	 * no pending DML whose outcome depends on commit/abort).
+	 */
+	flush_txn = (xact_boundary || lstats->trans == NULL);
+
 	/*
 	 * Ignore entries that didn't accumulate any actual counts, such as
-	 * indexes that were opened by the planner but not used.
+	 * indexes that were opened by the planner but not used.  With
+	 * in-transaction flushing an entry may be flushed multiple times, so keep
+	 * it pending if it has active transaction state and commit will merge
+	 * counters into it.
 	 */
 	if (pg_memory_is_all_zeros(&lstats->counts,
 							   sizeof(struct PgStat_TableCounts)))
-		return true;
+		return flush_txn ? PGSTAT_FLUSH_DONE : PGSTAT_FLUSH_PARTIAL;
 
 	if (!pgstat_lock_entry(entry_ref, nowait))
-		return false;
+		return PGSTAT_FLUSH_LOCK_CONFLICT;
 
-	/* add the values to the shared entry. */
+	/* Flush non-transactional counters. */
 	tabentry = &shtabstats->stats;
 
 	tabentry->numscans += lstats->counts.numscans;
 	if (lstats->counts.numscans)
 	{
-		TimestampTz t = GetCurrentTransactionStopTimestamp();
+		TimestampTz t = xact_boundary ?
+			GetCurrentTransactionStopTimestamp() :
+			GetCurrentStatementStartTimestamp();
 
 		if (t > tabentry->lastscan)
 			tabentry->lastscan = t;
 	}
 	tabentry->tuples_returned += lstats->counts.tuples_returned;
 	tabentry->tuples_fetched += lstats->counts.tuples_fetched;
-	tabentry->tuples_inserted += lstats->counts.tuples_inserted;
-	tabentry->tuples_updated += lstats->counts.tuples_updated;
-	tabentry->tuples_deleted += lstats->counts.tuples_deleted;
-	tabentry->tuples_hot_updated += lstats->counts.tuples_hot_updated;
-	tabentry->tuples_newpage_updated += lstats->counts.tuples_newpage_updated;
+	tabentry->blocks_fetched += lstats->counts.blocks_fetched;
+	tabentry->blocks_hit += lstats->counts.blocks_hit;
 
 	/*
-	 * If table was truncated/dropped, first reset the live/dead counters.
+	 * Flush transactional counters only when it is safe to do so.
+	 *
+	 * These counters are grouped because they have cross-counter invariants:
+	 * hot_updated + newpage_updated <= updated, ins_since_vacuum tracks
+	 * tuples_inserted, and live/dead/changed tuples are derived from
+	 * inserted/updated/deleted.  Flushing them separately could expose
+	 * inconsistent intermediate states (e.g. hot_updated > 0 with updated ==
+	 * 0 after an abort).
 	 */
-	if (lstats->counts.truncdropped)
+	if (flush_txn)
 	{
-		tabentry->live_tuples = 0;
-		tabentry->dead_tuples = 0;
-		tabentry->ins_since_vacuum = 0;
-	}
+		tabentry->tuples_inserted += lstats->counts.tuples_inserted;
+		tabentry->tuples_updated += lstats->counts.tuples_updated;
+		tabentry->tuples_deleted += lstats->counts.tuples_deleted;
+		tabentry->tuples_hot_updated += lstats->counts.tuples_hot_updated;
+		tabentry->tuples_newpage_updated += lstats->counts.tuples_newpage_updated;
 
-	tabentry->live_tuples += lstats->counts.delta_live_tuples;
-	tabentry->dead_tuples += lstats->counts.delta_dead_tuples;
-	tabentry->mod_since_analyze += lstats->counts.changed_tuples;
+		/*
+		 * If table was truncated/dropped, first reset the live/dead counters.
+		 */
+		if (lstats->counts.truncdropped)
+		{
+			tabentry->live_tuples = 0;
+			tabentry->dead_tuples = 0;
+			tabentry->ins_since_vacuum = 0;
+		}
 
-	/*
-	 * Using tuples_inserted to update ins_since_vacuum does mean that we'll
-	 * track aborted inserts too.  This isn't ideal, but otherwise probably
-	 * not worth adding an extra field for.  It may just amount to autovacuums
-	 * triggering for inserts more often than they maybe should, which is
-	 * probably not going to be common enough to be too concerned about here.
-	 */
-	tabentry->ins_since_vacuum += lstats->counts.tuples_inserted;
+		tabentry->live_tuples += lstats->counts.delta_live_tuples;
+		tabentry->dead_tuples += lstats->counts.delta_dead_tuples;
+		tabentry->mod_since_analyze += lstats->counts.changed_tuples;
 
-	tabentry->blocks_fetched += lstats->counts.blocks_fetched;
-	tabentry->blocks_hit += lstats->counts.blocks_hit;
+		/*
+		 * Using tuples_inserted to update ins_since_vacuum does mean that
+		 * we'll track aborted inserts too.  This isn't ideal, but otherwise
+		 * probably not worth adding an extra field for.  It may just amount
+		 * to autovacuums triggering for inserts more often than they maybe
+		 * should, which is probably not going to be common enough to be too
+		 * concerned about here.
+		 */
+		tabentry->ins_since_vacuum += lstats->counts.tuples_inserted;
 
-	/* Clamp live_tuples in case of negative delta_live_tuples */
-	tabentry->live_tuples = Max(tabentry->live_tuples, 0);
-	/* Likewise for dead_tuples */
-	tabentry->dead_tuples = Max(tabentry->dead_tuples, 0);
+		/* Clamp live_tuples in case of negative delta_live_tuples */
+		tabentry->live_tuples = Max(tabentry->live_tuples, 0);
+		/* Likewise for dead_tuples */
+		tabentry->dead_tuples = Max(tabentry->dead_tuples, 0);
+	}
 
 	pgstat_unlock_entry(entry_ref);
 
@@ -893,13 +929,29 @@ pgstat_relation_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	dbentry = pgstat_prep_database_pending(dboid);
 	dbentry->tuples_returned += lstats->counts.tuples_returned;
 	dbentry->tuples_fetched += lstats->counts.tuples_fetched;
-	dbentry->tuples_inserted += lstats->counts.tuples_inserted;
-	dbentry->tuples_updated += lstats->counts.tuples_updated;
-	dbentry->tuples_deleted += lstats->counts.tuples_deleted;
 	dbentry->blocks_fetched += lstats->counts.blocks_fetched;
 	dbentry->blocks_hit += lstats->counts.blocks_hit;
 
-	return true;
+	if (flush_txn)
+	{
+		dbentry->tuples_inserted += lstats->counts.tuples_inserted;
+		dbentry->tuples_updated += lstats->counts.tuples_updated;
+		dbentry->tuples_deleted += lstats->counts.tuples_deleted;
+		memset(&lstats->counts, 0, sizeof(lstats->counts));
+		return PGSTAT_FLUSH_DONE;
+	}
+
+	/*
+	 * This is a partial, in-transaction flush.  Zero out the counters we
+	 * already flushed so they aren't double-counted on the next flush.
+	 */
+	lstats->counts.numscans = 0;
+	lstats->counts.tuples_returned = 0;
+	lstats->counts.tuples_fetched = 0;
+	lstats->counts.blocks_fetched = 0;
+	lstats->counts.blocks_hit = 0;
+
+	return PGSTAT_FLUSH_PARTIAL;
 }
 
 void
diff --git a/src/backend/utils/activity/pgstat_subscription.c b/src/backend/utils/activity/pgstat_subscription.c
index 3eaf3e0390f..a79120d899a 100644
--- a/src/backend/utils/activity/pgstat_subscription.c
+++ b/src/backend/utils/activity/pgstat_subscription.c
@@ -114,10 +114,13 @@ pgstat_fetch_stat_subscription(Oid subid)
  * Flush out pending stats for the entry
  *
  * If nowait is true and the lock could not be immediately acquired, returns
- * false without flushing the entry.  Otherwise returns true.
+ * PGSTAT_FLUSH_LOCK_CONFLICT without flushing the entry.  Subscription stats
+ * are not transactional, so this always flushes everything and returns
+ * PGSTAT_FLUSH_DONE.
  */
-bool
-pgstat_subscription_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
+PgStat_FlushResult
+pgstat_subscription_flush_cb(PgStat_EntryRef *entry_ref, bool nowait,
+							 bool xact_boundary)
 {
 	PgStat_BackendSubEntry *localent;
 	PgStatShared_Subscription *shsubent;
@@ -128,7 +131,7 @@ pgstat_subscription_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	/* localent always has non-zero content */
 
 	if (!pgstat_lock_entry(entry_ref, nowait))
-		return false;
+		return PGSTAT_FLUSH_LOCK_CONFLICT;
 
 #define SUB_ACC(fld) shsubent->stats.fld += localent->fld
 	SUB_ACC(apply_error_count);
@@ -139,7 +142,10 @@ pgstat_subscription_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 #undef SUB_ACC
 
 	pgstat_unlock_entry(entry_ref);
-	return true;
+
+	memset(localent, 0, sizeof(*localent));
+
+	return PGSTAT_FLUSH_DONE;
 }
 
 void
diff --git a/src/include/utils/pgstat_internal.h b/src/include/utils/pgstat_internal.h
index b0a17691966..e8739efe41b 100644
--- a/src/include/utils/pgstat_internal.h
+++ b/src/include/utils/pgstat_internal.h
@@ -225,6 +225,34 @@ typedef struct PgStat_SubXactStatus
 } PgStat_SubXactStatus;
 
 
+/*
+ * Result of a flush_pending_cb call, describing what happened to the pending
+ * entry.  The caller combines this with the flush context to decide whether the
+ * entry can be removed from the pending list.
+ */
+typedef enum PgStat_FlushResult
+{
+	/*
+	 * The lock could not be acquired without waiting (nowait was true).  The
+	 * entry must stay pending and be retried later.
+	 *
+	 * Must be first; in previous versions flush_pending_cb returned a bool
+	 * where false meant lock conflict.
+	 */
+	PGSTAT_FLUSH_LOCK_CONFLICT = 0,
+
+	/* Fully flushed; the entry can be removed from the pending list. */
+	PGSTAT_FLUSH_DONE,
+
+	/*
+	 * Only part of the entry was flushed; some state was intentionally
+	 * retained because it cannot be flushed in the current context (e.g.
+	 * transactional counters flushed mid-transaction).  The entry must stay
+	 * pending and be flushed again at a suitable boundary.
+	 */
+	PGSTAT_FLUSH_PARTIAL,
+}			PgStat_FlushResult;
+
 /*
  * Metadata for a specific kind of statistics.
  */
@@ -297,8 +325,19 @@ typedef struct PgStat_KindInfo
 	 * For variable-numbered stats: flush pending stats. Required if pending
 	 * data is used. See flush_static_cb when dealing with stats data that
 	 * that cannot use PgStat_EntryRef->pending.
+	 *
+	 * If nowait is true and the entry's lock cannot be acquired without
+	 * waiting, the callback must return PGSTAT_FLUSH_LOCK_CONFLICT without
+	 * flushing.  xact_boundary tells the callback whether it is safe to flush
+	 * transactional state: it is true when called outside of a transaction
+	 * (e.g. at transaction end), and false when called mid-transaction.  When
+	 * xact_boundary is false, a callback holding state whose outcome depends
+	 * on the transaction must flush only its non-transactional counters,
+	 * retain the rest, and return PGSTAT_FLUSH_PARTIAL.  Otherwise it flushes
+	 * everything and returns PGSTAT_FLUSH_DONE.
 	 */
-	bool		(*flush_pending_cb) (PgStat_EntryRef *sr, bool nowait);
+				PgStat_FlushResult(*flush_pending_cb) (PgStat_EntryRef *sr, bool nowait,
+													   bool xact_boundary);
 
 	/*
 	 * For variable-numbered stats: delete pending stats. Optional.
@@ -743,7 +782,8 @@ extern void AtEOXact_PgStat_Database(bool isCommit, bool parallel);
 
 extern PgStat_StatDBEntry *pgstat_prep_database_pending(Oid dboid);
 extern void pgstat_reset_database_timestamp(Oid dboid, TimestampTz ts);
-extern bool pgstat_database_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);
+extern PgStat_FlushResult pgstat_database_flush_cb(PgStat_EntryRef *entry_ref,
+												   bool nowait, bool xact_boundary);
 extern void pgstat_database_reset_timestamp_cb(PgStatShared_Common *header, TimestampTz ts);
 
 
@@ -751,7 +791,8 @@ extern void pgstat_database_reset_timestamp_cb(PgStatShared_Common *header, Time
  * Functions in pgstat_function.c
  */
 
-extern bool pgstat_function_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);
+extern PgStat_FlushResult pgstat_function_flush_cb(PgStat_EntryRef *entry_ref,
+												   bool nowait, bool xact_boundary);
 extern void pgstat_function_reset_timestamp_cb(PgStatShared_Common *header, TimestampTz ts);
 
 
@@ -784,7 +825,8 @@ extern void AtEOSubXact_PgStat_Relations(PgStat_SubXactStatus *xact_state, bool
 extern void AtPrepare_PgStat_Relations(PgStat_SubXactStatus *xact_state);
 extern void PostPrepare_PgStat_Relations(PgStat_SubXactStatus *xact_state);
 
-extern bool pgstat_relation_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);
+extern PgStat_FlushResult pgstat_relation_flush_cb(PgStat_EntryRef *entry_ref,
+												   bool nowait, bool xact_boundary);
 extern void pgstat_relation_delete_pending_cb(PgStat_EntryRef *entry_ref);
 extern void pgstat_relation_reset_timestamp_cb(PgStatShared_Common *header, TimestampTz ts);
 
@@ -853,7 +895,8 @@ extern void pgstat_wal_snapshot_cb(void);
  * Functions in pgstat_subscription.c
  */
 
-extern bool pgstat_subscription_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);
+extern PgStat_FlushResult pgstat_subscription_flush_cb(PgStat_EntryRef *entry_ref,
+													   bool nowait, bool xact_boundary);
 extern void pgstat_subscription_reset_timestamp_cb(PgStatShared_Common *header, TimestampTz ts);
 
 
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 a39ada0b67c..9a29068a420 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
@@ -90,8 +90,9 @@ static dsa_area *custom_stats_description_dsa = NULL;
  */
 
 /* Flush callback: merge pending stats into shared memory */
-static bool test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref,
-												   bool nowait);
+static PgStat_FlushResult test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref,
+																 bool nowait,
+																 bool xact_boundary);
 
 /* Serialization callback: write auxiliary entry data */
 static bool test_custom_stats_var_to_serialized_data(const PgStat_HashKey *key,
@@ -151,10 +152,13 @@ _PG_init(void)
  * Called by pgstat collector to flush accumulated local statistics
  * to shared memory where other backends can read them.
  *
- * Returns false only if nowait=true and lock acquisition fails.
+ * These stats are not transactional, so xact_boundary is unused.  Returns
+ * PGSTAT_FLUSH_LOCK_CONFLICT only if nowait=true and lock acquisition fails,
+ * otherwise PGSTAT_FLUSH_DONE.
  */
-static bool
-test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref, bool nowait)
+static PgStat_FlushResult
+test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref, bool nowait,
+									   bool xact_boundary)
 {
 	PgStat_StatCustomVarEntry *pending_entry;
 	PgStatShared_CustomVarEntry *shared_entry;
@@ -163,14 +167,16 @@ test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref, bool nowait)
 	shared_entry = (PgStatShared_CustomVarEntry *) entry_ref->shared_stats;
 
 	if (!pgstat_lock_entry(entry_ref, nowait))
-		return false;
+		return PGSTAT_FLUSH_LOCK_CONFLICT;
 
 	/* Add pending counts to shared totals */
 	shared_entry->stats.numcalls += pending_entry->numcalls;
 
 	pgstat_unlock_entry(entry_ref);
 
-	return true;
+	memset(pending_entry, 0, sizeof(*pending_entry));
+
+	return PGSTAT_FLUSH_DONE;
 }
 
 /*
diff --git a/src/test/regress/expected/stats.out b/src/test/regress/expected/stats.out
index e230356de13..4b735625775 100644
--- a/src/test/regress/expected/stats.out
+++ b/src/test/regress/expected/stats.out
@@ -691,12 +691,6 @@ SELECT pg_stat_force_next_flush();
  
 (1 row)
 
-SELECT last_seq_scan, last_idx_scan FROM pg_stat_all_tables WHERE relid = 'test_last_scan'::regclass;
- last_seq_scan | last_idx_scan 
----------------+---------------
-               | 
-(1 row)
-
 COMMIT;
 SELECT stats_reset IS NOT NULL AS has_stats_reset
   FROM pg_stat_all_tables WHERE relid = 'test_last_scan'::regclass;
@@ -2068,4 +2062,393 @@ SELECT fastpath_exceeded > :backend_fastpath_exceeded_before
 (1 row)
 
 DROP TABLE part_test;
+--
+-- Test in-transaction flushes
+--
+CREATE TABLE partial_flush(id int);
+INSERT INTO partial_flush VALUES (1), (2), (3);
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+-- Record counters before the explicit transaction
+SELECT seq_scan AS seq_scan_before,
+       seq_tup_read AS seq_tup_read_before,
+       n_tup_ins AS n_tup_ins_before,
+       n_tup_upd AS n_tup_upd_before
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush' \gset
+BEGIN;
+SET LOCAL stats_fetch_consistency = none;
+-- Generate both transaction-safe and transaction-unsafe counters.
+SELECT count(*) FROM partial_flush;
+ count 
+-------
+     3
+(1 row)
+
+INSERT INTO partial_flush VALUES (4), (5);
+UPDATE partial_flush SET id = id WHERE id = 1;
+-- Flush in-transaction
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+-- Transaction-safe counters should be visible in-transaction.
+-- Transaction-unsafe counters (ins, upd) should NOT be flushed yet,
+-- since their final values depend on whether the transaction commits.
+-- Also verify that hot/newpage updates (which are subsets of n_tup_upd)
+-- were not flushed ahead of n_tup_upd, and tuple counts stay non-negative.
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_tup_ins - :n_tup_ins_before AS n_tup_ins_delta,
+       n_tup_upd - :n_tup_upd_before AS n_tup_upd_delta,
+       n_tup_hot_upd + n_tup_newpage_upd <= n_tup_upd AS upd_counts_ok
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush';
+ seq_scan_delta | seq_tup_read_delta | n_tup_ins_delta | n_tup_upd_delta | upd_counts_ok 
+----------------+--------------------+-----------------+-----------------+---------------
+              2 |                  8 |               0 |               0 | t
+(1 row)
+
+-- Generate more transaction-safe activity to verify no double counting.
+SELECT count(*) FROM partial_flush;
+ count 
+-------
+     5
+(1 row)
+
+-- Flush again in-transaction
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+-- Should show cumulative totals, not double-counted.
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_tup_ins - :n_tup_ins_before AS n_tup_ins_delta,
+       n_tup_upd - :n_tup_upd_before AS n_tup_upd_delta
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush';
+ seq_scan_delta | seq_tup_read_delta | n_tup_ins_delta | n_tup_upd_delta 
+----------------+--------------------+-----------------+-----------------
+              3 |                 13 |               0 |               0
+(1 row)
+
+COMMIT;
+-- After commit, all counters should be flushed.
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_tup_ins - :n_tup_ins_before AS n_tup_ins_delta,
+       n_tup_upd - :n_tup_upd_before AS n_tup_upd_delta,
+       n_tup_hot_upd + n_tup_newpage_upd <= n_tup_upd AS upd_counts_ok
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush';
+ seq_scan_delta | seq_tup_read_delta | n_tup_ins_delta | n_tup_upd_delta | upd_counts_ok 
+----------------+--------------------+-----------------+-----------------+---------------
+              3 |                 13 |               2 |               1 | t
+(1 row)
+
+DROP TABLE partial_flush;
+--
+-- Test in-transaction partial flush: rollback discards transactional counters
+-- but non-transactional counters already flushed to shared stats persist.
+--
+CREATE TABLE partial_flush_rollback(id int);
+INSERT INTO partial_flush_rollback SELECT generate_series(1, 50);
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT seq_scan AS seq_scan_before,
+       seq_tup_read AS seq_tup_read_before,
+       n_tup_ins AS n_tup_ins_before,
+       n_live_tup AS n_live_tup_before
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_rollback' \gset
+BEGIN;
+SET LOCAL stats_fetch_consistency = none;
+-- Generate both non-transactional (scan) and transactional (insert) activity.
+SELECT count(*) FROM partial_flush_rollback;
+ count 
+-------
+    50
+(1 row)
+
+INSERT INTO partial_flush_rollback SELECT generate_series(51, 100);
+-- Flush in-transaction: scans go to shared stats, inserts stay pending.
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+-- Verify: scans visible, inserts and live_tup deferred.  Also check that
+-- partial flush did not expose hot/newpage without their parent update count.
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_tup_ins - :n_tup_ins_before AS n_tup_ins_delta,
+       n_live_tup - :n_live_tup_before AS n_live_tup_delta,
+       n_tup_hot_upd + n_tup_newpage_upd <= n_tup_upd AS upd_counts_ok
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_rollback';
+ seq_scan_delta | seq_tup_read_delta | n_tup_ins_delta | n_live_tup_delta | upd_counts_ok 
+----------------+--------------------+-----------------+------------------+---------------
+              1 |                 50 |               0 |                0 | t
+(1 row)
+
+ROLLBACK;
+-- After rollback: scans persist (already in shared stats), inserts discarded,
+-- n_live_tup unchanged, and hot/newpage remain consistent with n_tup_upd.
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_live_tup - :n_live_tup_before AS n_live_tup_delta,
+       n_tup_hot_upd + n_tup_newpage_upd <= n_tup_upd AS upd_counts_ok
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_rollback';
+ seq_scan_delta | seq_tup_read_delta | n_live_tup_delta | upd_counts_ok 
+----------------+--------------------+------------------+---------------
+              1 |                 50 |                0 | t
+(1 row)
+
+DROP TABLE partial_flush_rollback;
+--
+-- Test in-transaction partial flush with TRUNCATE: the truncate's reset of
+-- live/dead counters is transactional and must not reach shared stats until
+-- commit.
+--
+CREATE TABLE partial_flush_truncate(id int);
+INSERT INTO partial_flush_truncate SELECT generate_series(1, 100);
+DELETE FROM partial_flush_truncate WHERE id <= 20;
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT n_live_tup AS n_live_tup_before,
+       n_dead_tup AS n_dead_tup_before,
+       seq_scan AS seq_scan_before
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate' \gset
+-- Case 1: DML + TRUNCATE + DML + ROLLBACK.  The truncate's zeroing and all
+-- transactional counters must not leak to shared stats.
+BEGIN;
+SET LOCAL stats_fetch_consistency = none;
+-- DML before truncate.
+SELECT count(*) FROM partial_flush_truncate;
+ count 
+-------
+    80
+(1 row)
+
+INSERT INTO partial_flush_truncate SELECT generate_series(101, 110);
+UPDATE partial_flush_truncate SET id = id WHERE id = 1;
+TRUNCATE partial_flush_truncate;
+-- DML after truncate.
+INSERT INTO partial_flush_truncate SELECT generate_series(1, 10);
+-- Mid-transaction flush: scan goes to shared stats, everything else deferred.
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       n_live_tup = :n_live_tup_before AS live_tup_unchanged,
+       n_dead_tup = :n_dead_tup_before AS dead_tup_unchanged
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate';
+ seq_scan_delta | live_tup_unchanged | dead_tup_unchanged 
+----------------+--------------------+--------------------
+              2 | t                  | t
+(1 row)
+
+ROLLBACK;
+-- After rollback: live_tup unchanged (rollback undid the inserts before they
+-- became live), but dead_tup increases because aborted inserts produce dead
+-- tuples that vacuum must clean up.
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT n_live_tup = :n_live_tup_before AS live_tup_unchanged,
+       n_dead_tup - :n_dead_tup_before AS dead_tup_delta
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate';
+ live_tup_unchanged | dead_tup_delta 
+--------------------+----------------
+ t                  |             10
+(1 row)
+
+-- Case 2: DML + TRUNCATE + INSERT + DELETE + COMMIT.
+-- Update baseline to account for changes from case 1.
+SELECT seq_scan AS seq_scan_before,
+       n_live_tup AS n_live_tup_before,
+       n_dead_tup AS n_dead_tup_before
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate' \gset
+BEGIN;
+SET LOCAL stats_fetch_consistency = none;
+-- DML before truncate.
+SELECT count(*) FROM partial_flush_truncate;
+ count 
+-------
+    80
+(1 row)
+
+INSERT INTO partial_flush_truncate SELECT generate_series(101, 110);
+UPDATE partial_flush_truncate SET id = id WHERE id = 1;
+TRUNCATE partial_flush_truncate;
+-- DML after truncate.
+INSERT INTO partial_flush_truncate SELECT generate_series(1, 10);
+DELETE FROM partial_flush_truncate WHERE id <= 3;
+-- Mid-transaction flush: scans visible, transactional counters deferred.
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       n_live_tup = :n_live_tup_before AS live_tup_unchanged,
+       n_dead_tup = :n_dead_tup_before AS dead_tup_unchanged
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate';
+ seq_scan_delta | live_tup_unchanged | dead_tup_unchanged 
+----------------+--------------------+--------------------
+              3 | t                  | t
+(1 row)
+
+COMMIT;
+-- After commit: TRUNCATE zeros live/dead, then only post-truncate DML counts.
+-- delta_live = inserted - deleted = 10 - 3 = 7
+-- delta_dead = updated + deleted = 0 + 3 = 3
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT n_live_tup, n_dead_tup
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate';
+ n_live_tup | n_dead_tup 
+------------+------------
+          7 |          3
+(1 row)
+
+DROP TABLE partial_flush_truncate;
+--
+-- Test that pg_stat_force_next_flush() called inside a function does not
+-- lose or double-count function call statistics.
+--
+SET track_functions TO 'all';
+CREATE FUNCTION flush_func_test() RETURNS VOID LANGUAGE plpgsql AS $$
+BEGIN
+    PERFORM pg_stat_force_next_flush();
+END;
+$$;
+SELECT 'flush_func_test()'::regprocedure::oid AS flush_func_test_oid \gset
+-- Call outside a transaction.
+SELECT flush_func_test();
+ flush_func_test 
+-----------------
+ 
+(1 row)
+
+SELECT flush_func_test();
+ flush_func_test 
+-----------------
+ 
+(1 row)
+
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT funcname, calls FROM pg_stat_user_functions WHERE funcid = :flush_func_test_oid;
+    funcname     | calls 
+-----------------+-------
+ flush_func_test |     2
+(1 row)
+
+-- Call inside an explicit transaction.
+BEGIN;
+SELECT flush_func_test();
+ flush_func_test 
+-----------------
+ 
+(1 row)
+
+SELECT flush_func_test();
+ flush_func_test 
+-----------------
+ 
+(1 row)
+
+SELECT flush_func_test();
+ flush_func_test 
+-----------------
+ 
+(1 row)
+
+COMMIT;
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT funcname, calls FROM pg_stat_user_functions WHERE funcid = :flush_func_test_oid;
+    funcname     | calls 
+-----------------+-------
+ flush_func_test |     5
+(1 row)
+
+-- Call with an explicit mid-transaction flush between function calls.
+BEGIN;
+SELECT flush_func_test();
+ flush_func_test 
+-----------------
+ 
+(1 row)
+
+SELECT flush_func_test();
+ flush_func_test 
+-----------------
+ 
+(1 row)
+
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT flush_func_test();
+ flush_func_test 
+-----------------
+ 
+(1 row)
+
+COMMIT;
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush 
+--------------------------
+ 
+(1 row)
+
+SELECT funcname, calls FROM pg_stat_user_functions WHERE funcid = :flush_func_test_oid;
+    funcname     | calls 
+-----------------+-------
+ flush_func_test |     8
+(1 row)
+
+DROP FUNCTION flush_func_test;
 -- End of Stats Test
diff --git a/src/test/regress/sql/stats.sql b/src/test/regress/sql/stats.sql
index 4c265d1245c..a65e2778297 100644
--- a/src/test/regress/sql/stats.sql
+++ b/src/test/regress/sql/stats.sql
@@ -316,7 +316,6 @@ BEGIN;
 CREATE TEMPORARY TABLE test_last_scan(idx_col int primary key, noidx_col int);
 INSERT INTO test_last_scan(idx_col, noidx_col) VALUES(1, 1);
 SELECT pg_stat_force_next_flush();
-SELECT last_seq_scan, last_idx_scan FROM pg_stat_all_tables WHERE relid = 'test_last_scan'::regclass;
 COMMIT;
 
 SELECT stats_reset IS NOT NULL AS has_stats_reset
@@ -1024,4 +1023,240 @@ SELECT fastpath_exceeded > :backend_fastpath_exceeded_before
 
 DROP TABLE part_test;
 
+--
+-- Test in-transaction flushes
+--
+CREATE TABLE partial_flush(id int);
+INSERT INTO partial_flush VALUES (1), (2), (3);
+SELECT pg_stat_force_next_flush();
+
+-- Record counters before the explicit transaction
+SELECT seq_scan AS seq_scan_before,
+       seq_tup_read AS seq_tup_read_before,
+       n_tup_ins AS n_tup_ins_before,
+       n_tup_upd AS n_tup_upd_before
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush' \gset
+
+BEGIN;
+SET LOCAL stats_fetch_consistency = none;
+
+-- Generate both transaction-safe and transaction-unsafe counters.
+SELECT count(*) FROM partial_flush;
+INSERT INTO partial_flush VALUES (4), (5);
+UPDATE partial_flush SET id = id WHERE id = 1;
+
+-- Flush in-transaction
+SELECT pg_stat_force_next_flush();
+
+-- Transaction-safe counters should be visible in-transaction.
+-- Transaction-unsafe counters (ins, upd) should NOT be flushed yet,
+-- since their final values depend on whether the transaction commits.
+-- Also verify that hot/newpage updates (which are subsets of n_tup_upd)
+-- were not flushed ahead of n_tup_upd, and tuple counts stay non-negative.
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_tup_ins - :n_tup_ins_before AS n_tup_ins_delta,
+       n_tup_upd - :n_tup_upd_before AS n_tup_upd_delta,
+       n_tup_hot_upd + n_tup_newpage_upd <= n_tup_upd AS upd_counts_ok
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush';
+
+-- Generate more transaction-safe activity to verify no double counting.
+SELECT count(*) FROM partial_flush;
+
+-- Flush again in-transaction
+SELECT pg_stat_force_next_flush();
+
+-- Should show cumulative totals, not double-counted.
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_tup_ins - :n_tup_ins_before AS n_tup_ins_delta,
+       n_tup_upd - :n_tup_upd_before AS n_tup_upd_delta
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush';
+
+COMMIT;
+
+-- After commit, all counters should be flushed.
+
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_tup_ins - :n_tup_ins_before AS n_tup_ins_delta,
+       n_tup_upd - :n_tup_upd_before AS n_tup_upd_delta,
+       n_tup_hot_upd + n_tup_newpage_upd <= n_tup_upd AS upd_counts_ok
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush';
+
+DROP TABLE partial_flush;
+
+--
+-- Test in-transaction partial flush: rollback discards transactional counters
+-- but non-transactional counters already flushed to shared stats persist.
+--
+CREATE TABLE partial_flush_rollback(id int);
+INSERT INTO partial_flush_rollback SELECT generate_series(1, 50);
+SELECT pg_stat_force_next_flush();
+
+SELECT seq_scan AS seq_scan_before,
+       seq_tup_read AS seq_tup_read_before,
+       n_tup_ins AS n_tup_ins_before,
+       n_live_tup AS n_live_tup_before
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_rollback' \gset
+
+BEGIN;
+SET LOCAL stats_fetch_consistency = none;
+
+-- Generate both non-transactional (scan) and transactional (insert) activity.
+SELECT count(*) FROM partial_flush_rollback;
+INSERT INTO partial_flush_rollback SELECT generate_series(51, 100);
+
+-- Flush in-transaction: scans go to shared stats, inserts stay pending.
+SELECT pg_stat_force_next_flush();
+
+-- Verify: scans visible, inserts and live_tup deferred.  Also check that
+-- partial flush did not expose hot/newpage without their parent update count.
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_tup_ins - :n_tup_ins_before AS n_tup_ins_delta,
+       n_live_tup - :n_live_tup_before AS n_live_tup_delta,
+       n_tup_hot_upd + n_tup_newpage_upd <= n_tup_upd AS upd_counts_ok
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_rollback';
+
+ROLLBACK;
+
+-- After rollback: scans persist (already in shared stats), inserts discarded,
+-- n_live_tup unchanged, and hot/newpage remain consistent with n_tup_upd.
+SELECT pg_stat_force_next_flush();
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       seq_tup_read - :seq_tup_read_before AS seq_tup_read_delta,
+       n_live_tup - :n_live_tup_before AS n_live_tup_delta,
+       n_tup_hot_upd + n_tup_newpage_upd <= n_tup_upd AS upd_counts_ok
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_rollback';
+
+DROP TABLE partial_flush_rollback;
+
+--
+-- Test in-transaction partial flush with TRUNCATE: the truncate's reset of
+-- live/dead counters is transactional and must not reach shared stats until
+-- commit.
+--
+CREATE TABLE partial_flush_truncate(id int);
+INSERT INTO partial_flush_truncate SELECT generate_series(1, 100);
+DELETE FROM partial_flush_truncate WHERE id <= 20;
+SELECT pg_stat_force_next_flush();
+
+SELECT n_live_tup AS n_live_tup_before,
+       n_dead_tup AS n_dead_tup_before,
+       seq_scan AS seq_scan_before
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate' \gset
+
+-- Case 1: DML + TRUNCATE + DML + ROLLBACK.  The truncate's zeroing and all
+-- transactional counters must not leak to shared stats.
+BEGIN;
+SET LOCAL stats_fetch_consistency = none;
+
+-- DML before truncate.
+SELECT count(*) FROM partial_flush_truncate;
+INSERT INTO partial_flush_truncate SELECT generate_series(101, 110);
+UPDATE partial_flush_truncate SET id = id WHERE id = 1;
+
+TRUNCATE partial_flush_truncate;
+
+-- DML after truncate.
+INSERT INTO partial_flush_truncate SELECT generate_series(1, 10);
+
+-- Mid-transaction flush: scan goes to shared stats, everything else deferred.
+SELECT pg_stat_force_next_flush();
+
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       n_live_tup = :n_live_tup_before AS live_tup_unchanged,
+       n_dead_tup = :n_dead_tup_before AS dead_tup_unchanged
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate';
+
+ROLLBACK;
+
+-- After rollback: live_tup unchanged (rollback undid the inserts before they
+-- became live), but dead_tup increases because aborted inserts produce dead
+-- tuples that vacuum must clean up.
+SELECT pg_stat_force_next_flush();
+SELECT n_live_tup = :n_live_tup_before AS live_tup_unchanged,
+       n_dead_tup - :n_dead_tup_before AS dead_tup_delta
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate';
+
+-- Case 2: DML + TRUNCATE + INSERT + DELETE + COMMIT.
+-- Update baseline to account for changes from case 1.
+SELECT seq_scan AS seq_scan_before,
+       n_live_tup AS n_live_tup_before,
+       n_dead_tup AS n_dead_tup_before
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate' \gset
+
+BEGIN;
+SET LOCAL stats_fetch_consistency = none;
+
+-- DML before truncate.
+SELECT count(*) FROM partial_flush_truncate;
+INSERT INTO partial_flush_truncate SELECT generate_series(101, 110);
+UPDATE partial_flush_truncate SET id = id WHERE id = 1;
+
+TRUNCATE partial_flush_truncate;
+
+-- DML after truncate.
+INSERT INTO partial_flush_truncate SELECT generate_series(1, 10);
+DELETE FROM partial_flush_truncate WHERE id <= 3;
+
+-- Mid-transaction flush: scans visible, transactional counters deferred.
+SELECT pg_stat_force_next_flush();
+
+SELECT seq_scan - :seq_scan_before AS seq_scan_delta,
+       n_live_tup = :n_live_tup_before AS live_tup_unchanged,
+       n_dead_tup = :n_dead_tup_before AS dead_tup_unchanged
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate';
+
+COMMIT;
+
+-- After commit: TRUNCATE zeros live/dead, then only post-truncate DML counts.
+-- delta_live = inserted - deleted = 10 - 3 = 7
+-- delta_dead = updated + deleted = 0 + 3 = 3
+SELECT pg_stat_force_next_flush();
+SELECT n_live_tup, n_dead_tup
+  FROM pg_stat_user_tables WHERE relname = 'partial_flush_truncate';
+
+DROP TABLE partial_flush_truncate;
+
+--
+-- Test that pg_stat_force_next_flush() called inside a function does not
+-- lose or double-count function call statistics.
+--
+SET track_functions TO 'all';
+CREATE FUNCTION flush_func_test() RETURNS VOID LANGUAGE plpgsql AS $$
+BEGIN
+    PERFORM pg_stat_force_next_flush();
+END;
+$$;
+SELECT 'flush_func_test()'::regprocedure::oid AS flush_func_test_oid \gset
+
+-- Call outside a transaction.
+SELECT flush_func_test();
+SELECT flush_func_test();
+SELECT pg_stat_force_next_flush();
+SELECT funcname, calls FROM pg_stat_user_functions WHERE funcid = :flush_func_test_oid;
+
+-- Call inside an explicit transaction.
+BEGIN;
+SELECT flush_func_test();
+SELECT flush_func_test();
+SELECT flush_func_test();
+COMMIT;
+SELECT pg_stat_force_next_flush();
+SELECT funcname, calls FROM pg_stat_user_functions WHERE funcid = :flush_func_test_oid;
+
+-- Call with an explicit mid-transaction flush between function calls.
+BEGIN;
+SELECT flush_func_test();
+SELECT flush_func_test();
+SELECT pg_stat_force_next_flush();
+SELECT flush_func_test();
+COMMIT;
+SELECT pg_stat_force_next_flush();
+SELECT funcname, calls FROM pg_stat_user_functions WHERE funcid = :flush_func_test_oid;
+
+DROP FUNCTION flush_func_test;
+
 -- End of Stats Test
-- 
2.47.3

