v17-0004-guarantee-to-collect-last-stats-messages.patch

diff --git a/src/backend/postmaster/pgstat.c b/src/backend/postmaster/pgstat.c
index 68eefb9722..86b8449193 100644
--- a/src/backend/postmaster/pgstat.c
+++ b/src/backend/postmaster/pgstat.c
@@ -329,6 +329,7 @@ static void pgstat_beshutdown_hook(int code, Datum arg);
 static PgStat_StatDBEntry *pgstat_get_db_entry(Oid databaseid, bool create);
 static PgStat_StatTabEntry *pgstat_get_tab_entry(PgStat_StatDBEntry *dbentry,
 												 Oid tableoid, bool create);
+static int	pgstat_process_message(void);
 static void pgstat_write_statsfiles(bool permanent, bool allDbs);
 static void pgstat_write_db_statsfile(PgStat_StatDBEntry *dbentry, bool permanent);
 static HTAB *pgstat_read_statsfiles(Oid onlydb, bool permanent, bool deep);
@@ -4810,8 +4811,6 @@ pgstat_send_slru(void)
 NON_EXEC_STATIC void
 PgstatCollectorMain(int argc, char *argv[])
 {
-	int			len;
-	PgStat_Msg	msg;
 	int			wr;
 	WaitEvent	event;
 	WaitEventSet *wes;
@@ -4896,158 +4895,10 @@ PgstatCollectorMain(int argc, char *argv[])
 				pgstat_write_statsfiles(false, false);
 
 			/*
-			 * Try to receive and process a message.  This will not block,
-			 * since the socket is set to non-blocking mode.
-			 *
-			 * XXX On Windows, we have to force pgwin32_recv to cooperate,
-			 * despite the previous use of pg_set_noblock() on the socket.
-			 * This is extremely broken and should be fixed someday.
+			 * Try to receive and process a message.
 			 */
-#ifdef WIN32
-			pgwin32_noblock = 1;
-#endif
-
-			len = recv(pgStatSock, (char *) &msg,
-					   sizeof(PgStat_Msg), 0);
-
-#ifdef WIN32
-			pgwin32_noblock = 0;
-#endif
-
-			if (len < 0)
-			{
-				if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
-					break;		/* out of inner loop */
-				ereport(ERROR,
-						(errcode_for_socket_access(),
-						 errmsg("could not read statistics message: %m")));
-			}
-
-			/*
-			 * We ignore messages that are smaller than our common header
-			 */
-			if (len < sizeof(PgStat_MsgHdr))
-				continue;
-
-			/*
-			 * The received length must match the length in the header
-			 */
-			if (msg.msg_hdr.m_size != len)
-				continue;
-
-			/*
-			 * O.K. - we accept this message.  Process it.
-			 */
-			switch (msg.msg_hdr.m_type)
-			{
-				case PGSTAT_MTYPE_DUMMY:
-					break;
-
-				case PGSTAT_MTYPE_INQUIRY:
-					pgstat_recv_inquiry(&msg.msg_inquiry, len);
-					break;
-
-				case PGSTAT_MTYPE_TABSTAT:
-					pgstat_recv_tabstat(&msg.msg_tabstat, len);
-					break;
-
-				case PGSTAT_MTYPE_TABPURGE:
-					pgstat_recv_tabpurge(&msg.msg_tabpurge, len);
-					break;
-
-				case PGSTAT_MTYPE_DROPDB:
-					pgstat_recv_dropdb(&msg.msg_dropdb, len);
-					break;
-
-				case PGSTAT_MTYPE_RESETCOUNTER:
-					pgstat_recv_resetcounter(&msg.msg_resetcounter, len);
-					break;
-
-				case PGSTAT_MTYPE_RESETSHAREDCOUNTER:
-					pgstat_recv_resetsharedcounter(&msg.msg_resetsharedcounter,
-												   len);
-					break;
-
-				case PGSTAT_MTYPE_RESETSINGLECOUNTER:
-					pgstat_recv_resetsinglecounter(&msg.msg_resetsinglecounter,
-												   len);
-					break;
-
-				case PGSTAT_MTYPE_RESETSLRUCOUNTER:
-					pgstat_recv_resetslrucounter(&msg.msg_resetslrucounter,
-												 len);
-					break;
-
-				case PGSTAT_MTYPE_RESETREPLSLOTCOUNTER:
-					pgstat_recv_resetreplslotcounter(&msg.msg_resetreplslotcounter,
-													 len);
-					break;
-
-				case PGSTAT_MTYPE_AUTOVAC_START:
-					pgstat_recv_autovac(&msg.msg_autovacuum_start, len);
-					break;
-
-				case PGSTAT_MTYPE_VACUUM:
-					pgstat_recv_vacuum(&msg.msg_vacuum, len);
-					break;
-
-				case PGSTAT_MTYPE_ANALYZE:
-					pgstat_recv_analyze(&msg.msg_analyze, len);
-					break;
-
-				case PGSTAT_MTYPE_ARCHIVER:
-					pgstat_recv_archiver(&msg.msg_archiver, len);
-					break;
-
-				case PGSTAT_MTYPE_BGWRITER:
-					pgstat_recv_bgwriter(&msg.msg_bgwriter, len);
-					break;
-
-				case PGSTAT_MTYPE_WAL:
-					pgstat_recv_wal(&msg.msg_wal, len);
-					break;
-
-				case PGSTAT_MTYPE_SLRU:
-					pgstat_recv_slru(&msg.msg_slru, len);
-					break;
-
-				case PGSTAT_MTYPE_FUNCSTAT:
-					pgstat_recv_funcstat(&msg.msg_funcstat, len);
-					break;
-
-				case PGSTAT_MTYPE_FUNCPURGE:
-					pgstat_recv_funcpurge(&msg.msg_funcpurge, len);
-					break;
-
-				case PGSTAT_MTYPE_RECOVERYCONFLICT:
-					pgstat_recv_recoveryconflict(&msg.msg_recoveryconflict,
-												 len);
-					break;
-
-				case PGSTAT_MTYPE_DEADLOCK:
-					pgstat_recv_deadlock(&msg.msg_deadlock, len);
-					break;
-
-				case PGSTAT_MTYPE_TEMPFILE:
-					pgstat_recv_tempfile(&msg.msg_tempfile, len);
-					break;
-
-				case PGSTAT_MTYPE_CHECKSUMFAILURE:
-					pgstat_recv_checksum_failure(&msg.msg_checksumfailure,
-												 len);
-					break;
-
-				case PGSTAT_MTYPE_REPLSLOT:
-					pgstat_recv_replslot(&msg.msg_replslot, len);
-					break;
-
-				case PGSTAT_MTYPE_CONNECTION:
-					pgstat_recv_connstat(&msg.msg_conn, len);
-					break;
-
-				default:
-					break;
-			}
+			if (pgstat_process_message() < 0)
+				break;			/* If an error occurred, go out of inner loop */
 		}						/* end of inner message-processing loop */
 
 		/* Sleep until there's something to do */
@@ -5077,6 +4928,21 @@ PgstatCollectorMain(int argc, char *argv[])
 			break;
 	}							/* end of outer loop */
 
+	/*
+	 * Try to receive and process remaining messages before the process exits.
+	 *
+	 * The reason is that there is no guarantee all messages were processed in
+	 * the above loop even though the stats collector is sent SIGQUIT signal
+	 * by the postmaster after other backend and background processes, which
+	 * sent their stats to the stats collector, exit if shutdown mode is smart
+	 * or fast.
+	 *
+	 * For example, there might be a case that messages are lost when there
+	 * are unprocessed messages, the postmaster send SIGQUIT signal to the
+	 * stats collector.
+	 */
+	while (pgstat_process_message() > 0);
+
 	/*
 	 * Save the final stats to reuse at next startup.
 	 */
@@ -5225,6 +5091,169 @@ pgstat_get_tab_entry(PgStat_StatDBEntry *dbentry, Oid tableoid, bool create)
 	return result;
 }
 
+/*
+ * Try to receive and process a message.  This will not block,
+ * since the socket is set to non-blocking mode.
+ *
+ * XXX On Windows, we have to force pgwin32_recv to cooperate,
+ * despite the previous use of pg_set_noblock() on the socket.
+ * This is extremely broken and should be fixed someday.
+ *
+ * Return the number of processed message. -1 if an error occurred.
+ */
+static int
+pgstat_process_message()
+{
+	int			len;
+	PgStat_Msg	msg;
+
+#ifdef WIN32
+	pgwin32_noblock = 1;
+#endif
+
+	len = recv(pgStatSock, (char *) &msg, sizeof(PgStat_Msg), 0);
+
+#ifdef WIN32
+	pgwin32_noblock = 0;
+#endif
+
+	if (len < 0)
+	{
+		if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
+			return -1;
+		ereport(ERROR,
+				(errcode_for_socket_access(),
+				 errmsg("could not read statistics message: %m")));
+	}
+
+	/*
+	 * We ignore messages that are smaller than our common header
+	 */
+	if (len < sizeof(PgStat_MsgHdr))
+		return 0;
+
+	/*
+	 * The received length must match the length in the header
+	 */
+	if (msg.msg_hdr.m_size != len)
+		return 0;
+
+	/*
+	 * O.K. - we accept this message.  Process it.
+	 */
+	switch (msg.msg_hdr.m_type)
+	{
+		case PGSTAT_MTYPE_DUMMY:
+			break;
+
+		case PGSTAT_MTYPE_INQUIRY:
+			pgstat_recv_inquiry(&msg.msg_inquiry, len);
+			break;
+
+		case PGSTAT_MTYPE_TABSTAT:
+			pgstat_recv_tabstat(&msg.msg_tabstat, len);
+			break;
+
+		case PGSTAT_MTYPE_TABPURGE:
+			pgstat_recv_tabpurge(&msg.msg_tabpurge, len);
+			break;
+
+		case PGSTAT_MTYPE_DROPDB:
+			pgstat_recv_dropdb(&msg.msg_dropdb, len);
+			break;
+
+		case PGSTAT_MTYPE_RESETCOUNTER:
+			pgstat_recv_resetcounter(&msg.msg_resetcounter, len);
+			break;
+
+		case PGSTAT_MTYPE_RESETSHAREDCOUNTER:
+			pgstat_recv_resetsharedcounter(&msg.msg_resetsharedcounter,
+										   len);
+			break;
+
+		case PGSTAT_MTYPE_RESETSINGLECOUNTER:
+			pgstat_recv_resetsinglecounter(&msg.msg_resetsinglecounter,
+										   len);
+			break;
+
+		case PGSTAT_MTYPE_RESETSLRUCOUNTER:
+			pgstat_recv_resetslrucounter(&msg.msg_resetslrucounter,
+										 len);
+			break;
+
+		case PGSTAT_MTYPE_RESETREPLSLOTCOUNTER:
+			pgstat_recv_resetreplslotcounter(&msg.msg_resetreplslotcounter,
+											 len);
+			break;
+
+		case PGSTAT_MTYPE_AUTOVAC_START:
+			pgstat_recv_autovac(&msg.msg_autovacuum_start, len);
+			break;
+
+		case PGSTAT_MTYPE_VACUUM:
+			pgstat_recv_vacuum(&msg.msg_vacuum, len);
+			break;
+
+		case PGSTAT_MTYPE_ANALYZE:
+			pgstat_recv_analyze(&msg.msg_analyze, len);
+			break;
+
+		case PGSTAT_MTYPE_ARCHIVER:
+			pgstat_recv_archiver(&msg.msg_archiver, len);
+			break;
+
+		case PGSTAT_MTYPE_BGWRITER:
+			pgstat_recv_bgwriter(&msg.msg_bgwriter, len);
+			break;
+
+		case PGSTAT_MTYPE_WAL:
+			pgstat_recv_wal(&msg.msg_wal, len);
+			break;
+
+		case PGSTAT_MTYPE_SLRU:
+			pgstat_recv_slru(&msg.msg_slru, len);
+			break;
+
+		case PGSTAT_MTYPE_FUNCSTAT:
+			pgstat_recv_funcstat(&msg.msg_funcstat, len);
+			break;
+
+		case PGSTAT_MTYPE_FUNCPURGE:
+			pgstat_recv_funcpurge(&msg.msg_funcpurge, len);
+			break;
+
+		case PGSTAT_MTYPE_RECOVERYCONFLICT:
+			pgstat_recv_recoveryconflict(&msg.msg_recoveryconflict,
+										 len);
+			break;
+
+		case PGSTAT_MTYPE_DEADLOCK:
+			pgstat_recv_deadlock(&msg.msg_deadlock, len);
+			break;
+
+		case PGSTAT_MTYPE_TEMPFILE:
+			pgstat_recv_tempfile(&msg.msg_tempfile, len);
+			break;
+
+		case PGSTAT_MTYPE_CHECKSUMFAILURE:
+			pgstat_recv_checksum_failure(&msg.msg_checksumfailure,
+										 len);
+			break;
+
+		case PGSTAT_MTYPE_REPLSLOT:
+			pgstat_recv_replslot(&msg.msg_replslot, len);
+			break;
+
+		case PGSTAT_MTYPE_CONNECTION:
+			pgstat_recv_connstat(&msg.msg_conn, len);
+			break;
+
+		default:
+			break;
+	}
+
+	return 1;
+}
 
 /* ----------
  * pgstat_write_statsfiles() -