1 files changed, 71 insertions, 27 deletions

diff --git a/src/backend/port/win32_latch.c b/src/backend/port/win32_latch.c
index 4bcf7b7a8f3..ef61b0184d1 100644
--- a/src/backend/port/win32_latch.c
+++ b/src/backend/port/win32_latch.c
@@ -23,6 +23,7 @@
 #include <unistd.h>
 
 #include "miscadmin.h"
+#include "postmaster/postmaster.h"
 #include "replication/walsender.h"
 #include "storage/latch.h"
 #include "storage/shmem.h"
@@ -81,43 +82,67 @@ DisownLatch(volatile Latch *latch)
 	latch->owner_pid = 0;
 }
 
-bool
-WaitLatch(volatile Latch *latch, long timeout)
+int
+WaitLatch(volatile Latch *latch, int wakeEvents, long timeout)
 {
-	return WaitLatchOrSocket(latch, PGINVALID_SOCKET, false, false, timeout) > 0;
+	return WaitLatchOrSocket(latch, wakeEvents, PGINVALID_SOCKET, timeout);
 }
 
 int
-WaitLatchOrSocket(volatile Latch *latch, SOCKET sock, bool forRead,
-				  bool forWrite, long timeout)
+WaitLatchOrSocket(volatile Latch *latch, int wakeEvents, SOCKET sock,
+				  long timeout)
 {
 	DWORD		rc;
-	HANDLE		events[3];
+	HANDLE		events[4];
 	HANDLE		latchevent;
-	HANDLE		sockevent = WSA_INVALID_EVENT;	/* silence compiler */
+	HANDLE		sockevent = WSA_INVALID_EVENT;
 	int			numevents;
 	int			result = 0;
+	int			pmdeath_eventno;
+	long		timeout_ms;
+
+	Assert(wakeEvents != 0);
+
+	/* Ignore WL_SOCKET_* events if no valid socket is given */
+	if (sock == PGINVALID_SOCKET)
+		wakeEvents &= ~(WL_SOCKET_READABLE | WL_SOCKET_WRITEABLE);
+
+	/* Convert timeout to milliseconds for WaitForMultipleObjects() */
+	if (wakeEvents & WL_TIMEOUT)
+	{
+		Assert(timeout >= 0);
+		timeout_ms = timeout / 1000;
+	}
+	else
+		timeout_ms = INFINITE;
 
+	/* Construct an array of event handles for WaitforMultipleObjects() */
 	latchevent = latch->event;
 
 	events[0] = latchevent;
 	events[1] = pgwin32_signal_event;
 	numevents = 2;
-	if (sock != PGINVALID_SOCKET && (forRead || forWrite))
+	if (((wakeEvents & WL_SOCKET_READABLE) ||
+		 (wakeEvents & WL_SOCKET_WRITEABLE)))
 	{
 		int			flags = 0;
 
-		if (forRead)
+		if (wakeEvents & WL_SOCKET_READABLE)
 			flags |= FD_READ;
-		if (forWrite)
+		if (wakeEvents & WL_SOCKET_WRITEABLE)
 			flags |= FD_WRITE;
 
 		sockevent = WSACreateEvent();
 		WSAEventSelect(sock, sockevent, flags);
 		events[numevents++] = sockevent;
 	}
+	if (wakeEvents & WL_POSTMASTER_DEATH)
+	{
+		pmdeath_eventno = numevents;
+		events[numevents++] = PostmasterHandle;
+	}
 
-	for (;;)
+	do
 	{
 		/*
 		 * Reset the event, and check if the latch is set already. If someone
@@ -127,45 +152,64 @@ WaitLatchOrSocket(volatile Latch *latch, SOCKET sock, bool forRead,
 		 */
 		if (!ResetEvent(latchevent))
 			elog(ERROR, "ResetEvent failed: error code %d", (int) GetLastError());
-		if (latch->is_set)
+		if (latch->is_set && (wakeEvents & WL_LATCH_SET))
 		{
-			result = 1;
+			result |= WL_LATCH_SET;
+			/*
+			 * Leave loop immediately, avoid blocking again. We don't attempt
+			 * to report any other events that might also be satisfied.
+			 */
 			break;
 		}
 
-		rc = WaitForMultipleObjects(numevents, events, FALSE,
-							   (timeout >= 0) ? (timeout / 1000) : INFINITE);
+		rc = WaitForMultipleObjects(numevents, events, FALSE, timeout_ms);
+
 		if (rc == WAIT_FAILED)
 			elog(ERROR, "WaitForMultipleObjects() failed: error code %d", (int) GetLastError());
+
+		/* Participate in Windows signal emulation */
+		else if (rc == WAIT_OBJECT_0 + 1)
+			pgwin32_dispatch_queued_signals();
+
+		else if ((wakeEvents & WL_POSTMASTER_DEATH) &&
+			rc == WAIT_OBJECT_0 + pmdeath_eventno)
+		{
+			/* Postmaster died */
+			result |= WL_POSTMASTER_DEATH;
+		}
 		else if (rc == WAIT_TIMEOUT)
 		{
-			result = 0;
-			break;
+			result |= WL_TIMEOUT;
 		}
-		else if (rc == WAIT_OBJECT_0 + 1)
-			pgwin32_dispatch_queued_signals();
-		else if (rc == WAIT_OBJECT_0 + 2)
+		else if ((wakeEvents & (WL_SOCKET_READABLE | WL_SOCKET_WRITEABLE)) != 0 &&
+				 rc == WAIT_OBJECT_0 + 2)	/* socket is at event slot 2 */
 		{
 			WSANETWORKEVENTS resEvents;
 
-			Assert(sock != PGINVALID_SOCKET);
-
 			ZeroMemory(&resEvents, sizeof(resEvents));
 			if (WSAEnumNetworkEvents(sock, sockevent, &resEvents) == SOCKET_ERROR)
 				ereport(FATAL,
 						(errmsg_internal("failed to enumerate network events: %i", (int) GetLastError())));
 
-			if ((forRead && resEvents.lNetworkEvents & FD_READ) ||
-				(forWrite && resEvents.lNetworkEvents & FD_WRITE))
-				result = 2;
-			break;
+			if ((wakeEvents & WL_SOCKET_READABLE) &&
+				(resEvents.lNetworkEvents & FD_READ))
+			{
+				result |= WL_SOCKET_READABLE;
+			}
+			if ((wakeEvents & WL_SOCKET_WRITEABLE) &&
+				(resEvents.lNetworkEvents & FD_WRITE))
+			{
+				result |= WL_SOCKET_WRITEABLE;
+			}
 		}
+		/* Otherwise it must be the latch event */
 		else if (rc != WAIT_OBJECT_0)
 			elog(ERROR, "unexpected return code from WaitForMultipleObjects(): %d", (int) rc);
 	}
+	while(result == 0);
 
 	/* Clean up the handle we created for the socket */
-	if (sock != PGINVALID_SOCKET && (forRead || forWrite))
+	if (sockevent != WSA_INVALID_EVENT)
 	{
 		WSAEventSelect(sock, sockevent, 0);
 		WSACloseEvent(sockevent);