[PATCH] immediately kill psql process if server is not running.

From 61221ab134ce06363a9f0c20aa3d5829db374a89 Mon Sep 17 00:00:00 2001
From: Srinath Reddy Sadipiralla <srinath.reddy@zohocorp.com>
Date: Fri, 8 Nov 2024 11:10:40 +0530
Subject: [PATCH] Kill psql process immediately if server crashed

---
 src/bin/psql/input.c        |  21 +++++++
 src/bin/psql/input.h        |   3 +
 src/bin/psql/startup.c      | 113 +++++++++++++++++++++++++++++++++---
 src/bin/psql/tab-complete.c |   7 +++
 4 files changed, 136 insertions(+), 8 deletions(-)

diff --git a/src/bin/psql/input.c b/src/bin/psql/input.c
index 01b7ef74c3..3912f00e38 100644
--- a/src/bin/psql/input.c
+++ b/src/bin/psql/input.c
@@ -47,6 +47,9 @@ static int	history_lines_added;
 #define NL_IN_HISTORY	0x01
 #endif
 
+bool waiting_for_input = true;
+pthread_mutex_t waiting_for_input_lock = PTHREAD_MUTEX_INITIALIZER;
+
 static void finishInput(void);
 
 
@@ -87,8 +90,17 @@ gets_interactive(const char *prompt, PQExpBuffer query_buf)
 
 		/* Enable SIGINT to longjmp to sigint_interrupt_jmp */
 		sigint_interrupt_enabled = true;
+		
+		/* Let the thread monitoring psql_fd know that main thread is waiting for user input */
+		pthread_mutex_lock(&waiting_for_input_lock);
+		waiting_for_input = true;
+		pthread_mutex_unlock(&waiting_for_input_lock);
 
 		result = readline(prompt);
+		
+		pthread_mutex_lock(&waiting_for_input_lock);
+		waiting_for_input = false;
+		pthread_mutex_unlock(&waiting_for_input_lock);
 
 		/* Disable SIGINT again */
 		sigint_interrupt_enabled = false;
@@ -201,9 +213,18 @@ gets_fromFile(FILE *source)
 		/* Enable SIGINT to longjmp to sigint_interrupt_jmp */
 		sigint_interrupt_enabled = true;
 
+		/* Let the thread monitoring psql_fd know that main thread is waiting for user input */
+		pthread_mutex_lock(&waiting_for_input_lock);
+		waiting_for_input = true;
+		pthread_mutex_unlock(&waiting_for_input_lock);
+
 		/* Get some data */
 		result = fgets(line, sizeof(line), source);
 
+		pthread_mutex_lock(&waiting_for_input_lock);
+		waiting_for_input = false;
+		pthread_mutex_unlock(&waiting_for_input_lock);
+
 		/* Disable SIGINT again */
 		sigint_interrupt_enabled = false;
 
diff --git a/src/bin/psql/input.h b/src/bin/psql/input.h
index 4c486d67d2..5afb6f58b3 100644
--- a/src/bin/psql/input.h
+++ b/src/bin/psql/input.h
@@ -37,6 +37,9 @@
 
 #include "pqexpbuffer.h"
 
+#include <pthread/pthread.h>
+extern bool waiting_for_input;
+extern pthread_mutex_t waiting_for_input_lock;
 
 extern char *gets_interactive(const char *prompt, PQExpBuffer query_buf);
 extern char *gets_fromFile(FILE *source);
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
index 036caaec2f..ab352eeb89 100644
--- a/src/bin/psql/startup.c
+++ b/src/bin/psql/startup.c
@@ -26,6 +26,13 @@
 #include "mainloop.h"
 #include "settings.h"
 
+#include <pthread.h>
+#include <sys/socket.h>
+
+pthread_t psql_fd_monitor_thread;
+bool main_thread_exited = false;
+pthread_mutex_t main_thread_exited_lock = PTHREAD_MUTEX_INITIALIZER;
+
 /*
  * Global psql options
  */
@@ -117,6 +124,87 @@ empty_signal_handler(SIGNAL_ARGS)
 }
 #endif
 
+static void clean_up()
+{
+	/* clean up */
+	if (pset.logfile)
+		fclose(pset.logfile);
+	if (pset.db)
+		PQfinish(pset.db);
+	if (pset.dead_conn)
+		PQfinish(pset.dead_conn);
+	setQFout(NULL);
+}
+
+static void terminate_process()
+{
+	pthread_mutex_lock(&main_thread_exited_lock);
+	if (!main_thread_exited)
+	{
+		pthread_mutex_unlock(&main_thread_exited_lock);
+		clean_up();
+		kill(getpid(), SIGKILL);
+	}
+	pthread_mutex_unlock(&main_thread_exited_lock);
+	pthread_exit(NULL);
+}
+
+/* Thread function for monitoring the psql's file descriptor */
+static void *psql_fd_monitor_thread_func()
+{
+	pgsocket psql_fd = PQsocket(pset.db);
+	fd_set read_fds;
+	int retval;
+
+	while (1)
+	{
+		/* Check for server process death only if psql is waiting for input */
+		pthread_mutex_lock(&waiting_for_input_lock);
+		if (waiting_for_input)
+		{
+			struct timeval timeout = {};
+			pthread_mutex_unlock(&waiting_for_input_lock);
+
+			FD_ZERO(&read_fds);
+			FD_SET(psql_fd, &read_fds);
+
+			retval = select(psql_fd + 1, &read_fds, NULL, NULL, &timeout);
+			if (retval == -1)
+			{
+				terminate_process();
+			}
+
+			if (FD_ISSET(psql_fd, &read_fds))
+			{
+				char buf;
+				int len = recv(psql_fd, &buf, 1, 0);
+				if (len <= 0)
+				{
+					// If recv fails, we consider it as a trigger to terminate the process
+					terminate_process();
+				}
+			}
+		}
+		else
+		{
+			pthread_mutex_unlock(&waiting_for_input_lock);
+		}
+	}
+
+	return NULL;
+}
+
+static void create_psql_fd_monitor_thread()
+{
+
+	if (pthread_create(&psql_fd_monitor_thread, NULL, psql_fd_monitor_thread_func, NULL) != 0)
+	{
+		pg_log_error("Failed to create fd monitor thread");
+		exit(EXIT_FAILURE);
+	}
+	
+}
+
 /*
  *
  * main
@@ -459,17 +547,26 @@ error:
 		if (!pset.quiet)
 			printf(_("Type \"help\" for help.\n\n"));
 		initializeInput(options.no_readline ? 0 : 1);
+		/* Create a thread to monitor the psql's file descriptor */
+		create_psql_fd_monitor_thread();
 		successResult = MainLoop(stdin);
 	}
 
-	/* clean up */
-	if (pset.logfile)
-		fclose(pset.logfile);
-	if (pset.db)
-		PQfinish(pset.db);
-	if (pset.dead_conn)
-		PQfinish(pset.dead_conn);
-	setQFout(NULL);
+	clean_up();
+
+	pthread_mutex_lock(&main_thread_exited_lock);
+	main_thread_exited = true;
+	pthread_mutex_unlock(&main_thread_exited_lock);
+
+	pthread_mutex_lock(&waiting_for_input_lock);
+	waiting_for_input = true;
+	pthread_mutex_unlock(&waiting_for_input_lock);
+
+	// Wait for the psql's fd monitoring thread to exit
+    if (pthread_join(psql_fd_monitor_thread, NULL) != 0)
+    {
+        pg_log_error("Failed to join psql's fd monitor thread");
+    }
 
 	return successResult;
 }
diff --git a/src/bin/psql/tab-complete.c b/src/bin/psql/tab-complete.c
index 2e55607b8a..023b5096cf 100644
--- a/src/bin/psql/tab-complete.c
+++ b/src/bin/psql/tab-complete.c
@@ -1752,6 +1752,10 @@ psql_completion(const char *text, int start, int end)
 	char	   *text_copy = pnstrdup(rl_line_buffer + start, end - start);
 	text = text_copy;
 
+	pthread_mutex_lock(&waiting_for_input_lock);
+	waiting_for_input = false;
+	pthread_mutex_unlock(&waiting_for_input_lock);
+
 	/* Remember last char of the given input word. */
 	completion_last_char = (end > start) ? text[end - start - 1] : '\0';
 
@@ -5060,6 +5064,9 @@ psql_completion(const char *text, int start, int end)
 	completion_ref_object = NULL;
 	free(completion_ref_schema);
 	completion_ref_schema = NULL;
+	pthread_mutex_lock(&waiting_for_input_lock);
+	waiting_for_input = true;
+	pthread_mutex_unlock(&waiting_for_input_lock);
 
 	/* Return our Grand List O' Matches */
 	return matches;
-- 
2.39.3 (Apple Git-146)

From f493be68e085fec4a7fb9518bb53d3e6da2acd18 Mon Sep 17 00:00:00 2001
From: Srinath Reddy Sadipiralla <srinath.reddy@zohocorp.com>
Date: Fri, 8 Nov 2024 23:53:46 +0530
Subject: [PATCH] Kill psql process with single process implementation

---
 src/bin/psql/input.c | 105 ++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 103 insertions(+), 2 deletions(-)

diff --git a/src/bin/psql/input.c b/src/bin/psql/input.c
index 01b7ef74c3..0f50e9c866 100644
--- a/src/bin/psql/input.c
+++ b/src/bin/psql/input.c
@@ -18,6 +18,8 @@
 #include "input.h"
 #include "settings.h"
 #include "tab-complete.h"
+#include <pthread.h>
+#include <sys/socket.h>
 
 #ifndef WIN32
 #define PSQLHISTORY ".psql_history"
@@ -47,8 +49,107 @@ static int	history_lines_added;
 #define NL_IN_HISTORY	0x01
 #endif
 
+static char *line_buffer = NULL;
+static bool received_line = false;
+char *full_query_buffer = NULL;
+
 static void finishInput(void);
 
+static void handle_line(char *line)
+{
+	if (line == NULL)
+	{
+		rl_callback_handler_remove();
+		exit(0);
+	}
+
+	if (line_buffer != NULL)
+	{
+		free(line_buffer);
+		line_buffer = NULL;
+	}
+
+	if (full_query_buffer == NULL)
+	{
+		full_query_buffer = strdup(line);
+	}
+	else
+	{
+		char *temp = malloc(strlen(full_query_buffer) + strlen(line) + 2);
+		sprintf(temp, "%s\n%s", full_query_buffer, line);
+		free(full_query_buffer);
+		full_query_buffer = temp;
+	}
+
+	if (full_query_buffer[strlen(full_query_buffer) - 1] == ';')
+	{
+		line_buffer = strdup(full_query_buffer);
+		free(full_query_buffer);
+		full_query_buffer = NULL;
+		received_line = true;
+	}
+	else
+	{
+		line_buffer = strdup(line);
+		received_line = true;
+	}
+	free(line);
+	line = NULL;
+}
+
+static char *monitor_stdin_psql_fds(const char *prompt)
+{
+	fd_set read_fds;
+	int max_fd;
+	struct timeval timeout = {};
+	int retval;
+
+	pgsocket psql_fd = PQsocket(pset.db);
+	int terminal_fd = fileno(stdin);
+
+	rl_callback_handler_install(prompt, handle_line);
+
+	while (1)
+	{
+		FD_ZERO(&read_fds);
+		FD_SET(terminal_fd, &read_fds);
+		FD_SET(psql_fd, &read_fds);
+		max_fd = (terminal_fd > psql_fd) ? terminal_fd : psql_fd;
+
+		retval = select(max_fd + 1, &read_fds, NULL, NULL, &timeout);
+
+		if (retval == -1 && errno != EINTR && errno != EAGAIN)
+		{
+			rl_callback_handler_remove();
+			return NULL;
+		}
+
+		if (FD_ISSET(terminal_fd, &read_fds))
+		{
+
+			rl_callback_read_char();
+			if (received_line)
+			{
+				char *result = line_buffer;
+				line_buffer = NULL;
+				received_line = false;
+				return result;
+			}
+		}
+
+		if (FD_ISSET(psql_fd, &read_fds))
+		{
+			char buf;
+			int len = recv(psql_fd, &buf, 1, 0);
+			if (len <= 0 && errno != EINTR && errno != EAGAIN)
+			{
+				rl_callback_handler_remove();
+				/* If recv fails, we consider it as a trigger to terminate the process */
+				return NULL;
+			}
+		}
+	}
+}
 
 /*
  * gets_interactive()
@@ -88,8 +189,8 @@ gets_interactive(const char *prompt, PQExpBuffer query_buf)
 		/* Enable SIGINT to longjmp to sigint_interrupt_jmp */
 		sigint_interrupt_enabled = true;
 
-		result = readline(prompt);
-
+		result = monitor_stdin_psql_fds(prompt);
+		
 		/* Disable SIGINT again */
 		sigint_interrupt_enabled = false;
 
-- 
2.39.3 (Apple Git-146)