>From 61221ab134ce06363a9f0c20aa3d5829db374a89 Mon Sep 17 00:00:00 2001 From: Srinath Reddy Sadipiralla 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 +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 +#include + +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)