From c1b332826070de4782b12bbb577d66406d439c8d Mon Sep 17 00:00:00 2001
From: Shawn Debnath <sdn@amazon.com>
Date: Tue, 12 Mar 2019 23:21:41 +0000
Subject: [PATCH] Introduce timeout capability for ConditionVariableSleep

This patch introduces ConditionVariableTimedSleep to enable timing out
of waiting for a condition variable to get signalled.
---
 src/backend/storage/lmgr/condition_variable.c | 41 +++++++++++++++++++++++++--
 src/include/storage/condition_variable.h      |  2 ++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/src/backend/storage/lmgr/condition_variable.c b/src/backend/storage/lmgr/condition_variable.c
index 58b7b51472..a0e14a5efa 100644
--- a/src/backend/storage/lmgr/condition_variable.c
+++ b/src/backend/storage/lmgr/condition_variable.c
@@ -121,9 +121,33 @@ ConditionVariablePrepareToSleep(ConditionVariable *cv)
  */
 void
 ConditionVariableSleep(ConditionVariable *cv, uint32 wait_event_info)
+{
+	(void) ConditionVariableTimedSleep(cv, -1 /* no timeout */, wait_event_info);
+}
+
+/*
+ * Wait for the given condition variable to be signaled or till timeout.
+ * This should be called in a predicate loop that tests for a specific exit
+ * condition and otherwise sleeps, like so:
+ *
+ *	 ConditionVariablePrepareToSleep(cv);  // optional
+ *	 while (condition for which we are waiting is not true)
+ *		 ConditionVariableSleep(cv, wait_event_info);
+ *	 ConditionVariableCancelSleep();
+ *
+ * wait_event_info should be a value from one of the WaitEventXXX enums
+ * defined in pgstat.h.  This controls the contents of pg_stat_activity's
+ * wait_event_type and wait_event columns while waiting.
+ *
+ * Returns 0 or -1 if timed out.
+ */
+int
+ConditionVariableTimedSleep(ConditionVariable *cv, long timeout,
+							uint32 wait_event_info)
 {
 	WaitEvent	event;
 	bool		done = false;
+	int			ret;
 
 	/*
 	 * If the caller didn't prepare to sleep explicitly, then do so now and
@@ -143,7 +167,7 @@ ConditionVariableSleep(ConditionVariable *cv, uint32 wait_event_info)
 	if (cv_sleep_target != cv)
 	{
 		ConditionVariablePrepareToSleep(cv);
-		return;
+		return 0;
 	}
 
 	do
@@ -154,12 +178,23 @@ ConditionVariableSleep(ConditionVariable *cv, uint32 wait_event_info)
 		 * Wait for latch to be set.  (If we're awakened for some other
 		 * reason, the code below will cope anyway.)
 		 */
-		(void) WaitEventSetWait(cv_wait_event_set, -1, &event, 1,
+		ret = WaitEventSetWait(cv_wait_event_set, timeout, &event, 1,
 								wait_event_info);
 
 		/* Reset latch before examining the state of the wait list. */
 		ResetLatch(MyLatch);
 
+		/* Timeout */
+		if (ret == 0 && timeout >= 0)
+		{
+			/*
+			 * In the event of a timeout, we simply return and the caller
+			 * calls ConditionVariableCancelSleep to remove themselves from the
+			 * wait queue
+			 */
+			return -1;
+		}
+
 		/*
 		 * If this process has been taken out of the wait list, then we know
 		 * that it has been signaled by ConditionVariableSignal (or
@@ -183,6 +218,8 @@ ConditionVariableSleep(ConditionVariable *cv, uint32 wait_event_info)
 		}
 		SpinLockRelease(&cv->mutex);
 	} while (!done);
+
+	return 0;
 }
 
 /*
diff --git a/src/include/storage/condition_variable.h b/src/include/storage/condition_variable.h
index 2a0249392c..f62164e483 100644
--- a/src/include/storage/condition_variable.h
+++ b/src/include/storage/condition_variable.h
@@ -43,6 +43,8 @@ extern void ConditionVariableInit(ConditionVariable *cv);
  * the condition variable.
  */
 extern void ConditionVariableSleep(ConditionVariable *cv, uint32 wait_event_info);
+extern int ConditionVariableTimedSleep(ConditionVariable *cv, long timeout,
+										uint32 wait_event_info);
 extern void ConditionVariableCancelSleep(void);
 
 /*
-- 
2.16.5