pgduck_server handles sigint/sigterm gracefully (#229)

Signed-off-by: Aykut Bozkurt <aykut.bozkurt@snowflake.com>

Author: sfc-gh-abozkurt

pgduck_server/include/pgserver/client_threadpool.h

@@ -41,5 +41,6 @@ extern int	pgclient_threadpool_cancel_thread(int cancellationProcId, uint8 *canc
 extern int	pgclient_threadpool_cancel_thread(int cancellationProcId, int32 cancellationToken);
 #endif
 extern void pgclient_threadpool_set_duckdb_conn(int threadIndex, duckdb_connection conn);
+extern int	pgclient_threadpool_cancel_all(void);
 
 #endif							/* PGDUCK_CLIENT_THREAD_H */

pgduck_server/include/pgserver/pgserver.h

@@ -50,7 +50,7 @@ extern int	pgserver_init(PGServer * pgServer,
 						  int unixSocketPermissions,
 						  int port);
 extern int	pgserver_run(PGServer * pgServer);
-extern int	pgserver_destroy(PGServer * pgServer);
+extern void pgserver_destroy(PGServer * pgServer);
 
 
 #endif							/* PGDUCK_PG_SERVER_H */

pgduck_server/src/main.c

@@ -88,8 +88,7 @@ main(int argc, char *argv[])
 	if (pgserver_run(&pgServer) != STATUS_OK)
 		return STATUS_ERROR;
 
-	if (pgserver_destroy(&pgServer) != STATUS_OK)
-		return STATUS_ERROR;
+	pgserver_destroy(&pgServer);
 
 	return STATUS_OK;
 }

pgduck_server/src/pgserver/client_threadpool.c

@@ -323,3 +323,36 @@ pgclient_threadpool_set_duckdb_conn(int threadIndex, duckdb_connection conn)
 
 	pthread_rwlock_unlock(&rwlock);
 }
+
+
+/*
+ * pgclient_threadpool_cancel_all interrupts every active DuckDB query.
+ *
+ * Called during server shutdown so that client threads get a clean
+ * interruption error instead of an abrupt connection reset when the
+ * process exits.  Only sets the DuckDB interrupt flag (an atomic bool),
+ * so this is cheap and safe to call from the main thread.
+ *
+ * Returns the number of active threads that were interrupted.
+ */
+int
+pgclient_threadpool_cancel_all(void)
+{
+	int			interrupted = 0;
+
+	pthread_rwlock_rdlock(&rwlock);
+
+	for (int i = 0; i < MaxThreads; i++)
+	{
+		if (ClientThreadPool[i].isStarted &&
+			ClientThreadPool[i].duckdbConnection != NULL)
+		{
+			duckdb_interrupt(ClientThreadPool[i].duckdbConnection);
+			interrupted++;
+		}
+	}
+
+	pthread_rwlock_unlock(&rwlock);
+
+	return interrupted;
+}

pgduck_server/src/pgserver/pgserver.c

@@ -30,6 +30,7 @@
 #include <netdb.h>
 #include <common/ip.h>
 #include <pthread.h>
+#include <signal.h>
 #include <sys/fcntl.h>
 #include <sys/file.h>
 #include <sys/stat.h>
@@ -325,45 +326,98 @@ set_unix_socket_permissions(char *unixSocketPath, char *groupName, int permissio
 
 static volatile sig_atomic_t running = 1;
 
-/* basic sigint handler */
 static void
-handle_signal(int sig)
+handle_shutdown_signal(int sig)
 {
 	running = 0;
 }
 
 
 /*
- * pgserver_run is the main loop for the PostgreSQL wire compatible server.
+ * install_shutdown_signal_handlers -- install SIGINT/SIGTERM handlers.
+ *
+ * The handler simply sets `running = 0`.  SA_RESTART is deliberately
+ * *not* set so that accept() in the main loop returns -1/EINTR, giving
+ * the loop a chance to notice the flag and exit promptly.
  */
-int
-pgserver_run(PGServer * pgServer)
+static int
+install_shutdown_signal_handlers(void)
 {
-	/* install signal handlers */
 	struct sigaction sa;
 
-	/* Use our custom handler */
-	sa.sa_handler = handle_signal;
-
-	/* Do not block any other signals during handling */
+	sa.sa_handler = handle_shutdown_signal;
 	sigemptyset(&sa.sa_mask);
+	sa.sa_flags = 0;			/* no SA_RESTART — accept() must return
+								 * EINTR */
 
-	/* CRITICAL: Do NOT set the SA_RESTART flag. */
-	/* This ensures that system calls like accept() are interrupted. */
-	sa.sa_flags = 0;
+	if (sigaction(SIGINT, &sa, NULL) == -1 ||
+		sigaction(SIGTERM, &sa, NULL) == -1)
+	{
+		PGDUCK_SERVER_ERROR("sigaction failed: %s", strerror(errno));
+		return STATUS_ERROR;
+	}
+
+	return STATUS_OK;
+}
 
-	/* Install the handler for SIGINT and SIGTERM */
-	if (sigaction(SIGINT, &sa, NULL) == -1)
+
+/*
+ * disable_shutdown_signals -- block SIGINT/SIGTERM in the calling thread.
+ *
+ * Used before pthread_create() so the child thread inherits a blocked
+ * mask and never receives these signals.
+ */
+static int
+disable_shutdown_signals(void)
+{
+	sigset_t	sigs;
+
+	sigemptyset(&sigs);
+	sigaddset(&sigs, SIGINT);
+	sigaddset(&sigs, SIGTERM);
+
+	if (pthread_sigmask(SIG_BLOCK, &sigs, NULL) != 0)
 	{
-		perror("sigaction for SIGINT failed");
-		exit(STATUS_ERROR);
+		PGDUCK_SERVER_ERROR("pthread_sigmask failed: %s", strerror(errno));
+		return STATUS_ERROR;
 	}
-	if (sigaction(SIGTERM, &sa, NULL) == -1)
+
+	return STATUS_OK;
+}
+
+
+/*
+ * enable_shutdown_signals -- unblock SIGINT/SIGTERM and re-install the
+ * shutdown signal handlers so accept() can be interrupted again.
+ */
+static int
+enable_shutdown_signals(void)
+{
+	sigset_t	sigs;
+
+	sigemptyset(&sigs);
+	sigaddset(&sigs, SIGINT);
+	sigaddset(&sigs, SIGTERM);
+
+	if (pthread_sigmask(SIG_UNBLOCK, &sigs, NULL) != 0)
 	{
-		perror("sigaction for SIGTERM failed");
-		exit(STATUS_ERROR);
+		PGDUCK_SERVER_ERROR("pthread_sigmask failed: %s", strerror(errno));
+		return STATUS_ERROR;
 	}
 
+	return STATUS_OK;
+}
+
+
+/*
+ * pgserver_run is the main loop for the PostgreSQL wire compatible server.
+ */
+int
+pgserver_run(PGServer * pgServer)
+{
+	if (install_shutdown_signal_handlers() != STATUS_OK)
+		return STATUS_ERROR;
+
 	while (running)
 	{
 		PGClient   *client = (PGClient *) pg_malloc0(sizeof(PGClient));
@@ -375,8 +429,21 @@ pgserver_run(PGServer * pgServer)
 
 		if (client->clientSocket < 0)
 		{
+			int			save_errno = errno;
+
+			pg_free(client);
+
+			/*
+			 * EINTR can come from our shutdown handler (running == 0) or from
+			 * unrelated sources like a debugger attaching (ptrace). In either
+			 * case, just retry the loop — the while-condition takes care of
+			 * the shutdown case.
+			 */
+			if (save_errno == EINTR)
+				continue;
+
 			PGDUCK_SERVER_ERROR("Could not accept the client: %s",
-								strerror(errno));
+								strerror(save_errno));
 
 			/*
 			 * TODO: We can probably recover from this error, but lets handle
@@ -422,6 +489,9 @@ pgserver_run(PGServer * pgServer)
 		initState->startFunction = pgServer->startFunction;
 		initState->pgClient = client;
 
+		if (disable_shutdown_signals() != STATUS_OK)
+			exit(STATUS_ERROR);
+
 		if (pgserver_create_client_thread(initState) != OK)
 		{
 			PGDUCK_SERVER_ERROR("Thread creation failed for client %d", client->clientSocket);
@@ -430,19 +500,61 @@ pgserver_run(PGServer * pgServer)
 			pg_free(client);
 			pg_free(initState);
 			pgclient_threadpool_free_slot(threadIndex);
-			continue;
 		}
-	}
 
-	PGDUCK_SERVER_LOG("Done running");
+		if (enable_shutdown_signals() != STATUS_OK)
+			exit(STATUS_ERROR);
+	}
 
 	return STATUS_OK;
 }
 
 
+/*
+ * pgserver_destroy performs a graceful shutdown of the server.
+ *
+ * 1. Close the listening socket so no new connections are accepted.
+ * 2. Interrupt every active DuckDB query so client threads can send a
+ *    proper error to their clients instead of an abrupt TCP reset.
+ * 3. Brief grace period to let interrupted threads finish their error
+ *    path and close their sockets cleanly.
+ *
+ * We don't join the client threads (they are detached), so the grace
+ * period is best-effort.  When main() returns, exit() will tear down
+ * any remaining threads.
+ */
+void
+pgserver_destroy(PGServer * pgServer)
+{
+	PGDUCK_SERVER_LOG("Shutting down: closing listening socket");
+	closesocket(pgServer->listeningSocket);
+	pgServer->listeningSocket = -1;
+
+	int			interrupted = pgclient_threadpool_cancel_all();
+
+	if (interrupted > 0)
+	{
+		PGDUCK_SERVER_LOG("Shutting down: interrupted %d active connection(s), "
+						  "waiting briefly for them to finish", interrupted);
+
+		/*
+		 * 2 seconds is generous for the threads to send an error to their
+		 * client and run through pgclient_thread_cleanup.
+		 */
+		pg_usleep(2 * 1000000L);
+	}
+
+	PGDUCK_SERVER_LOG("Done running");
+}
+
+
 /*
  * pgserver_create_client_thread creates a new thread for the client.
  * We use PTHREAD_CREATE_DETACHED so that we don't have to join the threads.
+ *
+ * The caller must block shutdown signals before calling this function
+ * so the new thread inherits a blocked mask and never receives
+ * SIGINT/SIGTERM.
  */
 static int
 pgserver_create_client_thread(const PgClientThreadInitState * initState)
@@ -487,6 +599,12 @@ pgclient_thread_main(void *arg)
 {
 	PgClientThreadInitState *initState = (PgClientThreadInitState *) arg;
 
+	/*
+	 * SIGINT/SIGTERM are already blocked — pgserver_create_client_thread()
+	 * blocks them before pthread_create(), so this thread inherits a blocked
+	 * mask.  No per-thread sigmask call needed.
+	 */
+
 	/* cleanup handler */
 	pthread_cleanup_push(pgclient_thread_cleanup, arg);
 
@@ -523,18 +641,6 @@ pgclient_thread_cleanup(void *arg)
 }
 
 
-/*
- * cleanup on successful exists.
- *
- * TODO: not called ever yet
- */
-int
-pgserver_destroy(PGServer * pgServer)
-{
-	closesocket(pgServer->listeningSocket);
-
-	return STATUS_OK;
-}
 
 
 /*

pgduck_server/tests/pytests/test_s3.py

@@ -198,4 +198,7 @@ def test_s3_get_region_invalid(pgduck_conn):
         pgduck_conn,
         raise_error=False,
     )
-    assert "Could not establish connection error" in error
+    assert (
+        "Could not establish connection error" in error
+        or "server closed the connection" in error
+    )