pgstat: Bring up pgstat in BaseInit() to fix uninitialized use of pgstat by AV.

Previously pgstat_initialize() was called in InitPostgres() and
AuxiliaryProcessMain(). As it turns out there was at least one case where we
reported stats before pgstat_initialize() was called, see
AutoVacWorkerMain()'s intentionally early call to pgstat_report_autovac().

This turns out to not be a problem with the current pgstat implementation as
pgstat_initialize() only registers a shutdown callback. But in the shared
memory based stats implementation we are working towards pgstat_initialize()
has to do more work.

After b406478b87e BaseInit() is a central place where initialization shared by
normal backends and auxiliary backends can be put. Obviously BaseInit() is
called before InitPostgres() registers ShutdownPostgres. Previously
ShutdownPostgres was the first before_shmem_exit callback, now that's commonly
pgstats. That should be fine.

Previously pgstat_initialize() was not called in bootstrap mode, but there
does not appear to be a need for that. It's now done unconditionally.

To detect future issues like this, assertions are added to a few places
verifying that the pgstat subsystem is initialized and not yet shut down.

Author: Andres Freund <andres@anarazel.de>
Discussion: https://postgr.es/m/20210405092914.mmxqe7j56lsjfsej@alap3.anarazel.de
Discussion: https://postgr.es/m/20210802164124.ufo5buo4apl6yuvs@alap3.anarazel.de

3 files changed, 65 insertions, 13 deletions

diff --git a/src/backend/postmaster/auxprocess.c b/src/backend/postmaster/auxprocess.c
index 16d87e91402..7452f908b23 100644
--- a/src/backend/postmaster/auxprocess.c
+++ b/src/backend/postmaster/auxprocess.c
@@ -125,8 +125,6 @@ AuxiliaryProcessMain(AuxProcType auxtype)
 	 */
 	CreateAuxProcessResourceOwner();
 
-	/* Initialize statistics reporting */
-	pgstat_initialize();
 
 	/* Initialize backend status information */
 	pgstat_beinit();
diff --git a/src/backend/postmaster/pgstat.c b/src/backend/postmaster/pgstat.c
index 56755cb92b6..9664e7f256f 100644
--- a/src/backend/postmaster/pgstat.c
+++ b/src/backend/postmaster/pgstat.c
@@ -295,6 +295,15 @@ static List *pending_write_requests = NIL;
  */
 static instr_time total_func_time;
 
+/*
+ * For assertions that check pgstat is not used before initialization / after
+ * shutdown.
+ */
+#ifdef USE_ASSERT_CHECKING
+static bool pgstat_is_initialized = false;
+static bool pgstat_is_shutdown = false;
+#endif
+
 
 /* ----------
  * Local function forward declarations
@@ -330,6 +339,7 @@ static void pgstat_send_connstats(bool disconnect, TimestampTz last_report);
 static PgStat_TableStatus *get_tabstat_entry(Oid rel_id, bool isshared);
 
 static void pgstat_setup_memcxt(void);
+static void pgstat_assert_is_up(void);
 
 static void pgstat_setheader(PgStat_MsgHdr *hdr, StatMsgType mtype);
 static void pgstat_send(void *msg, int len);
@@ -854,6 +864,8 @@ pgstat_report_stat(bool disconnect)
 	TabStatusArray *tsa;
 	int			i;
 
+	pgstat_assert_is_up();
+
 	/*
 	 * Don't expend a clock check if nothing to do.
 	 *
@@ -1960,6 +1972,8 @@ pgstat_init_function_usage(FunctionCallInfo fcinfo,
 PgStat_BackendFunctionEntry *
 find_funcstat_entry(Oid func_id)
 {
+	pgstat_assert_is_up();
+
 	if (pgStatFunctions == NULL)
 		return NULL;
 
@@ -2078,6 +2092,8 @@ get_tabstat_entry(Oid rel_id, bool isshared)
 	TabStatusArray *tsa;
 	bool		found;
 
+	pgstat_assert_is_up();
+
 	/*
 	 * Create hash table if we don't have it already.
 	 */
@@ -2938,6 +2954,8 @@ pgstat_fetch_replslot(NameData slotname)
 static void
 pgstat_shutdown_hook(int code, Datum arg)
 {
+	Assert(!pgstat_is_shutdown);
+
 	/*
 	 * If we got as far as discovering our own database ID, we can report what
 	 * we did to the collector.  Otherwise, we'd be sending an invalid
@@ -2946,13 +2964,17 @@ pgstat_shutdown_hook(int code, Datum arg)
 	 */
 	if (OidIsValid(MyDatabaseId))
 		pgstat_report_stat(true);
+
+#ifdef USE_ASSERT_CHECKING
+	pgstat_is_shutdown = true;
+#endif
 }
 
 /* ----------
  * pgstat_initialize() -
  *
- *	Initialize pgstats state, and set up our on-proc-exit hook.
- *	Called from InitPostgres and AuxiliaryProcessMain.
+ *	Initialize pgstats state, and set up our on-proc-exit hook. Called from
+ *	BaseInit().
  *
  *	NOTE: MyDatabaseId isn't set yet; so the shutdown hook has to be careful.
  * ----------
@@ -2960,6 +2982,8 @@ pgstat_shutdown_hook(int code, Datum arg)
 void
 pgstat_initialize(void)
 {
+	Assert(!pgstat_is_initialized);
+
 	/*
 	 * Initialize prevWalUsage with pgWalUsage so that pgstat_send_wal() can
 	 * calculate how much pgWalUsage counters are increased by subtracting
@@ -2969,6 +2993,10 @@ pgstat_initialize(void)
 
 	/* Set up a process-exit hook to clean up */
 	on_shmem_exit(pgstat_shutdown_hook, 0);
+
+#ifdef USE_ASSERT_CHECKING
+	pgstat_is_initialized = true;
+#endif
 }
 
 /* ------------------------------------------------------------
@@ -3001,6 +3029,8 @@ pgstat_send(void *msg, int len)
 {
 	int			rc;
 
+	pgstat_assert_is_up();
+
 	if (pgStatSock == PGINVALID_SOCKET)
 		return;
 
@@ -3053,6 +3083,8 @@ pgstat_send_bgwriter(void)
 	/* We assume this initializes to zeroes */
 	static const PgStat_MsgBgWriter all_zeroes;
 
+	pgstat_assert_is_up();
+
 	/*
 	 * This function can be called even if nothing at all has happened. In
 	 * this case, avoid sending a completely empty message to the stats
@@ -4629,6 +4661,8 @@ backend_read_statsfile(void)
 	Oid			inquiry_db;
 	int			count;
 
+	pgstat_assert_is_up();
+
 	/* already read it? */
 	if (pgStatDBHash)
 		return;
@@ -4765,6 +4799,17 @@ pgstat_setup_memcxt(void)
 												   ALLOCSET_SMALL_SIZES);
 }
 
+/*
+ * Stats should only be reported after pgstat_initialize() and before
+ * pgstat_shutdown(). This check is put in a few central places to catch
+ * violations of this rule more easily.
+ */
+static void
+pgstat_assert_is_up(void)
+{
+	Assert(pgstat_is_initialized && !pgstat_is_shutdown);
+}
+
 
 /* ----------
  * pgstat_clear_snapshot() -
@@ -4779,6 +4824,8 @@ pgstat_setup_memcxt(void)
 void
 pgstat_clear_snapshot(void)
 {
+	pgstat_assert_is_up();
+
 	/* Release memory, if any was allocated */
 	if (pgStatLocalContext)
 		MemoryContextDelete(pgStatLocalContext);
@@ -5897,6 +5944,8 @@ pgstat_slru_name(int slru_idx)
 static inline PgStat_MsgSLRU *
 slru_entry(int slru_idx)
 {
+	pgstat_assert_is_up();
+
 	/*
 	 * The postmaster should never register any SLRU statistics counts; if it
 	 * did, the counts would be duplicated into child processes via fork().
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
index 420b246fb13..e37b86494e1 100644
--- a/src/backend/utils/init/postinit.c
+++ b/src/backend/utils/init/postinit.c
@@ -517,6 +517,14 @@ BaseInit(void)
 	 */
 	DebugFileOpen();
 
+	/*
+	 * Initialize statistics reporting. This needs to happen early to ensure
+	 * that pgstat's shutdown callback runs after the shutdown callbacks of
+	 * all subsystems that can produce stats (like e.g. transaction commits
+	 * can).
+	 */
+	pgstat_initialize();
+
 	/* Do local initialization of file, storage and buffer managers */
 	InitFileAccess();
 	InitSync();
@@ -646,10 +654,6 @@ InitPostgres(const char *in_dbname, Oid dboid, const char *username,
 	/* Initialize portal manager */
 	EnablePortalManager();
 
-	/* Initialize stats collection --- must happen before first xact */
-	if (!bootstrap)
-		pgstat_initialize();
-
 	/* Initialize status reporting */
 	if (!bootstrap)
 		pgstat_beinit();
@@ -662,11 +666,12 @@ InitPostgres(const char *in_dbname, Oid dboid, const char *username,
 
 	/*
 	 * Set up process-exit callback to do pre-shutdown cleanup.  This is the
-	 * first before_shmem_exit callback we register; thus, this will be the
-	 * last thing we do before low-level modules like the buffer manager begin
-	 * to close down.  We need to have this in place before we begin our first
-	 * transaction --- if we fail during the initialization transaction, as is
-	 * entirely possible, we need the AbortTransaction call to clean up.
+	 * one of the first before_shmem_exit callbacks we register; thus, this
+	 * will be one the last things we do before low-level modules like the
+	 * buffer manager begin to close down.  We need to have this in place
+	 * before we begin our first transaction --- if we fail during the
+	 * initialization transaction, as is entirely possible, we need the
+	 * AbortTransaction call to clean up.
 	 */
 	before_shmem_exit(ShutdownPostgres, 0);