From 5548477e192abce02a95249f4c1e0c9cf465d668 Mon Sep 17 00:00:00 2001
From: "yizhi.fzh" <yizhi.fzh@alibaba-inc.com>
Date: Thu, 11 Jan 2024 11:12:00 +0800
Subject: [PATCH v2] Detect more misuse of spin lock automatically

spin lock are intended for *very* short-term locks, but it is possible
to be misused in many cases. e.g. Acquiring another LWLocks or regular
locks, memory allocation. In this patch, all of such cases will be
automatically detected in an ASSERT_CHECKING build.

CHECK_FOR_INTERRUPTS should be avoided when holding a spin lock because
it is nearly impossible to release the spin lock correctly if an
error/fatal happens would cause the code jump to some other places.
---
 src/backend/access/table/tableam.c  |  1 +
 src/backend/storage/buffer/bufmgr.c |  1 +
 src/backend/storage/ipc/barrier.c   |  2 ++
 src/backend/storage/lmgr/lock.c     |  2 ++
 src/backend/storage/lmgr/lwlock.c   |  2 ++
 src/backend/storage/lmgr/s_lock.c   | 15 ++++++++-------
 src/backend/utils/hash/dynahash.c   |  1 +
 src/backend/utils/init/globals.c    |  1 +
 src/backend/utils/mmgr/mcxt.c       |  9 +++++++++
 src/include/miscadmin.h             |  9 +++++++++
 src/include/storage/buf_internals.h |  1 +
 src/include/storage/s_lock.h        |  2 ++
 src/include/storage/shm_toc.h       |  1 +
 src/include/storage/spin.h          | 17 ++++++++++++++---
 14 files changed, 54 insertions(+), 10 deletions(-)

diff --git a/src/backend/access/table/tableam.c b/src/backend/access/table/tableam.c
index 6ed8cca05a..6c6a65764c 100644
--- a/src/backend/access/table/tableam.c
+++ b/src/backend/access/table/tableam.c
@@ -24,6 +24,7 @@
 #include "access/syncscan.h"
 #include "access/tableam.h"
 #include "access/xact.h"
+#include "miscadmin.h"
 #include "optimizer/plancat.h"
 #include "port/pg_bitutils.h"
 #include "storage/bufmgr.h"
diff --git a/src/backend/storage/buffer/bufmgr.c b/src/backend/storage/buffer/bufmgr.c
index 7d601bef6d..c600a113cf 100644
--- a/src/backend/storage/buffer/bufmgr.c
+++ b/src/backend/storage/buffer/bufmgr.c
@@ -5409,6 +5409,7 @@ LockBufHdr(BufferDesc *desc)
 
 	init_local_spin_delay(&delayStatus);
 
+	START_SPIN_LOCK();
 	while (true)
 	{
 		/* set BM_LOCKED flag */
diff --git a/src/backend/storage/ipc/barrier.c b/src/backend/storage/ipc/barrier.c
index 5b52e060ba..513835ccf2 100644
--- a/src/backend/storage/ipc/barrier.c
+++ b/src/backend/storage/ipc/barrier.c
@@ -84,6 +84,8 @@
  */
 
 #include "postgres.h"
+
+#include "miscadmin.h"
 #include "storage/barrier.h"
 
 static inline bool BarrierDetachImpl(Barrier *barrier, bool arrive);
diff --git a/src/backend/storage/lmgr/lock.c b/src/backend/storage/lmgr/lock.c
index c70a1adb9a..f896201244 100644
--- a/src/backend/storage/lmgr/lock.c
+++ b/src/backend/storage/lmgr/lock.c
@@ -776,6 +776,8 @@ LockAcquireExtended(const LOCKTAG *locktag,
 	bool		found_conflict;
 	bool		log_lock = false;
 
+	Assert(SpinLockCount == 0);
+
 	if (lockmethodid <= 0 || lockmethodid >= lengthof(LockMethods))
 		elog(ERROR, "unrecognized lock method: %d", lockmethodid);
 	lockMethodTable = LockMethods[lockmethodid];
diff --git a/src/backend/storage/lmgr/lwlock.c b/src/backend/storage/lmgr/lwlock.c
index b4b989ac56..974a7c26ea 100644
--- a/src/backend/storage/lmgr/lwlock.c
+++ b/src/backend/storage/lmgr/lwlock.c
@@ -1205,6 +1205,8 @@ LWLockAcquire(LWLock *lock, LWLockMode mode)
 
 	Assert(mode == LW_SHARED || mode == LW_EXCLUSIVE);
 
+	Assert(SpinLockCount == 0);
+
 	PRINT_LWDEBUG("LWLockAcquire", lock, mode);
 
 #ifdef LWLOCK_STATS
diff --git a/src/backend/storage/lmgr/s_lock.c b/src/backend/storage/lmgr/s_lock.c
index 0e5f7ab0b9..43a05b68be 100644
--- a/src/backend/storage/lmgr/s_lock.c
+++ b/src/backend/storage/lmgr/s_lock.c
@@ -71,18 +71,18 @@ static int	spins_per_delay = DEFAULT_SPINS_PER_DELAY;
  * s_lock_stuck() - complain about a stuck spinlock
  */
 static void
-s_lock_stuck(const char *file, int line, const char *func)
+s_lock_stuck(const char *file, int line, const char *func, uint32_t delay_ms)
 {
 	if (!func)
 		func = "(unknown)";
 #if defined(S_LOCK_TEST)
-	fprintf(stderr,
-			"\nStuck spinlock detected at %s, %s:%d.\n",
-			func, file, line);
+	fprintf(stder,
+			"\nStuck spinlock detected at %s, %s:%d. after waiting for %u ms\n",
+			func, file, line, delay_ms);
 	exit(1);
 #else
-	elog(PANIC, "stuck spinlock detected at %s, %s:%d",
-		 func, file, line);
+	elog(PANIC, "stuck spinlock detected at %s, %s:%d after waiting for %u ms",
+		 func, file, line, delay_ms);
 #endif
 }
 
@@ -132,7 +132,7 @@ perform_spin_delay(SpinDelayStatus *status)
 	if (++(status->spins) >= spins_per_delay)
 	{
 		if (++(status->delays) > NUM_DELAYS)
-			s_lock_stuck(status->file, status->line, status->func);
+			s_lock_stuck(status->file, status->line, status->func, status->delay_ms);
 
 		if (status->cur_delay == 0) /* first time to delay? */
 			status->cur_delay = MIN_DELAY_USEC;
@@ -148,6 +148,7 @@ perform_spin_delay(SpinDelayStatus *status)
 		pgstat_report_wait_start(WAIT_EVENT_SPIN_DELAY);
 		pg_usleep(status->cur_delay);
 		pgstat_report_wait_end();
+		status->delay_ms += (status->cur_delay / 1000);
 
 #if defined(S_LOCK_TEST)
 		fprintf(stdout, "*");
diff --git a/src/backend/utils/hash/dynahash.c b/src/backend/utils/hash/dynahash.c
index a4152080b5..504dae4d6c 100644
--- a/src/backend/utils/hash/dynahash.c
+++ b/src/backend/utils/hash/dynahash.c
@@ -98,6 +98,7 @@
 
 #include "access/xact.h"
 #include "common/hashfn.h"
+#include "miscadmin.h"
 #include "port/pg_bitutils.h"
 #include "storage/shmem.h"
 #include "storage/spin.h"
diff --git a/src/backend/utils/init/globals.c b/src/backend/utils/init/globals.c
index 88b03e8fa3..913f775d2c 100644
--- a/src/backend/utils/init/globals.c
+++ b/src/backend/utils/init/globals.c
@@ -40,6 +40,7 @@ volatile sig_atomic_t IdleStatsUpdateTimeoutPending = false;
 volatile uint32 InterruptHoldoffCount = 0;
 volatile uint32 QueryCancelHoldoffCount = 0;
 volatile uint32 CritSectionCount = 0;
+volatile uint32 SpinLockCount = 0;
 
 int			MyProcPid;
 pg_time_t	MyStartTime;
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 1336944084..5acb03ac60 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -26,6 +26,7 @@
 #include "storage/proc.h"
 #include "storage/procarray.h"
 #include "storage/procsignal.h"
+#include "storage/spin.h"
 #include "utils/fmgrprotos.h"
 #include "utils/memdebug.h"
 #include "utils/memutils.h"
@@ -1028,6 +1029,8 @@ MemoryContextAlloc(MemoryContext context, Size size)
 	if (!AllocSizeIsValid(size))
 		elog(ERROR, "invalid memory alloc request size %zu", size);
 
+	Assert(SpinLockCount == 0);
+
 	context->isReset = false;
 
 	ret = context->methods->alloc(context, size);
@@ -1071,6 +1074,8 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 	if (!AllocSizeIsValid(size))
 		elog(ERROR, "invalid memory alloc request size %zu", size);
 
+	Assert(SpinLockCount == 0);
+
 	context->isReset = false;
 
 	ret = context->methods->alloc(context, size);
@@ -1197,6 +1202,8 @@ palloc(Size size)
 	if (!AllocSizeIsValid(size))
 		elog(ERROR, "invalid memory alloc request size %zu", size);
 
+	Assert(SpinLockCount == 0);
+
 	context->isReset = false;
 
 	ret = context->methods->alloc(context, size);
@@ -1228,6 +1235,8 @@ palloc0(Size size)
 	if (!AllocSizeIsValid(size))
 		elog(ERROR, "invalid memory alloc request size %zu", size);
 
+	Assert(SpinLockCount == 0);
+
 	context->isReset = false;
 
 	ret = context->methods->alloc(context, size);
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index 0b01c1f093..91d256a943 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -103,6 +103,7 @@ extern PGDLLIMPORT volatile sig_atomic_t ClientConnectionLost;
 extern PGDLLIMPORT volatile uint32 InterruptHoldoffCount;
 extern PGDLLIMPORT volatile uint32 QueryCancelHoldoffCount;
 extern PGDLLIMPORT volatile uint32 CritSectionCount;
+extern PGDLLIMPORT volatile uint32 SpinLockCount;
 
 /* in tcop/postgres.c */
 extern void ProcessInterrupts(void);
@@ -120,6 +121,7 @@ extern void ProcessInterrupts(void);
 /* Service interrupt, if one is pending and it's safe to service it now */
 #define CHECK_FOR_INTERRUPTS() \
 do { \
+	Assert(SpinLockCount == 0);	\
 	if (INTERRUPTS_PENDING_CONDITION()) \
 		ProcessInterrupts(); \
 } while(0)
@@ -153,6 +155,13 @@ do { \
 	CritSectionCount--; \
 } while(0)
 
+#define START_SPIN_LOCK() (SpinLockCount++)
+#define END_SPIN_LOCK() \
+do { \
+	Assert(SpinLockCount > 0); \
+	SpinLockCount--; \
+} while(0)
+
 
 /*****************************************************************************
  *	  globals.h --															 *
diff --git a/src/include/storage/buf_internals.h b/src/include/storage/buf_internals.h
index e43e616579..ad7e1d0a21 100644
--- a/src/include/storage/buf_internals.h
+++ b/src/include/storage/buf_internals.h
@@ -363,6 +363,7 @@ UnlockBufHdr(BufferDesc *desc, uint32 buf_state)
 {
 	pg_write_barrier();
 	pg_atomic_write_u32(&desc->state, buf_state & (~BM_LOCKED));
+	END_SPIN_LOCK();
 }
 
 /* in bufmgr.c */
diff --git a/src/include/storage/s_lock.h b/src/include/storage/s_lock.h
index aa06e49da2..ac875ab3dd 100644
--- a/src/include/storage/s_lock.h
+++ b/src/include/storage/s_lock.h
@@ -846,6 +846,7 @@ typedef struct
 	const char *file;
 	int			line;
 	const char *func;
+	uint32_t	delay_ms;
 } SpinDelayStatus;
 
 static inline void
@@ -858,6 +859,7 @@ init_spin_delay(SpinDelayStatus *status,
 	status->file = file;
 	status->line = line;
 	status->func = func;
+	status->delay_ms = 0;
 }
 
 #define init_local_spin_delay(status) init_spin_delay(status, __FILE__, __LINE__, __func__)
diff --git a/src/include/storage/shm_toc.h b/src/include/storage/shm_toc.h
index a46470d68b..dd0e591340 100644
--- a/src/include/storage/shm_toc.h
+++ b/src/include/storage/shm_toc.h
@@ -23,6 +23,7 @@
 #define SHM_TOC_H
 
 #include "storage/shmem.h"		/* for add_size() */
+#include "miscadmin.h"
 
 /* shm_toc is an opaque type known only within shm_toc.c */
 typedef struct shm_toc shm_toc;
diff --git a/src/include/storage/spin.h b/src/include/storage/spin.h
index c0679c5999..99bced0f93 100644
--- a/src/include/storage/spin.h
+++ b/src/include/storage/spin.h
@@ -56,12 +56,23 @@
 #include "storage/pg_sema.h"
 #endif
 
-
 #define SpinLockInit(lock)	S_INIT_LOCK(lock)
 
-#define SpinLockAcquire(lock) S_LOCK(lock)
+#define SpinLockAcquire(lock) \
+	do \
+	{ \
+		START_SPIN_LOCK(); \
+		S_LOCK(lock); \
+	} while (false);
+
+
+#define SpinLockRelease(lock) \
+	do \
+	{ \
+		S_UNLOCK(lock); \
+		END_SPIN_LOCK(); \
+	} while (false);
 
-#define SpinLockRelease(lock) S_UNLOCK(lock)
 
 #define SpinLockFree(lock)	S_LOCK_FREE(lock)
 
-- 
2.34.1

