From 2222901c1947bb4c72f26b5706f5f3248cec96b8 Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Sun, 30 Jan 2022 13:01:13 +1300
Subject: [PATCH 3/3] Switch to per-socket Windows event handles.

Previously, each WaitEventSet would create a new operating system
'event' object for each socket.  Create a dedicated event and some state
flags for each socket, independent of WaitEventSets.  This fixes three
problems:

1.  When a socket is used in a new WaitEventSet (for example, one after
another when using temporary WaitEventSets), queued but not yet reported
FD_CLOSE events could be lost by the operating system.  This is now
avoided with a single long-lived event for the socket.

2.  We need to report Winsock's FD_CLOSE event as WL_SOCKET_READABLE
with level-triggered semantics, but it's only reported once by Winsock
on graceful socket shutdown, so could be lost if a caller waits twice in
a row without reading the socket.  This is avoided by adding a flag to
remember that we've already received FD_CLOSE.

3.  It should also be possible to use a socket in two WaitEventSet
objects at the same time (though not from two threads without more
work), since the network event mask will be adjusted on each wait if
required when you switch between sets containing the same socket.
Though this isn't currently needed by existing PostgreSQL code AFAIK,
it's a nice restriction to lift, to match Unix.

Manage the new extra per-socket state in a hash table indexed by socket,
and require clients of WaitEventSet to register and deregister sockets
with a new API.  On Unix this is a no-op, other than correctness checks
in assertion builds.  For libwalreceiver and postgres_fdw, a new
libpq-events callback automatically manages socket registration.

XXX Proof of concept

XXX There is some more temporary event-switching going on inside
src/backend/port/win32/socket.c, but not in code paths that were
affected by the recent graceful shutdown brokenness.  They may need to
be integrated with this, potentially widening the set of sockets that
need to be registered with socket_table.c to 'all of them'?

Discussion: https://postgr.es/m/CA%2BhUKG%2BOeoETZQ%3DQw5Ub5h3tmwQhBmDA%3DnuNO3KG%3DzWfUypFAw%40mail.gmail.com
---
 contrib/postgres_fdw/connection.c             |  25 +++
 src/backend/libpq/pqcomm.c                    |   3 +
 src/backend/port/Makefile                     |   3 +-
 src/backend/port/socket_table.c               | 160 ++++++++++++++++++
 src/backend/postmaster/pgstat.c               |  38 +++--
 src/backend/postmaster/syslogger.c            |   3 +
 .../libpqwalreceiver/libpqwalreceiver.c       |  33 ++++
 src/backend/storage/ipc/latch.c               |  73 ++++----
 src/include/port.h                            |  18 ++
 src/include/storage/latch.h                   |   2 +-
 10 files changed, 302 insertions(+), 56 deletions(-)
 create mode 100644 src/backend/port/socket_table.c

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 29fcb6a76e..783b9a3a74 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -17,6 +17,7 @@
 #include "catalog/pg_user_mapping.h"
 #include "commands/defrem.h"
 #include "funcapi.h"
+#include "libpq-events.h"
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
 #include "pgstat.h"
@@ -109,6 +110,7 @@ static void pgfdw_abort_cleanup(ConnCacheEntry *entry, const char *sql,
 								bool toplevel);
 static bool UserMappingPasswordRequired(UserMapping *user);
 static bool disconnect_cached_connections(Oid serverid);
+static int	pgfdw_eventproc(PGEventId evtId, void *evtInfo, void *passThrough);
 
 /*
  * Get a PGconn which can be used to execute queries on the remote PostgreSQL
@@ -469,6 +471,10 @@ connect_pg_server(ForeignServer *server, UserMapping *user)
 							server->servername),
 					 errdetail_internal("%s", pchomp(PQerrorMessage(conn)))));
 
+		/* Make sure that all sockets are registered for use by latch.c. */
+		if (!PQregisterEventProc(conn, pgfdw_eventproc, "postgres_fdw", NULL))
+			elog(ERROR, "postgres_fdw could not register for libpq events");
+
 		/*
 		 * Check that non-superuser has used password to establish connection;
 		 * otherwise, he's piggybacking on the postgres server's user
@@ -1709,3 +1715,22 @@ disconnect_cached_connections(Oid serverid)
 
 	return result;
 }
+
+/*
+ * A callback for libpq to allow its sockets to be used in WaitEventSet on Windows.
+ */
+static int
+pgfdw_eventproc(PGEventId evtId, void *evtInfo, void *passThrough)
+{
+	if (evtId == PGEVT_SOCKET)
+	{
+		PGEventSocket *evt = (PGEventSocket *) evtInfo;
+		return SocketTableAdd(evt->socket, true) ? 1 : 0;
+	}
+	else if (evtId == PGEVT_SOCKETCLOSE)
+	{
+		PGEventSocket *evt = (PGEventSocket *) evtInfo;
+		SocketTableDrop(evt->socket);
+	}
+	return 1;
+}
diff --git a/src/backend/libpq/pqcomm.c b/src/backend/libpq/pqcomm.c
index f05723dc92..bb35618af4 100644
--- a/src/backend/libpq/pqcomm.c
+++ b/src/backend/libpq/pqcomm.c
@@ -204,6 +204,9 @@ pq_init(void)
 				(errmsg("could not set socket to nonblocking mode: %m")));
 #endif
 
+	/* Make sure the socket can be used in a WaitEventSet on Windows. */
+	SocketTableAdd(MyProcPort->sock, false);
+
 	FeBeWaitSet = CreateWaitEventSet(TopMemoryContext, 3);
 	socket_pos = AddWaitEventToSet(FeBeWaitSet, WL_SOCKET_WRITEABLE,
 								   MyProcPort->sock, NULL, NULL);
diff --git a/src/backend/port/Makefile b/src/backend/port/Makefile
index 2d00b4f05a..03df581e06 100644
--- a/src/backend/port/Makefile
+++ b/src/backend/port/Makefile
@@ -25,7 +25,8 @@ OBJS = \
 	$(TAS) \
 	atomics.o \
 	pg_sema.o \
-	pg_shmem.o
+	pg_shmem.o\
+	socket_table.o
 
 ifeq ($(PORTNAME), win32)
 SUBDIRS += win32
diff --git a/src/backend/port/socket_table.c b/src/backend/port/socket_table.c
new file mode 100644
index 0000000000..c2bbc2e0cb
--- /dev/null
+++ b/src/backend/port/socket_table.c
@@ -0,0 +1,160 @@
+/*-------------------------------------------------------------------------
+ *
+ * socket_table.c
+ *	  Routines for dealing with extra per-socket state on Windows.
+ *
+ * Windows sockets can only safely be associate with one 'event', or they risk
+ * losing FD_CLOSE events.  FD_CLOSE is also edge-triggered and not
+ * resettable, so we need space for state to make it level-triggered.
+ *
+ * These functions are no-ops on Unix systems, except in assertion builds
+ * where we do the minimum book keeping required to report failure to register
+ * sockets appropriately.
+ *
+ * Portions Copyright (c) 1996-2022, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  src/backend/port/socket_table.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "common/hashfn.h"
+
+typedef struct SocketTableEntry
+{
+	pgsocket	sock;
+	int			status;
+	ExtraSocketState *extra;
+} SocketTableEntry;
+
+#define SH_PREFIX		socktab
+#define SH_ELEMENT_TYPE SocketTableEntry
+#define SH_KEY_TYPE		pgsocket
+#define SH_KEY			sock
+#define SH_HASH_KEY(tb, key)	hash_bytes((uint8 *) &(key), sizeof(key))
+#define SH_EQUAL(tb, a, b)		((a) == (b))
+#define SH_SCOPE static inline
+#define SH_NO_OOM
+#define SH_DECLARE
+#define SH_DEFINE
+#include "lib/simplehash.h"
+
+#if defined(WIN32) || defined(USE_ASSERT_CHECKING)
+static socktab_hash *socket_table;
+#endif
+
+/*
+ * Must be called exactly once for each socket before including it in code
+ * that requires ExtraSocketState, such as WaitEventSet.  The socket must be
+ * not already registered, and must be dropped before closing it.  Return NULL
+ * on out-of-memory, if no_oom is true, otherwise raises errors.
+ */
+ExtraSocketState *
+SocketTableAdd(pgsocket sock, bool no_oom)
+{
+#if defined(WIN32) || defined(USE_ASSERT_CHECKING)
+	SocketTableEntry *ste;
+	ExtraSocketState *ess;
+	bool		found;
+
+	/*
+	 * Create a new SocketTableEntry object.  We can't use a pointer into the
+	 * hash table because it would move around, and we want WaitEventSet to be
+	 * able to hold pointers.
+	 */
+	if (!(ess = malloc(sizeof(*ess))))
+		goto out_of_memory;
+	memset(ess, 0, sizeof(*ess));
+
+#ifdef WIN32
+	/* Set up Windows kernel event for this socket. */
+	ess->event_handle = WSACreateEvent();
+	if (ess->event_handle == WSA_INVALID_EVENT)
+		goto out_of_memory;
+#endif
+
+	/* Create socket_table on demand. */
+	/* XXX dedicated memory context for debugging? */
+	if (!socket_table &&
+		(!(socket_table = socktab_create(TopMemoryContext, 128, NULL))))
+		goto out_of_memory;
+
+	/* Try to insert. */
+	ste = socktab_insert(socket_table, sock, &found);
+	if (!ste)
+		goto out_of_memory;
+	if (found)
+	{
+#ifdef WIN32
+		WSACloseEvent(ess->event_handle);
+#endif
+		free(ess);
+		elog(ERROR, "cannot register socket, already registered");
+	}
+
+	/* Success. */
+	ste->extra = ess;
+	return ess;
+
+ out_of_memory:
+	if (ess)
+	{
+#ifdef WIN32
+		if (ess->event_handle != WSA_INVALID_EVENT)
+			WSACloseEvent(ess->event_handle);
+#endif
+		free(ess);
+	}
+	if (!no_oom)
+		elog(ERROR, "out of memory");
+	return NULL;
+#else
+	return NULL;
+#endif
+}
+
+/*
+ * Unregister a socket that was registered.  This must be done before closing
+ * the socket.
+ */
+void
+SocketTableDrop(pgsocket sock)
+{
+#if defined(WIN32) || defined(USE_ASSERT_CHECKING)
+	SocketTableEntry *ste = NULL;
+	ExtraSocketState *ess;
+
+	if (socket_table)
+		ste = socktab_lookup(socket_table, sock);
+	if (!ste)
+		elog(ERROR, "cannot unregister socket that is not registered");
+	ess = ste->extra;
+
+#ifdef WIN32
+	WSACloseEvent(ess->event_handle);
+#endif
+
+	free(ess);
+	socktab_delete_item(socket_table, ste);
+#endif
+}
+
+ExtraSocketState *
+SocketTableGet(pgsocket sock)
+{
+#if defined(WIN32) || defined(USE_ASSERT_CHECKING)
+	SocketTableEntry *ste = NULL;
+
+	if (socket_table)
+		ste = socktab_lookup(socket_table, sock);
+	if (!ste)
+		elog(ERROR, "socket is not registered");
+	return ste->extra;
+#else
+	return NULL;
+#endif
+}
diff --git a/src/backend/postmaster/pgstat.c b/src/backend/postmaster/pgstat.c
index 0646f53098..dd95353880 100644
--- a/src/backend/postmaster/pgstat.c
+++ b/src/backend/postmaster/pgstat.c
@@ -382,6 +382,7 @@ static void pgstat_recv_replslot(PgStat_MsgReplSlot *msg, int len);
 static void pgstat_recv_tempfile(PgStat_MsgTempFile *msg, int len);
 static void pgstat_recv_subscription_purge(PgStat_MsgSubscriptionPurge *msg, int len);
 static void pgstat_recv_subworker_error(PgStat_MsgSubWorkerError *msg, int len);
+static void pgstat_release_socket(void);
 
 /* ------------------------------------------------------------
  * Public functions called from postmaster follow
@@ -483,8 +484,7 @@ pgstat_init(void)
 			ereport(LOG,
 					(errcode_for_socket_access(),
 					 errmsg("could not bind socket for statistics collector: %m")));
-			closesocket(pgStatSock);
-			pgStatSock = PGINVALID_SOCKET;
+			pgstat_release_socket();
 			continue;
 		}
 
@@ -494,8 +494,7 @@ pgstat_init(void)
 			ereport(LOG,
 					(errcode_for_socket_access(),
 					 errmsg("could not get address of socket for statistics collector: %m")));
-			closesocket(pgStatSock);
-			pgStatSock = PGINVALID_SOCKET;
+			pgstat_release_socket();
 			continue;
 		}
 
@@ -510,8 +509,7 @@ pgstat_init(void)
 			ereport(LOG,
 					(errcode_for_socket_access(),
 					 errmsg("could not connect socket for statistics collector: %m")));
-			closesocket(pgStatSock);
-			pgStatSock = PGINVALID_SOCKET;
+			pgstat_release_socket();
 			continue;
 		}
 
@@ -531,8 +529,7 @@ retry1:
 			ereport(LOG,
 					(errcode_for_socket_access(),
 					 errmsg("could not send test message on socket for statistics collector: %m")));
-			closesocket(pgStatSock);
-			pgStatSock = PGINVALID_SOCKET;
+			pgstat_release_socket();
 			continue;
 		}
 
@@ -557,8 +554,7 @@ retry1:
 			ereport(LOG,
 					(errcode_for_socket_access(),
 					 errmsg("select() failed in statistics collector: %m")));
-			closesocket(pgStatSock);
-			pgStatSock = PGINVALID_SOCKET;
+			pgstat_release_socket();
 			continue;
 		}
 		if (sel_res == 0 || !FD_ISSET(pgStatSock, &rset))
@@ -572,8 +568,7 @@ retry1:
 			ereport(LOG,
 					(errcode(ERRCODE_CONNECTION_FAILURE),
 					 errmsg("test message did not get through on socket for statistics collector")));
-			closesocket(pgStatSock);
-			pgStatSock = PGINVALID_SOCKET;
+			pgstat_release_socket();
 			continue;
 		}
 
@@ -587,8 +582,7 @@ retry2:
 			ereport(LOG,
 					(errcode_for_socket_access(),
 					 errmsg("could not receive test message on socket for statistics collector: %m")));
-			closesocket(pgStatSock);
-			pgStatSock = PGINVALID_SOCKET;
+			pgstat_release_socket();
 			continue;
 		}
 
@@ -597,8 +591,7 @@ retry2:
 			ereport(LOG,
 					(errcode(ERRCODE_INTERNAL_ERROR),
 					 errmsg("incorrect test message transmission on socket for statistics collector")));
-			closesocket(pgStatSock);
-			pgStatSock = PGINVALID_SOCKET;
+			pgstat_release_socket();
 			continue;
 		}
 
@@ -669,7 +662,7 @@ startup_failed:
 		pg_freeaddrinfo_all(hints.ai_family, addrs);
 
 	if (pgStatSock != PGINVALID_SOCKET)
-		closesocket(pgStatSock);
+		pgstat_release_socket();
 	pgStatSock = PGINVALID_SOCKET;
 
 	/*
@@ -3527,6 +3520,9 @@ PgstatCollectorMain(int argc, char *argv[])
 	pgStatRunningInCollector = true;
 	pgStatDBHash = pgstat_read_statsfiles(InvalidOid, true, true);
 
+	/* Make sure pgStatSock can be used in a WaitEventSet on Windows. */
+	SocketTableAdd(pgStatSock, false);
+
 	/* Prepare to wait for our latch or data in our socket. */
 	wes = CreateWaitEventSet(CurrentMemoryContext, 3);
 	AddWaitEventToSet(wes, WL_LATCH_SET, PGINVALID_SOCKET, MyLatch, NULL);
@@ -6418,3 +6414,11 @@ pgstat_count_slru_truncate(int slru_idx)
 {
 	slru_entry(slru_idx)->m_truncate += 1;
 }
+
+static void
+pgstat_release_socket(void)
+{
+	SocketTableDrop(pgStatSock);
+	closesocket(pgStatSock);
+	pgStatSock = PGINVALID_SOCKET;
+}
diff --git a/src/backend/postmaster/syslogger.c b/src/backend/postmaster/syslogger.c
index 25e2131e31..aa551b2719 100644
--- a/src/backend/postmaster/syslogger.c
+++ b/src/backend/postmaster/syslogger.c
@@ -301,6 +301,9 @@ SysLoggerMain(int argc, char *argv[])
 	 */
 	whereToSendOutput = DestNone;
 
+	/* Make sure the pipe can be used in a WaitEventSet on Windows. */
+	SocketTableAdd((pgsocket) syslogPipe[0], false);
+
 	/*
 	 * Set up a reusable WaitEventSet object we'll use to wait for our latch,
 	 * and (except on Windows) our socket.
diff --git a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
index 0d89db4e6a..8508f93aa6 100644
--- a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
+++ b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
@@ -23,6 +23,7 @@
 #include "catalog/pg_type.h"
 #include "common/connect.h"
 #include "funcapi.h"
+#include "libpq-events.h"
 #include "libpq-fe.h"
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
@@ -83,6 +84,10 @@ static WalRcvExecResult *libpqrcv_exec(WalReceiverConn *conn,
 									   const Oid *retTypes);
 static void libpqrcv_disconnect(WalReceiverConn *conn);
 
+static int libpqrcv_eventproc(PGEventId evtId,
+							  void *evtInfo,
+							  void *passThrough);
+
 static WalReceiverFunctionsType PQWalReceiverFunctions = {
 	libpqrcv_connect,
 	libpqrcv_check_conninfo,
@@ -182,6 +187,15 @@ libpqrcv_connect(const char *conninfo, bool logical, const char *appname,
 		return NULL;
 	}
 
+	if (PQregisterEventProc(conn->streamConn, libpqrcv_eventproc,
+							"libwalrcv", NULL) == 0)
+	{
+		PQfinish(conn->streamConn);
+		conn->streamConn = NULL;
+		*err = pstrdup("could not register for libpq events");
+		return NULL;
+	}
+
 	/*
 	 * Poll connection until we have OK or FAILED status.
 	 *
@@ -1165,3 +1179,22 @@ stringlist_to_identifierstr(PGconn *conn, List *strings)
 
 	return res.data;
 }
+
+/*
+ * A callback for libpq to allow its sockets to be used in WaitEventSet on Windows.
+ */
+static int
+libpqrcv_eventproc(PGEventId evtId, void *evtInfo, void *passThrough)
+{
+	if (evtId == PGEVT_SOCKET)
+	{
+		PGEventSocket *evt = (PGEventSocket *) evtInfo;
+		return SocketTableAdd(evt->socket, true) ? 1 : 0;
+	}
+	else if (evtId == PGEVT_SOCKETCLOSE)
+	{
+		PGEventSocket *evt = (PGEventSocket *) evtInfo;
+		SocketTableDrop(evt->socket);
+	}
+	return 1;
+}
diff --git a/src/backend/storage/ipc/latch.c b/src/backend/storage/ipc/latch.c
index 61c876beff..5a6fb36fa2 100644
--- a/src/backend/storage/ipc/latch.c
+++ b/src/backend/storage/ipc/latch.c
@@ -60,6 +60,7 @@
 #include "storage/shmem.h"
 #include "utils/memutils.h"
 
+
 /*
  * Select the fd readiness primitive to use. Normally the "most modern"
  * primitive supported by the OS will be used, but for testing it can be
@@ -803,28 +804,6 @@ FreeWaitEventSet(WaitEventSet *set)
 #elif defined(WAIT_USE_KQUEUE)
 	close(set->kqueue_fd);
 	ReleaseExternalFD();
-#elif defined(WAIT_USE_WIN32)
-	WaitEvent  *cur_event;
-
-	for (cur_event = set->events;
-		 cur_event < (set->events + set->nevents);
-		 cur_event++)
-	{
-		if (cur_event->events & WL_LATCH_SET)
-		{
-			/* uses the latch's HANDLE */
-		}
-		else if (cur_event->events & WL_POSTMASTER_DEATH)
-		{
-			/* uses PostmasterHandle */
-		}
-		else
-		{
-			/* Clean up the event object we created for the socket */
-			WSAEventSelect(cur_event->fd, NULL, 0);
-			WSACloseEvent(set->handles[cur_event->pos + 1]);
-		}
-	}
 #endif
 
 	pfree(set);
@@ -897,9 +876,17 @@ AddWaitEventToSet(WaitEventSet *set, uint32 events, pgsocket fd, Latch *latch,
 	event->fd = fd;
 	event->events = events;
 	event->user_data = user_data;
+
+	if (events & WL_SOCKET_MASK)
+	{
 #ifdef WIN32
-	event->reset = false;
+		/* Point to our Windows event tracking state for the socket. */
+		event->extra_socket_state = SocketTableGet(fd);
+#else
+		/* On Unix, check it's registered to avoid portability bugs. */
+		(void) SocketTableGet(fd);
 #endif
+	}
 
 	if (events == WL_LATCH_SET)
 	{
@@ -1265,6 +1252,7 @@ WaitEventAdjustWin32(WaitEventSet *set, WaitEvent *event)
 	}
 	else
 	{
+		ExtraSocketState *ess;
 		int			flags = FD_CLOSE;	/* always check for errors/EOF */
 
 		if (event->events & WL_SOCKET_READABLE)
@@ -1274,18 +1262,19 @@ WaitEventAdjustWin32(WaitEventSet *set, WaitEvent *event)
 		if (event->events & WL_SOCKET_CONNECTED)
 			flags |= FD_CONNECT;
 
-		if (*handle == WSA_INVALID_EVENT)
-		{
-			*handle = WSACreateEvent();
-			if (*handle == WSA_INVALID_EVENT)
-				elog(ERROR, "failed to create event for socket: error code %d",
-					 WSAGetLastError());
-		}
-		if (WSAEventSelect(event->fd, *handle, flags) != 0)
+		/* Find the event associated with this socket. */
+		ess = event->extra_socket_state;
+		*handle = ess->event_handle;
+
+		/*
+		 * Adjust the flags for this event, if they don't already match what we
+		 * want to wait for.
+		 */
+		if (ess->flags != flags &&
+			WSAEventSelect(event->fd, *handle, flags) != 0)
 			elog(ERROR, "failed to set up event for socket: error code %d",
 				 WSAGetLastError());
-
-		Assert(event->fd != PGINVALID_SOCKET);
+		ess->flags = flags;
 	}
 }
 #endif
@@ -1845,10 +1834,19 @@ WaitEventSetWaitBlock(WaitEventSet *set, int cur_timeout,
 		 cur_event < (set->events + set->nevents);
 		 cur_event++)
 	{
-		if (cur_event->reset)
-		{
+		/* Reset the event mask if necessary. */
+		if (cur_event->events & WL_SOCKET_MASK)
 			WaitEventAdjustWin32(set, cur_event);
-			cur_event->reset = false;
+
+		/* If we've already seen FD_CLOSE, keep reporting readiness. */
+		if ((cur_event->events & WL_SOCKET_MASK) &&
+			cur_event->extra_socket_state->seen_fd_close)
+		{
+			occurred_events->pos = cur_event->pos;
+			occurred_events->user_data = cur_event->user_data;
+			occurred_events->events = cur_event->events;
+			occurred_events->fd = cur_event->fd;
+			return 1;
 		}
 
 		/*
@@ -1984,7 +1982,7 @@ WaitEventSetWaitBlock(WaitEventSet *set, int cur_timeout,
 			 * for the behavior of socket events.
 			 *------
 			 */
-			cur_event->reset = true;
+			cur_event->extra_socket_state->flags = 0;
 		}
 		if ((cur_event->events & WL_SOCKET_WRITEABLE) &&
 			(resEvents.lNetworkEvents & FD_WRITE))
@@ -2002,6 +2000,7 @@ WaitEventSetWaitBlock(WaitEventSet *set, int cur_timeout,
 		{
 			/* EOF/error, so signal all caller-requested socket flags */
 			occurred_events->events |= (cur_event->events & WL_SOCKET_MASK);
+			cur_event->extra_socket_state->seen_fd_close = true;
 		}
 
 		if (occurred_events->events != 0)
diff --git a/src/include/port.h b/src/include/port.h
index 3d103a2b31..e780852bd1 100644
--- a/src/include/port.h
+++ b/src/include/port.h
@@ -543,4 +543,22 @@ extern char *wait_result_to_str(int exit_status);
 extern bool wait_result_is_signal(int exit_status, int signum);
 extern bool wait_result_is_any_signal(int exit_status, bool include_command_not_found);
 
+/* backend/port/socket_table.c */
+#if !defined(FRONTEND)
+struct ExtraSocketState
+{
+#ifdef WIN32
+	HANDLE		event_handle;		/* one event for the life of the socket */
+	int			flags;				/* most recent WSAEventSelect() flags */
+	bool		seen_fd_close;		/* has FD_CLOSE been received? */
+#else
+	int			dummy;				/* none of this is needed for Unix */
+#endif
+};
+typedef struct ExtraSocketState ExtraSocketState;
+extern ExtraSocketState *SocketTableAdd(pgsocket sock, bool no_oom);
+extern ExtraSocketState *SocketTableGet(pgsocket sock);
+extern void SocketTableDrop(pgsocket sock);
+#endif
+
 #endif							/* PG_PORT_H */
diff --git a/src/include/storage/latch.h b/src/include/storage/latch.h
index 3aa7b33834..576d339829 100644
--- a/src/include/storage/latch.h
+++ b/src/include/storage/latch.h
@@ -146,7 +146,7 @@ typedef struct WaitEvent
 	pgsocket	fd;				/* socket fd associated with event */
 	void	   *user_data;		/* pointer provided in AddWaitEventToSet */
 #ifdef WIN32
-	bool		reset;			/* Is reset of the event required? */
+	ExtraSocketState *extra_socket_state;
 #endif
 } WaitEvent;
 
-- 
2.33.1

