pgstat_reject_inconsistent_consistency_mode_change.diff

diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index 60fc4e761f..5568b9d6dc 100644
--- a/src/backend/utils/activity/pgstat.c
+++ b/src/backend/utils/activity/pgstat.c
@@ -172,8 +172,8 @@ static void pgstat_reset_after_failure(void);
 
 static bool pgstat_flush_pending_entries(bool nowait);
 
-static void pgstat_prep_snapshot(void);
-static void pgstat_build_snapshot(void);
+static PgStat_FetchConsistency pgstat_prep_snapshot(void);
+static void pgstat_build_snapshot(PgStat_FetchConsistency fetch_consistency);
 static void pgstat_build_snapshot_fixed(PgStat_Kind kind);
 
 static inline bool pgstat_is_kind_valid(int ikind);
@@ -796,23 +796,25 @@ pgstat_fetch_entry(PgStat_Kind kind, Oid dboid, Oid objoid)
 	PgStat_EntryRef *entry_ref;
 	void	   *stats_data;
 	const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+	PgStat_FetchConsistency fetch_consistency;
 
 	/* should be called from backends */
 	Assert(IsUnderPostmaster || !IsPostmasterEnvironment);
 	Assert(!kind_info->fixed_amount);
 
-	pgstat_prep_snapshot();
-
 	key.kind = kind;
 	key.dboid = dboid;
 	key.objoid = objoid;
 
 	/* if we need to build a full snapshot, do so */
-	if (pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_SNAPSHOT)
-		pgstat_build_snapshot();
+	fetch_consistency = pgstat_prep_snapshot();
+	/* if we need to build a full snapshot, do so */
+	if (fetch_consistency == PGSTAT_FETCH_CONSISTENCY_SNAPSHOT)
+		pgstat_build_snapshot(fetch_consistency);
+
 
 	/* if caching is desired, look up in cache */
-	if (pgstat_fetch_consistency > PGSTAT_FETCH_CONSISTENCY_NONE)
+	if (fetch_consistency > PGSTAT_FETCH_CONSISTENCY_NONE)
 	{
 		PgStat_SnapshotEntry *entry = NULL;
 
@@ -825,18 +827,16 @@ pgstat_fetch_entry(PgStat_Kind kind, Oid dboid, Oid objoid)
 		 * If we built a full snapshot and the key is not in
 		 * pgStatLocal.snapshot.stats, there are no matching stats.
 		 */
-		if (pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_SNAPSHOT)
+		if (fetch_consistency == PGSTAT_FETCH_CONSISTENCY_SNAPSHOT)
 			return NULL;
 	}
 
-	pgStatLocal.snapshot.mode = pgstat_fetch_consistency;
-
 	entry_ref = pgstat_get_entry_ref(kind, dboid, objoid, false, NULL);
 
 	if (entry_ref == NULL || entry_ref->shared_entry->dropped)
 	{
 		/* create empty entry when using PGSTAT_FETCH_CONSISTENCY_CACHE */
-		if (pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_CACHE)
+		if (fetch_consistency == PGSTAT_FETCH_CONSISTENCY_CACHE)
 		{
 			PgStat_SnapshotEntry *entry = NULL;
 			bool		found;
@@ -853,7 +853,7 @@ pgstat_fetch_entry(PgStat_Kind kind, Oid dboid, Oid objoid)
 	 * otherwise we could quickly end up with a fair bit of memory used due to
 	 * repeated accesses.
 	 */
-	if (pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_NONE)
+	if (fetch_consistency == PGSTAT_FETCH_CONSISTENCY_NONE)
 		stats_data = palloc(kind_info->shared_data_len);
 	else
 		stats_data = MemoryContextAlloc(pgStatLocal.snapshot.context,
@@ -865,7 +865,7 @@ pgstat_fetch_entry(PgStat_Kind kind, Oid dboid, Oid objoid)
 		   kind_info->shared_data_len);
 	pgstat_unlock_entry(entry_ref);
 
-	if (pgstat_fetch_consistency > PGSTAT_FETCH_CONSISTENCY_NONE)
+	if (fetch_consistency > PGSTAT_FETCH_CONSISTENCY_NONE)
 	{
 		PgStat_SnapshotEntry *entry = NULL;
 		bool		found;
@@ -915,23 +915,45 @@ pgstat_have_entry(PgStat_Kind kind, Oid dboid, Oid objoid)
 void
 pgstat_snapshot_fixed(PgStat_Kind kind)
 {
+	PgStat_FetchConsistency fetch_consistency;
+
 	Assert(pgstat_is_kind_valid(kind));
 	Assert(pgstat_get_kind_info(kind)->fixed_amount);
 
-	if (pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_SNAPSHOT)
-		pgstat_build_snapshot();
+	fetch_consistency = pgstat_prep_snapshot();
+	if (fetch_consistency == PGSTAT_FETCH_CONSISTENCY_SNAPSHOT)
+		pgstat_build_snapshot(fetch_consistency);
 	else
 		pgstat_build_snapshot_fixed(kind);
 
 	Assert(pgStatLocal.snapshot.fixed_valid[kind]);
 }
 
-static void
+/*
+ * Prepare snapshot hash if needed.
+ *
+ *  Returns the actual snapshot consistency mode to use.
+ */
+static PgStat_FetchConsistency
 pgstat_prep_snapshot(void)
 {
+	/*
+	 * To maintain consistency, reject consistency changes if we have created a
+	 * snapshot.  Non-null stats and PGSTAT_FETCH_CONSISTENCY_NONE mode means
+	 * the snapshot has been prepared and is waiting for being filled.
+	 */
+	if (pgStatLocal.snapshot.stats != NULL &&
+		pgStatLocal.snapshot.mode != PGSTAT_FETCH_CONSISTENCY_NONE &&
+		pgstat_fetch_consistency != pgStatLocal.snapshot.mode)
+	{
+		ereport(WARNING,
+				errmsg("ignored in-transaction change of statistics consistency"));
+		return pgStatLocal.snapshot.mode;
+	}
+
 	if (pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_NONE ||
 		pgStatLocal.snapshot.stats != NULL)
-		return;
+		return pgstat_fetch_consistency;
 
 	if (!pgStatLocal.snapshot.context)
 		pgStatLocal.snapshot.context = AllocSetContextCreate(TopMemoryContext,
@@ -942,22 +964,34 @@ pgstat_prep_snapshot(void)
 		pgstat_snapshot_create(pgStatLocal.snapshot.context,
 							   PGSTAT_SNAPSHOT_HASH_SIZE,
 							   NULL);
+
+	/*
+	 * Enter PGSTAT_FETCH_CONSISTENCY_CACHE mode now. However, we don't enter
+	 * PGSTAT_FETCH_CONSISTENCY_SNAPSHOT mode until pgstat_build_snapshot() has
+	 * completed the remaining task later.
+	 */
+	if (pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_CACHE)
+		pgStatLocal.snapshot.mode = pgstat_fetch_consistency;
+
+	return pgstat_fetch_consistency;
 }
 
 static void
-pgstat_build_snapshot(void)
+pgstat_build_snapshot(PgStat_FetchConsistency fetch_consistency)
 {
 	dshash_seq_status hstat;
 	PgStatShared_HashEntry *p;
 
 	/* should only be called when we need a snapshot */
-	Assert(pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_SNAPSHOT);
+	Assert(fetch_consistency == PGSTAT_FETCH_CONSISTENCY_SNAPSHOT);
 
 	/* snapshot already built */
 	if (pgStatLocal.snapshot.mode == PGSTAT_FETCH_CONSISTENCY_SNAPSHOT)
 		return;
 
-	pgstat_prep_snapshot();
+	/* return immediately if we don't enter snapshot mode */
+	if (pgstat_prep_snapshot() != PGSTAT_FETCH_CONSISTENCY_SNAPSHOT)
+		return;
 
 	Assert(pgStatLocal.snapshot.stats->members == 0);
 
@@ -1014,7 +1048,8 @@ pgstat_build_snapshot(void)
 	dshash_seq_term(&hstat);
 
 	/*
-	 * Build snapshot of all fixed-numbered stats.
+	 * Build snapshot of all fixed-numbered stats.  snapshot mode is required
+	 * to be PGSTAT_FETCH_CONSISTENCY_NONE tentatively while initialization.
 	 */
 	for (int kind = PGSTAT_KIND_FIRST_VALID; kind <= PGSTAT_KIND_LAST; kind++)
 	{
@@ -1040,7 +1075,7 @@ pgstat_build_snapshot_fixed(PgStat_Kind kind)
 	Assert(kind_info->fixed_amount);
 	Assert(kind_info->snapshot_cb != NULL);
 
-	if (pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_NONE)
+	if (pgStatLocal.snapshot.mode == PGSTAT_FETCH_CONSISTENCY_NONE)
 	{
 		/* rebuild every time */
 		pgStatLocal.snapshot.fixed_valid[kind] = false;
@@ -1048,7 +1083,7 @@ pgstat_build_snapshot_fixed(PgStat_Kind kind)
 	else if (pgStatLocal.snapshot.fixed_valid[kind])
 	{
 		/* in snapshot mode we shouldn't get called again */
-		Assert(pgstat_fetch_consistency == PGSTAT_FETCH_CONSISTENCY_CACHE);
+		Assert(pgStatLocal.snapshot.mode == PGSTAT_FETCH_CONSISTENCY_CACHE);
 		return;
 	}