fix spurious cancel on pgduck_server

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

Author: sfc-gh-abozkurt

pgduck_server/include/pgserver/client_threadpool.h

@@ -41,6 +41,7 @@ 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 void pgclient_threadpool_set_executing(int threadIndex, bool executing);
 extern int	pgclient_threadpool_cancel_all(void);
 
 #endif							/* PGDUCK_CLIENT_THREAD_H */

pgduck_server/src/pgserver/client_threadpool.c

@@ -58,6 +58,13 @@ typedef struct PgClientThreadState
 	/* DuckDB connection to interrupt */
 	duckdb_connection duckdbConnection;
 
+	/*
+	 * Whether the thread is currently executing a DuckDB query.  Set by the
+	 * session thread (under wrlock) before/after DuckDB execution so that the
+	 * cancel handler can skip duckdb_interrupt() on idle sessions.
+	 */
+	bool		isExecutingQuery;
+
 }			PgClientThreadState;
 
 int			MaxAllowedClients;
@@ -221,6 +228,7 @@ pgclient_threadpool_free_slot(int threadIndex)
 	threadState->cancellationTokenSize = 0;
 #endif
 	threadState->duckdbConnection = NULL;
+	threadState->isExecutingQuery = false;
 
 	--ActiveClientThreadCount;
 
@@ -252,18 +260,18 @@ pgclient_threadpool_cancel_thread(int cancellationProcId, uint8 *cancellationTok
 			ClientThreadPool[threadIndex].cancellationProcId == cancellationProcId)
 		{
 			usedThreadIndex = threadIndex;
-			duckdb_connection conn = ClientThreadPool[threadIndex].duckdbConnection;
 
 			/*
-			 * As per DuckDB docs, duckdb connections are thread safe so we
-			 * can safely interrupt it from another thread.
-			 *
-			 * We do the interrupt while holding the threadpool lock.
-			 * Otherwise, a thread could end just before we call
-			 * duckdb_interrupt. Luckily, duckdb_interrupt will only set an
-			 * atomic<bool> flag.
+			 * Only interrupt when the session is actually executing a DuckDB
+			 * query.  Interrupting an idle session sets an atomic flag that
+			 * persists and poisons the next query on that connection.
 			 */
-			duckdb_interrupt(conn);
+			if (ClientThreadPool[threadIndex].isExecutingQuery)
+			{
+				duckdb_connection conn = ClientThreadPool[threadIndex].duckdbConnection;
+
+				duckdb_interrupt(conn);
+			}
 			break;
 		}
 	}
@@ -286,18 +294,18 @@ pgclient_threadpool_cancel_thread(int cancellationProcId, int32 cancellationToke
 			ClientThreadPool[threadIndex].cancellationProcId == cancellationProcId)
 		{
 			usedThreadIndex = threadIndex;
-			duckdb_connection conn = ClientThreadPool[threadIndex].duckdbConnection;
 
 			/*
-			 * As per DuckDB docs, duckdb connections are thread safe so we
-			 * can safely interrupt it from another thread.
-			 *
-			 * We do the interrupt while holding the threadpool lock.
-			 * Otherwise, a thread could end just before we call
-			 * duckdb_interrupt. Luckily, duckdb_interrupt will only set an
-			 * atomic<bool> flag.
+			 * Only interrupt when the session is actually executing a DuckDB
+			 * query.  Interrupting an idle session sets an atomic flag that
+			 * persists and poisons the next query on that connection.
 			 */
-			duckdb_interrupt(conn);
+			if (ClientThreadPool[threadIndex].isExecutingQuery)
+			{
+				duckdb_connection conn = ClientThreadPool[threadIndex].duckdbConnection;
+
+				duckdb_interrupt(conn);
+			}
 			break;
 		}
 	}
@@ -325,6 +333,20 @@ pgclient_threadpool_set_duckdb_conn(int threadIndex, duckdb_connection conn)
 }
 
 
+/*
+ * pgclient_threadpool_set_executing marks whether the given thread is
+ * currently executing a DuckDB query.  The cancel handler checks this
+ * flag (under rdlock) to avoid interrupting idle sessions.
+ */
+void
+pgclient_threadpool_set_executing(int threadIndex, bool executing)
+{
+	pthread_rwlock_wrlock(&rwlock);
+	ClientThreadPool[threadIndex].isExecutingQuery = executing;
+	pthread_rwlock_unlock(&rwlock);
+}
+
+
 /*
  * pgclient_threadpool_cancel_all interrupts every active DuckDB query.
  *

pgduck_server/src/pgsession/pgsession.c

@@ -346,6 +346,15 @@ pgsession_handle_connection(void *input)
 
 
 finally:
+	/*
+	 * Clear the DuckDB connection pointer from the thread pool BEFORE
+	 * destroying the session.  Otherwise a cancel request arriving between
+	 * duckdb_disconnect (inside pgsession_destroy) and
+	 * pgclient_threadpool_free_slot would call duckdb_interrupt on a freed
+	 * pointer.
+	 */
+	pgclient_threadpool_set_duckdb_conn(pgClient->threadIndex, NULL);
+
 	pgsession_prepared_statement_deallocate(&pgSession);
 	pgsession_destroy(&pgSession);
 	pg_free(inputMessage.data);
@@ -392,8 +401,12 @@ process_query_message(PGSession * pgSession, StringInfo inputMessage)
 	}
 
 	char	   *errorMessage = NULL;
+	int			threadIndex = pgSession->pgClient->threadIndex;
+
+	pgclient_threadpool_set_executing(threadIndex, true);
 	DuckDBStatus status = duckdb_session_run_command(&pgSession->duckSession, queryString,
 													 &responseFormat, &errorMessage);
+	pgclient_threadpool_set_executing(threadIndex, false);
 
 	if (status == DUCKDB_SUCCESS)
 	{
@@ -506,8 +519,12 @@ process_parse_message(PGSession * pgSession, StringInfo inputMessage)
 	char	   *queryStringCopy = pstrdup(queryString);
 
 	char	   *errorMessage = NULL;
+	int			threadIndex = pgSession->pgClient->threadIndex;
+
+	pgclient_threadpool_set_executing(threadIndex, true);
 	DuckDBStatus status = duckdb_session_prepare(&pgSession->duckSession, queryStringCopy,
 												 &errorMessage);
+	pgclient_threadpool_set_executing(threadIndex, false);
 
 	/*
 	 * make sure we destroy the prepared statement, even in case of failure
@@ -813,9 +830,13 @@ process_execute_message(PGSession * pgSession, StringInfo inputMessage)
 	ResponseFormat *responseFormat = &pgSession->pgSessionPreparedStmt.responseFormat;
 
 	char	   *errorMessage = NULL;
+	int			threadIndex = pgSession->pgClient->threadIndex;
+
+	pgclient_threadpool_set_executing(threadIndex, true);
 	DuckDBStatus status = duckdb_session_execute_prepared(&pgSession->duckSession,
 														  responseFormat,
 														  &errorMessage);
+	pgclient_threadpool_set_executing(threadIndex, false);
 
 	if (status == DUCKDB_SUCCESS)
 	{