Add multi-waiter support for TCP accept

include/NetAPI.h

@@ -34,6 +34,29 @@ struct NetworkAddress
 	} kind;
 };
 
+/**
+ * Enumeration that defines futex types. Each value corresponds
+ * to an index in the socket's futex array.
+ */
+enum SocketEventType : uint8_t
+{
+	/// Triggered when a new TCP connection
+	/// is accepted on a listening socket.
+	/// The value of the futex corresponding to
+	/// SocketAcceptEvent, means the number of
+	/// pending connections.
+	SocketAcceptEvent = 0,
+	// SocketReceiveEvent = 1, // Triggered when data is received on a socket
+	// SocketSendEvent = 2     // Triggered when a socket has space available
+	// for sending
+};
+
+/// Number of distinct socket event futex types.
+static constexpr size_t NumFutexTypes = SocketEventType::SocketAcceptEvent + 1;
+/// Sentinel value stored in a socket event futex after the socket has been
+/// torn down.
+static constexpr int32_t SocketNotAvailable = INT32_MIN;
+
 /**
  * Enumeration defining the connection type.
  */
@@ -269,6 +292,14 @@ Socket __cheri_compartment("TCPIP")
                      AllocatorCapability mallocCapability,
                      bool                isIPv6);
 
+/**
+ * Return the event source associated with a socket.
+ *
+ * The returned capability is read-only and bounded to four bytes.
+ */
+uint32_t *__cheri_compartment("TCPIP")
+  network_socket_get_event_source(Socket sealedSocket, SocketEventType type);
+
 /**
  * Authorise a UDP socket to send packets to a specific host.  This opens a
  * firewall hole allowing the socket to send and receive packets to the host.

lib/tcpip/network_wrapper.cc

@@ -412,8 +412,50 @@ namespace
 
 		return ret;
 	}
+
 } // namespace
 
+int SealedSocket::signal_event_futex(SocketEventType type)
+{
+	auto &futex = eventFutexState[type];
+	if (heap_claim_ephemeral(TimeoutWaitForever, &futex) != 0)
+	{
+		return -EINVAL;
+	}
+	int32_t current = futex.load();
+	while (current != SocketNotAvailable)
+	{
+		/*
+		 * If the futex's current value equals current,
+		 * store current + 1 and return true. Otherwise
+		 * write the actual current value into current
+		 * (by reference) and return false.
+		 */
+		if (futex.compare_exchange_strong(current, current + 1))
+		{
+			futex.notify_all();
+			return 0;
+		}
+	}
+
+	return -EINVAL;
+}
+
+int SealedSocket::consume_event_futex(SocketEventType type)
+{
+	auto   &futex   = eventFutexState[type];
+	int32_t current = futex.load();
+	while (current != SocketNotAvailable)
+	{
+		if (futex.compare_exchange_strong(current, current - 1))
+		{
+			return 0;
+		}
+	}
+
+	return -EINVAL;
+}
+
 /**
  * Callback called by FreeRTOS+TCP when a TCP connection is created or
  * terminated.
@@ -468,6 +510,24 @@ static void on_tcp_connect(Socket_t socket, BaseType_t isConnected)
 			  address.sin_address.ulIP_IPv4, localPort, address.sin_port);
 		}
 	}
+	else
+	{
+		// Update eventFutexState in the listening socket and wake up the
+		// threads sleeping on the corresponding futex.
+		// Use the FreeRTOS getter to get the wrapper capability we stored in
+		// it.
+		auto *wrapper =
+		  static_cast<SealedSocket *>(pvSocketGetSocketID(socket));
+		if (wrapper != nullptr)
+		{
+			/*
+			 * Ignore the return value. If the socket has been invalidated,
+			 * the subsequent accept call will detect that it is unavailable.
+			 */
+			wrapper->signal_event_futex(SocketEventType::SocketAcceptEvent);
+		}
+		return;
+	}
 }
 F_TCP_UDP_Handler_t onTCPConnectCallback = {on_tcp_connect};
 
@@ -495,6 +555,11 @@ Socket network_socket_create_and_bind(Timeout            *timeout,
 
 		  // Set the socket epoch
 		  socketWrapper->socketEpoch = currentSocketEpoch.load();
+		  // Multi-waiter: initialize the futexes.
+		  for (auto &eventState : socketWrapper->eventFutexState)
+		  {
+			  eventState.store(0);
+		  }
 
 		  const auto Family = isIPv6 ? FREERTOS_AF_INET6 : FREERTOS_AF_INET;
 		  Socket_t   socket =
@@ -512,7 +577,14 @@ Socket network_socket_create_and_bind(Timeout            *timeout,
 			  token_obj_destroy(mallocCapability, socket_key(), sealedSocket);
 			  return -ENOMEM;
 		  }
+
+		  /*
+		   * Create a bidirectional association between the FreeRTOS socket and
+		   * the sealed socket wrapper. This is used to retrieve the sealed
+		   * socket wrapper from the FreeRTOS socket in the callbacks.
+		   */
 		  socketWrapper->socket = socket;
+		  xSocketSetSocketID(socket, socketWrapper);
 
 		  // Claim the socket so that it counts towards the caller's quota.  The
 		  // network stack also keeps a claim to it.  We will drop this claim on
@@ -643,6 +715,13 @@ Socket network_socket_accept_tcp(Timeout            *timeout,
 		  }
 
 		  socketWrapper->socketEpoch = currentSocketEpoch.load();
+		  /*
+		   * For the newly allocated child socket, initialize the futexes.
+		   */
+		  for (auto &eventState : socketWrapper->eventFutexState)
+		  {
+			  eventState.store(0);
+		  }
 
 		  struct freertos_sockaddr addressTmp;
 		  uint32_t                 addressLength = sizeof(addressTmp);
@@ -673,6 +752,24 @@ Socket network_socket_accept_tcp(Timeout            *timeout,
 			  token_obj_destroy(mallocCapability, socket_key(), sealedSocket);
 			  return acceptResult;
 		  }
+		  /*
+		   * Note that here we update the futex, but no need to call
+		   * notify_all() on the waiting threads. Because the threads are
+		   * waiting on multiwaiter for arriving connected socket, but here the
+		   * semantics is not having an arrving connection. So there is no need
+		   * to wake the threads up, since they will probably find that there is
+		   * no available sockets and go back to sleep again.
+		   */
+		  int futexConsumeResult = listeningSocket->consume_event_futex(
+		    SocketEventType::SocketAcceptEvent);
+		  if (futexConsumeResult != 0)
+		  {
+			  // Return -EINVAL, which represents the
+			  // socket will be freed soon.
+			  close_socket_retry(timeout, rawSocket);
+			  token_obj_destroy(mallocCapability, socket_key(), sealedSocket);
+			  return futexConsumeResult;
+		  }
 		  socketWrapper->socket = rawSocket;
 
 		  // Claim the socket so that it counts towards the caller's quota.  The
@@ -708,6 +805,7 @@ Socket network_socket_accept_tcp(Timeout            *timeout,
 			  token_obj_destroy(mallocCapability, socket_key(), sealedSocket);
 			  return -EINVAL;
 		  }
+		  xSocketSetSocketID(rawSocket, socketWrapper);
 
 		  // Set `address`.
 		  if ((heap_claim_ephemeral(timeout, address) < 0) ||
@@ -831,6 +929,36 @@ Socket network_socket_udp(Timeout            *timeout,
 	  timeout, mallocCapability, isIPv6, ConnectionTypeUDP);
 }
 
+/**
+ * Getter to get the read only capability of a specific futex in
+ * specific socket. If the futex is invalid, return a untagged nullptr.
+ */
+uint32_t *network_socket_get_event_source(Socket          sealedSocket,
+                                          SocketEventType type)
+{
+	uint32_t *result = nullptr;
+	if (type < NumFutexTypes)
+	{
+		with_sealed_socket(
+		  [&](SealedSocket *socket) {
+			  auto *futex = &socket->eventFutexState[type];
+
+			  Capability readOnlyEventSource{futex};
+			  // Expose the futex as read-only. The caller may observe its raw
+			  // 32-bit representation but may modify it only through the
+			  // socket callbacks.
+			  readOnlyEventSource.bounds() = sizeof(*futex);
+			  readOnlyEventSource.permissions() &=
+			    {Permission::Load, Permission::Global};
+
+			  result = reinterpret_cast<uint32_t *>(readOnlyEventSource.get());
+			  return 0;
+		  },
+		  sealedSocket);
+	}
+	return result;
+}
+
 int network_socket_close(Timeout            *t,
                          AllocatorCapability mallocCapability,
                          Socket              sealedSocket)
@@ -893,6 +1021,12 @@ int network_socket_close(Timeout            *t,
 			  if (socketEpoch == currentSocketEpoch.load())
 			  {
 				  bool isTCP = rawSocket->ucProtocol == FREERTOS_IPPROTO_TCP;
+
+				  // Set the back pointer (pointing to the
+				  // wrapper) to nullptr, so that we will not
+				  // have a dangling pointer.
+				  xSocketSetSocketID(rawSocket, nullptr);
+
 				  // Nothing to do if `FreeRTOS_shutdown`
 				  // fails: this happens only if the TCP
 				  // connection is dead, which is likely to
@@ -1040,6 +1174,14 @@ int network_socket_close(Timeout            *t,
 					  // freed next time we try.
 					  return -ETIMEDOUT;
 				  }
+				  // Wake any thread waiting on this socket's event futex so
+				  // that it does not sleep forever on memory that is about to
+				  // be freed.
+				  for (auto &eventState : socket->eventFutexState)
+				  {
+					  eventState.store(SocketNotAvailable);
+					  eventState.notify_all();
+				  }
 
 				  g.release();
 				  socket->socketLock.upgrade_for_destruction();

lib/tcpip/tcpip-internal.h

@@ -3,6 +3,8 @@
 
 #pragma once
 #include <FreeRTOS_IP.h>
+#include <NetAPI.h>
+#include <atomic>
 #include <ds/linked_list.h>
 #include <function_wrapper.hh>
 #include <locks.hh>
@@ -42,6 +44,33 @@ struct SealedSocket
 	 * to the current instance of the network stack.
 	 */
 	uint64_t socketEpoch;
+	/**
+	 * Event waiter source futex array. This supports the multi-waiter
+	 * feature. Different events increment different futexes in the array and
+	 * wake the corresponding waiting threads.
+	 */
+	std::atomic<int32_t> eventFutexState[NumFutexTypes];
+	/**
+	 * Increments the futex and notifies all waiters if the futex is still
+	 * valid.
+	 *
+	 * Returns 0 on success, or `-EINVAL` if the futex has been invalidated.
+	 */
+	int signal_event_futex(SocketEventType type);
+	/**
+	 * Records one successfully accepted child socket by decrementing the accept
+	 * event counter.
+	 *
+	 * The caller must hold `socketLock`, which prevents the socket from being
+	 * freed while the futex is accessed. The futex value may become negative
+	 * if this helper is invoked after `FreeRTOS_accept()` successfully dequeue
+	 * a child, but before `on_tcp_connect()` increment the futex. The delayed
+	 * increment will repay this temporary debt.
+	 *
+	 * Returns 0 on success, or `-EINVAL` if the futex has been invalidated by
+	 * being set to `SocketNotAvailable`.
+	 */
+	int consume_event_futex(SocketEventType type);
 	/**
 	 * The lock protecting this socket.
 	 */

lib/tcpip/tcpip_error_handler.h

@@ -195,6 +195,13 @@ extern "C" void reset_network_stack_state(bool isIpThread)
 				DebugErrorHandler::log("Ignoring corrupted socket lock {}.",
 				                       lock);
 			}
+			// Notify all threads waiting on the socket's futexes so that they
+			// do not sleep forever on this old socket.
+			for (auto &eventState : socket->eventFutexState)
+			{
+				eventState.store(SocketNotAvailable);
+				eventState.notify_all();
+			}
 
 			FreeRTOS_Socket_t *s = socket->socket;
 			if (Capability{s}.is_valid() &&