From 9be7360e49db424c45c53e85efe8a4f5359b5244 Mon Sep 17 00:00:00 2001
From: Melih Mutlu <m.melihmutlu@gmail.com>
Date: Wed, 26 Apr 2023 18:21:32 +0300
Subject: [PATCH v2] Add timeout to flush stats during startup's main replay
 loop

---
 src/backend/access/transam/xlogrecovery.c |  7 +++++
 src/backend/postmaster/startup.c          | 34 +++++++++++++++++++++++
 src/backend/utils/activity/pgstat.c       | 14 ----------
 src/include/pgstat.h                      | 14 ++++++++++
 src/include/postmaster/startup.h          |  3 ++
 5 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/src/backend/access/transam/xlogrecovery.c b/src/backend/access/transam/xlogrecovery.c
index 188f6d6f85..53a4939be9 100644
--- a/src/backend/access/transam/xlogrecovery.c
+++ b/src/backend/access/transam/xlogrecovery.c
@@ -63,6 +63,7 @@
 #include "utils/pg_lsn.h"
 #include "utils/ps_status.h"
 #include "utils/pg_rusage.h"
+#include "utils/timeout.h"
 
 /* Unsupported old recovery command file names (relative to $PGDATA) */
 #define RECOVERY_COMMAND_FILE	"recovery.conf"
@@ -1628,6 +1629,10 @@ PerformWalRecovery(void)
 	 */
 	CheckRecoveryConsistency();
 
+	/* Register the timeout to flush wal stats when it's needed. */
+	RegisterTimeout(IDLE_STATS_UPDATE_TIMEOUT,
+					idle_stats_update_timeout_handler);
+
 	/*
 	 * Find the first record that logically follows the checkpoint --- it
 	 * might physically precede it, though.
@@ -1666,6 +1671,8 @@ PerformWalRecovery(void)
 		if (!StandbyMode)
 			begin_startup_progress_phase();
 
+		enable_idle_stats_update_timeout();
+
 		/*
 		 * main redo apply loop
 		 */
diff --git a/src/backend/postmaster/startup.c b/src/backend/postmaster/startup.c
index efc2580536..842394bc8f 100644
--- a/src/backend/postmaster/startup.c
+++ b/src/backend/postmaster/startup.c
@@ -72,6 +72,9 @@ static TimestampTz startup_progress_phase_start_time;
  */
 static volatile sig_atomic_t startup_progress_timer_expired = false;
 
+/* Indicates whether flushing stats is needed. */
+static volatile sig_atomic_t idle_stats_update_pending = false;
+
 /*
  * Time between progress updates for long-running startup operations.
  */
@@ -206,6 +209,18 @@ HandleStartupProcInterrupts(void)
 	/* Perform logging of memory contexts of this process */
 	if (LogMemoryContextPending)
 		ProcessLogMemoryContextInterrupt();
+
+	if (idle_stats_update_pending)
+	{
+		/* It's time to report wal stats. */
+		pgstat_report_wal(true);
+		idle_stats_update_pending = false;
+	}
+	else if (!get_timeout_active(IDLE_STATS_UPDATE_TIMEOUT))
+	{
+		/* Set the next timeout. */
+		enable_idle_stats_update_timeout();
+	}
 }
 
 
@@ -385,3 +400,22 @@ has_startup_progress_timeout_expired(long *secs, int *usecs)
 
 	return true;
 }
+
+/* Set a flag indicating that it's time to flush wal stats. */
+void
+idle_stats_update_timeout_handler(void)
+{
+	idle_stats_update_pending = true;
+	WakeupRecovery();
+}
+
+/* Enable the timeout set for wal stat flush. */
+void
+enable_idle_stats_update_timeout(void)
+{
+	TimestampTz fin_time;
+
+	fin_time = TimestampTzPlusMilliseconds(GetCurrentTimestamp(),
+										   PGSTAT_MIN_INTERVAL);
+	enable_timeout_at(IDLE_STATS_UPDATE_TIMEOUT, fin_time);
+}
\ No newline at end of file
diff --git a/src/backend/utils/activity/pgstat.c b/src/backend/utils/activity/pgstat.c
index b125802b21..a3ca3e648d 100644
--- a/src/backend/utils/activity/pgstat.c
+++ b/src/backend/utils/activity/pgstat.c
@@ -109,20 +109,6 @@
 #include "utils/timestamp.h"
 
 
-/* ----------
- * Timer definitions.
- *
- * In milliseconds.
- * ----------
- */
-
-/* minimum interval non-forced stats flushes.*/
-#define PGSTAT_MIN_INTERVAL			1000
-/* how long until to block flushing pending stats updates */
-#define PGSTAT_MAX_INTERVAL			60000
-/* when to call pgstat_report_stat() again, even when idle */
-#define PGSTAT_IDLE_INTERVAL		10000
-
 /* ----------
  * Initial size hints for the hash tables used in statistics.
  * ----------
diff --git a/src/include/pgstat.h b/src/include/pgstat.h
index fff4ad5b6d..069e08006c 100644
--- a/src/include/pgstat.h
+++ b/src/include/pgstat.h
@@ -31,6 +31,20 @@
 /* Default directory to store temporary statistics data in */
 #define PG_STAT_TMP_DIR		"pg_stat_tmp"
 
+/* ----------
+ * Timer definitions.
+ *
+ * In milliseconds.
+ * ----------
+ */
+
+/* minimum interval non-forced stats flushes.*/
+#define PGSTAT_MIN_INTERVAL			1000
+/* how long until to block flushing pending stats updates */
+#define PGSTAT_MAX_INTERVAL			60000
+/* when to call pgstat_report_stat() again, even when idle */
+#define PGSTAT_IDLE_INTERVAL		10000
+
 /* The types of statistics entries */
 typedef enum PgStat_Kind
 {
diff --git a/src/include/postmaster/startup.h b/src/include/postmaster/startup.h
index 6a2e4c4526..f1e051811f 100644
--- a/src/include/postmaster/startup.h
+++ b/src/include/postmaster/startup.h
@@ -38,4 +38,7 @@ extern void begin_startup_progress_phase(void);
 extern void startup_progress_timeout_handler(void);
 extern bool has_startup_progress_timeout_expired(long *secs, int *usecs);
 
+extern void enable_idle_stats_update_timeout(void);
+extern void idle_stats_update_timeout_handler(void);
+
 #endif							/* _STARTUP_H */
-- 
2.25.1