Add callback for backend initialization in pgstats

michaelpq · michaelpq · commit 1b373aed20e6 · 2024-09-05T16:05:21.000+09:00
pgstat_initialize() is currently used by the WAL stats as a code path to take some custom actions when a backend starts. A callback is added to generalize the concept so as all stats kinds can do the same, for builtin and custom kinds, if set. Reviewed-by: Bertrand Drouvot, Kyotaro Horiguchi Discussion: https://postgr.es/m/ZtZr1K4PLdeWclXY@paquier.xyz
diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
@@ -441,6 +441,7 @@ static const PgStat_KindInfo pgstat_kind_builtin_infos[PGSTAT_KIND_BUILTIN_SIZE]
 		.shared_data_off = offsetof(PgStatShared_Wal, stats),
 		.shared_data_len = sizeof(((PgStatShared_Wal *) 0)->stats),
 
+		.init_backend_cb = pgstat_wal_init_backend_cb,
 		.init_shmem_cb = pgstat_wal_init_shmem_cb,
 		.reset_all_cb = pgstat_wal_reset_all_cb,
 		.snapshot_cb = pgstat_wal_snapshot_cb,
@@ -604,10 +605,19 @@ pgstat_initialize(void)
 
 	pgstat_attach_shmem();
 
-	pgstat_init_wal();
-
 	pgstat_init_snapshot_fixed();
 
+	/* Backend initialization callbacks */
+	for (PgStat_Kind kind = PGSTAT_KIND_MIN; kind <= PGSTAT_KIND_MAX; kind++)
+	{
+		const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
+
+		if (kind_info == NULL || kind_info->init_backend_cb == NULL)
+			continue;
+
+		kind_info->init_backend_cb();
+	}
+
 	/* Set up a process-exit hook to clean up */
 	before_shmem_exit(pgstat_shutdown_hook, 0);
 
diff --git a/src/backend/utils/activity/pgstat_wal.c b/src/backend/utils/activity/pgstat_wal.c
@@ -138,7 +138,7 @@ pgstat_flush_wal(bool nowait)
 }
 
 void
-pgstat_init_wal(void)
+pgstat_wal_init_backend_cb(void)
 {
 	/*
 	 * Initialize prevWalUsage with pgWalUsage so that pgstat_flush_wal() can
diff --git a/src/include/utils/pgstat_internal.h b/src/include/utils/pgstat_internal.h
@@ -229,6 +229,12 @@ typedef struct PgStat_KindInfo
 	 */
 	uint32		pending_size;
 
+	/*
+	 * Perform custom actions when initializing a backend (standalone or under
+	 * postmaster). Optional.
+	 */
+	void		(*init_backend_cb) (void);
+
 	/*
 	 * For variable-numbered stats: flush pending stats. Required if pending
 	 * data is used.
@@ -673,9 +679,9 @@ extern void pgstat_slru_snapshot_cb(void);
  */
 
 extern bool pgstat_flush_wal(bool nowait);
-extern void pgstat_init_wal(void);
 extern bool pgstat_have_pending_wal(void);
 
+extern void pgstat_wal_init_backend_cb(void);
 extern void pgstat_wal_init_shmem_cb(void *stats);
 extern void pgstat_wal_reset_all_cb(TimestampTz ts);
 extern void pgstat_wal_snapshot_cb(void);

Original file line number	Diff line number	Diff line change
`@@ -138,7 +138,7 @@ pgstat_flush_wal(bool nowait)`
`138`	`138`	`}`
`139`	`139`
`140`	`140`	`void`
`141`		`-pgstat_init_wal(void)`
	`141`	`+pgstat_wal_init_backend_cb(void)`
`142`	`142`	`{`
`143`	`143`	`/*`
`144`	`144`	`* Initialize prevWalUsage with pgWalUsage so that pgstat_flush_wal() can`