Speed up mg_stop() by using a socket-pair to cause all mg_poll() calls to return immediately on server shutdown

src/civetweb.c

@@ -2438,6 +2438,9 @@ struct mg_context {
 	int lua_bg_log_available;     /* Use Lua background state for access log */
 #endif
 
+	int user_shutdown_notification_socket;   /* mg_stop() will close this socket... */
+	int thread_shutdown_notification_socket; /* to cause poll() in all threads to return immediately */
+
 	/* Server nonce */
 	pthread_mutex_t nonce_mutex; /* Protects ssl_ctx, handlers,
 	                              * ssl_cert_last_mtime, nonce_count, and
@@ -6143,12 +6146,15 @@ push_inner(struct mg_context *ctx,
 			mg_sleep(5);
 		} else {
 			/* For sockets, wait for the socket using poll */
-			struct mg_pollfd pfd[1];
+			struct mg_pollfd pfd[2];
 			int pollres;
 
 			pfd[0].fd = sock;
 			pfd[0].events = POLLOUT;
-			pollres = mg_poll(pfd, 1, (int)(ms_wait), &(ctx->stop_flag));
+
+			pfd[1].fd = ctx->thread_shutdown_notification_socket;
+			pfd[1].events = POLLIN;
+			pollres = mg_poll(pfd, 2, (int)(ms_wait), &(ctx->stop_flag));
 			if (!STOP_FLAG_IS_ZERO(&ctx->stop_flag)) {
 				return -2;
 			}
@@ -6256,7 +6262,7 @@ pull_inner(FILE *fp,
 
 #if defined(USE_MBEDTLS)
 	} else if (conn->ssl != NULL) {
-		struct mg_pollfd pfd[1];
+		struct mg_pollfd pfd[2];
 		int to_read;
 		int pollres;
 
@@ -6274,10 +6280,13 @@ pull_inner(FILE *fp,
 			pfd[0].fd = conn->client.sock;
 			pfd[0].events = POLLIN;
 
+			pfd[1].fd = conn->phys_ctx->thread_shutdown_notification_socket;
+			pfd[1].events = POLLIN;
+
 			to_read = len;
 
 			pollres = mg_poll(pfd,
-			                  1,
+			                  2,
 			                  (int)(timeout * 1000.0),
 			                  &(conn->phys_ctx->stop_flag));
 
@@ -6312,7 +6321,7 @@ pull_inner(FILE *fp,
 #elif !defined(NO_SSL)
 	} else if (conn->ssl != NULL) {
 		int ssl_pending;
-		struct mg_pollfd pfd[1];
+		struct mg_pollfd pfd[2];
 		int pollres;
 
 		if ((ssl_pending = SSL_pending(conn->ssl)) > 0) {
@@ -6326,8 +6335,11 @@ pull_inner(FILE *fp,
 		} else {
 			pfd[0].fd = conn->client.sock;
 			pfd[0].events = POLLIN;
+			pfd[1].fd = conn->phys_ctx->thread_shutdown_notification_socket;
+			pfd[1].events = POLLIN;
+
 			pollres = mg_poll(pfd,
-			                  1,
+			                  2,
 			                  (int)(timeout * 1000.0),
 			                  &(conn->phys_ctx->stop_flag));
 			if (!STOP_FLAG_IS_ZERO(&conn->phys_ctx->stop_flag)) {
@@ -6365,13 +6377,17 @@ pull_inner(FILE *fp,
 #endif
 
 	} else {
-		struct mg_pollfd pfd[1];
+		struct mg_pollfd pfd[2];
 		int pollres;
 
 		pfd[0].fd = conn->client.sock;
 		pfd[0].events = POLLIN;
+
+		pfd[1].fd = conn->phys_ctx->thread_shutdown_notification_socket;
+		pfd[1].events = POLLIN;
+
 		pollres = mg_poll(pfd,
-		                  1,
+		                  2,
 		                  (int)(timeout * 1000.0),
 		                  &(conn->phys_ctx->stop_flag));
 		if (!STOP_FLAG_IS_ZERO(&conn->phys_ctx->stop_flag)) {
@@ -9526,7 +9542,7 @@ connect_socket(
 #endif
 
 		/* Data for poll */
-		struct mg_pollfd pfd[1];
+		struct mg_pollfd pfd[2];
 		int pollres;
 		int ms_wait = 10000;     /* 10 second timeout */
 		stop_flag_t nonstop = 0; /* STOP_FLAG_ASSIGN(&nonstop, 0); */
@@ -9538,7 +9554,11 @@ connect_socket(
 		 */
 		pfd[0].fd = *sock;
 		pfd[0].events = POLLOUT;
-		pollres = mg_poll(pfd, 1, ms_wait, ctx ? &(ctx->stop_flag) : &nonstop);
+
+		pfd[1].fd = ctx ? ctx->thread_shutdown_notification_socket : -1;
+		pfd[1].events = POLLIN;
+
+		pollres = mg_poll(pfd, ctx ? 2 : 1, ms_wait, ctx ? &(ctx->stop_flag) : &nonstop);
 
 		if (pollres != 1) {
 			/* Not connected */
@@ -16028,9 +16048,13 @@ set_ports_option(struct mg_context *phys_ctx)
 			continue;
 		}
 
+		/* The +2 below includes the original +1 (for the socket we're about to add),
+		 * plus another +1 for the thread_shutdown_notification_socket that we'll
+		 * also want to poll() on so that mg_stop() can return quickly
+		 */
 		if ((pfd = (struct mg_pollfd *)
 		         mg_realloc_ctx(phys_ctx->listening_socket_fds,
-		                        (phys_ctx->num_listening_sockets + 1)
+		                        (phys_ctx->num_listening_sockets + 2)
 		                            * sizeof(phys_ctx->listening_socket_fds[0]),
 		                        phys_ctx))
 		    == NULL) {
@@ -16552,15 +16576,18 @@ sslize(struct mg_connection *conn,
 					/* Need to retry the function call "later".
 					 * See https://linux.die.net/man/3/ssl_get_error
 					 * This is typical for non-blocking sockets. */
-					struct mg_pollfd pfd;
+					struct mg_pollfd pfd[2];
 					int pollres;
-					pfd.fd = conn->client.sock;
-					pfd.events = ((err == SSL_ERROR_WANT_CONNECT)
-					              || (err == SSL_ERROR_WANT_WRITE))
-					                 ? POLLOUT
-					                 : POLLIN;
+					pfd[0].fd = conn->client.sock;
+					pfd[0].events = ((err == SSL_ERROR_WANT_CONNECT)
+					                || (err == SSL_ERROR_WANT_WRITE))
+					                   ? POLLOUT
+					                   : POLLIN;
+
+					pfd[1].fd = conn->phys_ctx->thread_shutdown_notification_socket;
+					pfd[1].events = POLLIN;
 					pollres =
-					    mg_poll(&pfd, 1, 50, &(conn->phys_ctx->stop_flag));
+					    mg_poll(pfd, 2, 50, &(conn->phys_ctx->stop_flag));
 					if (pollres < 0) {
 						/* Break if error occurred (-1)
 						 * or server shutdown (-2) */
@@ -20144,8 +20171,14 @@ master_thread_run(struct mg_context *ctx)
 			pfd[i].events = POLLIN;
 		}
 
+		/* We listen on this socket just so that mg_stop() can cause mg_poll() to return ASAP.
+		 * Don't worry, we did allocate an extra slot at the end of listening_socket_fds[] just to hold this
+		 */
+		pfd[ctx->num_listening_sockets].fd = ctx->thread_shutdown_notification_socket;
+		pfd[ctx->num_listening_sockets].events = POLLIN;
+
 		if (mg_poll(pfd,
-		            ctx->num_listening_sockets,
+		            ctx->num_listening_sockets+1,   // +1 for the thread_shutdown_notification_socket
 		            SOCKET_TIMEOUT_QUANTUM,
 		            &(ctx->stop_flag))
 		    > 0) {
@@ -20303,6 +20336,14 @@ free_context(struct mg_context *ctx)
 	(void)pthread_mutex_destroy(&ctx->lua_bg_mutex);
 #endif
 
+	/* Deallocate shutdown-triggering socket-pair */
+	if (ctx->user_shutdown_notification_socket >= 0) {
+		closesocket(ctx->user_shutdown_notification_socket);
+	}
+	if (ctx->thread_shutdown_notification_socket >= 0) {
+		closesocket(ctx->thread_shutdown_notification_socket);
+	}
+
 	/* Deallocate config parameters */
 	for (i = 0; i < NUM_OPTIONS; i++) {
 		if (ctx->dd.config[i] != NULL) {
@@ -20379,6 +20420,10 @@ mg_stop(struct mg_context *ctx)
 	/* Set stop flag, so all threads know they have to exit. */
 	STOP_FLAG_ASSIGN(&ctx->stop_flag, 1);
 
+	/* Closing this socket will cause mg_poll() in all the I/O threads to return immediately */
+	closesocket(ctx->user_shutdown_notification_socket);
+	ctx->user_shutdown_notification_socket = -1;  /* to avoid calling closesocket() again in free_context() */
+
 	/* Join timer thread */
 #if defined(USE_TIMERS)
 	timers_exit(ctx);
@@ -20476,6 +20521,64 @@ legacy_init(const char **options)
 	}
 }
 
+/* we'll assume it's only Windows that doesn't have socketpair() available */
+#if !defined(HAVE_SOCKETPAIR) && !defined(_WIN32)
+# define HAVE_SOCKETPAIR 1
+#endif
+
+static int
+mg_socketpair(int * sockA,
+	      int * sockB)
+{
+	int temp[2] = {-1, -1};
+
+	/** Default to unallocated */
+	*sockA = -1;
+	*sockB = -1;
+
+#if defined(HAVE_SOCKETPAIR)
+	int ret = socketpair(AF_UNIX, SOCK_STREAM, 0, temp);
+	if (ret == 0) {
+		*sockA = temp[0];
+		*sockB = temp[1];
+	}
+	return ret;
+#else
+	/** No socketpair() call is available, so we'll have to roll our own implementation */
+	int asock = socket(PF_INET, SOCK_STREAM, 0);
+	if (asock >= 0) {
+		struct sockaddr_in addr;
+		struct sockaddr * pa = (struct sockaddr *) &addr;
+		socklen_t addrLen = sizeof(addr);
+
+		memset(&addr, 0, sizeof(addr));
+		addr.sin_family = AF_INET;
+		addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
+		addr.sin_port = 0;
+
+		if ((bind(asock, pa, sizeof(addr)) == 0)
+		  &&(getsockname(asock, pa, &addrLen) == 0)
+		  &&(listen(asock, 1) == 0)) {
+			temp[0] = socket(PF_INET, SOCK_STREAM, 0);
+			if ((temp[0] >= 0)&&(connect(temp[0], pa, sizeof(addr)) == 0)) {
+				temp[1] = accept(asock, pa, &addrLen);
+				if (temp[1] >= 0) {
+					closesocket(asock);
+					*sockA = temp[0];
+					*sockB = temp[1];
+					return 0;  /* success! */
+				}
+			}
+		}
+	}
+
+	/* Cleanup */
+	if (asock   >= 0) closesocket(asock);
+	if (temp[0] >= 0) closesocket(temp[0]);
+	if (temp[1] >= 0) closesocket(temp[1]);
+	return -1;  /* fail! */
+#endif
+}
 
 CIVETWEB_API struct mg_context *
 mg_start2(struct mg_init_data *init, struct mg_error_data *error)
@@ -20559,6 +20662,12 @@ mg_start2(struct mg_init_data *init, struct mg_error_data *error)
 #if defined(USE_LUA)
 	ok &= (0 == pthread_mutex_init(&ctx->lua_bg_mutex, &pthread_mutex_attr));
 #endif
+
+	/** mg_stop() will close the user_shutdown_notification_socket, and that will cause poll()
+	  * to return immediately in the master-thread, so that mg_stop() can also return immediately.
+	  */
+	ok &= (0 == mg_socketpair(&ctx->user_shutdown_notification_socket, &ctx->thread_shutdown_notification_socket));
+
 	if (!ok) {
 		unsigned error_id = (unsigned)ERRNO;
 		const char *err_msg =