Refactor Connection Reset Handling

* Introduced startup parameters in PgSQL_Connection (removed default session parameters from PgSQL_Session)
* Startup parameters are populated during both frontend and backend connection creation
* Parameters provided via connection options are set as startup parameters
* Backend connection parameter handling updated: only critical variables are now set via connection options to prevent interference with DISCARD ALL during connection reset; remaining parameters will be applied using individual SET commands

Author: rahim-kanji

include/PgSQL_Connection.h

@@ -244,7 +244,7 @@ class PgSQL_Connection_userinfo {
 
 class PgSQL_Connection {
 public:
-	PgSQL_Connection();
+	PgSQL_Connection(bool is_client_conn);
 	~PgSQL_Connection();
 
 	PG_ASYNC_ST handler(short event);
@@ -312,7 +312,7 @@ class PgSQL_Connection {
 		if (error_info.severity == PGSQL_ERROR_SEVERITY::ERRSEVERITY_FATAL ||
 			error_info.severity == PGSQL_ERROR_SEVERITY::ERRSEVERITY_ERROR ||
 			error_info.severity == PGSQL_ERROR_SEVERITY::ERRSEVERITY_PANIC) {
-				return true;
+			return true;
 		}
 		return false;
 	}
@@ -442,6 +442,42 @@ class PgSQL_Connection {
 
 	bool IsKeepMultiplexEnabledVariables(char* query_digest_text);
 
+	/**
+	 * @brief Retrieves startup parameter and it's hash
+	 *
+	 * This function tries to retrieve value and hash of startup paramters if present (provided in connection parameters).
+	 * If value is not found, it falls back to the thread-specific default variables.
+	 *
+	 * @param idx The index of startup parameter to retrieve.
+	 * @return The value and hash of startup parameter.
+	 *
+	 */
+	std::pair<const char*, uint32_t> get_startup_parameter_and_hash(enum pgsql_variable_name idx);
+	
+	/**
+	 * @brief Copies tracked PgSQL session variables to startup parameters
+	 *
+	 * This function synchronizes the current tracked session variables (in `variables` and `var_hash`)
+	 * to the startup parameters arrays (`startup_parameters` and `startup_parameters_hash`). If `copy_only_critical_param` 
+	 * is true, only the critical parameters (indices 0 to PGSQL_NAME_LAST_LOW_WM-1) are copied. 
+	 * Otherwise, all tracked variables up to PGSQL_NAME_LAST_HIGH_WM are copied.
+	 *
+	 * @param copy_only_critical_param If true, only critical parameters are copied; otherwise, all tracked variables.
+	 */
+	void copy_pgsql_variables_to_startup_parameters(bool copy_only_critical_param);
+
+	/**
+	 * @brief Copies startup parameters to tracked PgSQL session variables.
+	 *
+	 * This function synchronizes the startup parameters arrays (`startup_parameters` and `startup_parameters_hash`)
+	 * to the tracked session variables (`variables` and `var_hash`). If `copy_only_critical_param` is true,
+	 * only the critical parameters (indices 0 to PGSQL_NAME_LAST_LOW_WM-1) are copied. Otherwise, all tracked
+	 * variables up to PGSQL_NAME_LAST_HIGH_WM are copied.
+	 *
+	 * @param copy_only_critical_param If true, only critical parameters are copied; otherwise, all tracked variables.
+	 */
+	void copy_startup_parameters_to_pgsql_variables(bool copy_only_critical_param);
+
 	struct {
 		unsigned long length;
 		char* ptr;
@@ -471,6 +507,9 @@ class PgSQL_Connection {
 	bool var_absent[PGSQL_NAME_LAST_HIGH_WM] = { false };
 	std::vector<uint32_t> dynamic_variables_idx;
 
+	uint32_t startup_parameters_hash[PGSQL_NAME_LAST_HIGH_WM] = {};
+	char* startup_parameters[PGSQL_NAME_LAST_HIGH_WM] = {};
+
 	/**
 	 * @brief Keeps tracks of the 'server_status'. Do not confuse with the 'server_status' from the
 	 *  'MYSQL' connection itself. This flag keeps track of the configured server status from the
@@ -499,7 +538,7 @@ class PgSQL_Connection {
 	bool reusable;
 	bool processing_multi_statement;
 	bool multiplex_delayed;
-
+	bool is_client_connection; // true if this is a client connection, false if it is a server connection
 
 	PgSQL_SrvC *parent;
 	PgSQL_Connection_userinfo* userinfo;

include/PgSQL_Session.h

@@ -369,7 +369,6 @@ class PgSQL_Session : public Base_Session<PgSQL_Session, PgSQL_Data_Stream, PgSQ
 	PtrSize_t pkt;
 	std::string untracked_option_parameters;
 	PgSQL_DateStyle_t current_datestyle = {};
-	char* default_session_variables[PGSQL_NAME_LAST_HIGH_WM] = {};
 
 #if 0
 	// uint64_t
@@ -526,22 +525,6 @@ class PgSQL_Session : public Base_Session<PgSQL_Session, PgSQL_Data_Stream, PgSQ
 	void detected_broken_connection(const char* file, unsigned int line, const char* func, const char* action, PgSQL_Connection* myconn, bool verbose = false);
 	void generate_status_one_hostgroup(int hid, std::string& s);
 	void set_previous_status_mode3(bool allow_execute = true);
-
-	void set_default_session_variable(enum pgsql_variable_name idx, const char* value);
-
-	/**
-	 * @brief Retrieves default session variable
-	 *
-	 * This function tries to retrieve value of default session variable if present (provided in connection parameters). 
-	 * If value is not found, it falls back to the thread-specific default variables.
-	 *
-	 * @param idx The index of the session variable to retrieve.
-	 * @return The value of the session variable 
-	 * 
-	 */
-	const char* get_default_session_variable(enum pgsql_variable_name idx);
-	
-	void reset_default_session_variable(enum pgsql_variable_name idx);
 };
 
 #define PgSQL_KILL_QUERY       1

lib/Base_Thread.cpp

@@ -132,7 +132,7 @@ S Base_Thread::create_new_session_and_client_data_stream(int _fd) {
 	sess->client_myds->myprot.dump_pkt = true;
 #endif
 	if constexpr (std::is_same_v<T, PgSQL_Thread>) {
-		PgSQL_Connection* myconn = new PgSQL_Connection();
+		PgSQL_Connection* myconn = new PgSQL_Connection(true);
 		sess->client_myds->attach_connection(myconn);
 	} else if constexpr (std::is_same_v<T, MySQL_Thread>) {
 		MySQL_Connection* myconn = new MySQL_Connection();

lib/PgSQL_Connection.cpp

@@ -147,8 +147,9 @@ void print_backtrace(void);
 
 #define NEXT_IMMEDIATE(new_st) do { async_state_machine = new_st; goto handler_again; } while (0)
 
-PgSQL_Connection::PgSQL_Connection() {
+PgSQL_Connection::PgSQL_Connection(bool is_client_conn) {
 	proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 4, "Creating new PgSQL_Connection %p\n", this);
+	is_client_connection = is_client_conn;
 	pgsql_conn = NULL;
 	result_type = 0;
 	pgsql_result = NULL;
@@ -215,13 +216,6 @@ PgSQL_Connection::~PgSQL_Connection() {
 		delete query_result_reuse;
 		query_result_reuse = NULL;
 	}
-	for (int i = 0; i < PGSQL_NAME_LAST_HIGH_WM; i++) {
-		if (variables[i].value) {
-			free(variables[i].value);
-			variables[i].value = NULL;
-			var_hash[i] = 0;
-		}
-	}
 
 	if (connected_host_details.hostname) {
 		free(connected_host_details.hostname);
@@ -233,6 +227,21 @@ PgSQL_Connection::~PgSQL_Connection() {
 	}
 
 	if (options.init_connect) free(options.init_connect);
+
+	for (int i = 0; i < PGSQL_NAME_LAST_HIGH_WM; ++i) {
+		if (variables[i].value) {
+			free(variables[i].value);
+			variables[i].value = NULL;
+			var_hash[i] = 0;
+		}
+	}
+
+	for (int i = 0; i < PGSQL_NAME_LAST_HIGH_WM; ++i) {
+		if (startup_parameters[i]) {
+			free(startup_parameters[i]);
+			startup_parameters_hash[i] = 0;
+		}
+	}
 }
 
 void PgSQL_Connection::next_event(PG_ASYNC_ST new_st) {
@@ -756,44 +765,28 @@ void PgSQL_Connection::connect_start() {
 		// charset validation is already done 
 		pgsql_variables.server_set_hash_and_value(myds->sess, PGSQL_CLIENT_ENCODING, client_charset, client_charset_hash);
 
-		std::vector<unsigned int> client_options;
-		client_options.reserve(PGSQL_NAME_LAST_LOW_WM + myds->sess->client_myds->myconn->dynamic_variables_idx.size());
+		// optimized way to set client parameters on backend connection when creating a new connection
+		conninfo << "options='";
+		// excluding client_encoding, which is already set above
+		for (int idx = 1; idx < PGSQL_NAME_LAST_LOW_WM; idx++) {
+			const char* value = pgsql_variables.client_get_value(myds->sess, idx);
+			const char* escaped_str = escape_string_backslash_spaces(value);
+			conninfo << "-c " << pgsql_tracked_variables[idx].set_variable_name << "=" << escaped_str << " ";
+			if (escaped_str != value)
+				free((char*)escaped_str);
 
-		// excluding PGSQL_CLIENT_ENCODING
-		for (unsigned int idx = 1; idx < PGSQL_NAME_LAST_LOW_WM; idx++) {
-			if (pgsql_variables.client_get_hash(myds->sess, idx) == 0) continue;
-			client_options.push_back(idx);
+			const uint32_t hash = pgsql_variables.client_get_hash(myds->sess, idx);
+			pgsql_variables.server_set_hash_and_value(myds->sess, idx, value, hash);
 		}
 
-		for (uint32_t idx : myds->sess->client_myds->myconn->dynamic_variables_idx) {
-			assert(pgsql_variables.client_get_hash(myds->sess, idx));
-			client_options.push_back(idx);
-		}
+		myds->sess->mybe->server_myds->myconn->copy_pgsql_variables_to_startup_parameters(true);
 
-		if (client_options.empty() == false ||
-			myds->sess->untracked_option_parameters.empty() == false) {
-
-			// optimized way to set client parameters on backend connection when creating a new connection
-			conninfo << "options='";
-			for (int idx : client_options) {
-				const char* value = pgsql_variables.client_get_value(myds->sess, idx);
-				const char* escaped_str = escape_string_backslash_spaces(value);
-				conninfo << "-c " << pgsql_tracked_variables[idx].set_variable_name << "=" << escaped_str << " ";
-				if (escaped_str != value)
-					free((char*)escaped_str);
-
-				const uint32_t hash = pgsql_variables.client_get_hash(myds->sess, idx);
-				pgsql_variables.server_set_hash_and_value(myds->sess, idx, value, hash);
-			}
-
-			myds->sess->mybe->server_myds->myconn->reorder_dynamic_variables_idx();
-
-			// if there are untracked parameters, the session should lock on the host group
-			if (myds->sess->untracked_option_parameters.empty() == false) {
-				conninfo << myds->sess->untracked_option_parameters;
-			}
-			conninfo << "'";
+		// if there are untracked parameters, the session should lock on the host group
+		if (myds->sess->untracked_option_parameters.empty() == false) {
+			conninfo << myds->sess->untracked_option_parameters;
 		}
+		conninfo << "'";
+		
 	}
 
 	/*conninfo << "postgres://";
@@ -1881,16 +1874,19 @@ void PgSQL_Connection::reset() {
 	reusable = true;
 	creation_time = monotonic_time();
 
-	for (int i = 0; i < PGSQL_NAME_LAST_HIGH_WM; i++) {
+	for (int i = (is_client_connection ? 0 : PGSQL_NAME_LAST_LOW_WM + 1);
+		i < PGSQL_NAME_LAST_HIGH_WM; 
+		i++) {
 		var_hash[i] = 0;
 		if (variables[i].value) {
 			free(variables[i].value);
 			variables[i].value = NULL;
-			var_hash[i] = 0;
 		}
 	}
 	dynamic_variables_idx.clear();
 
+	if (!is_client_connection) copy_startup_parameters_to_pgsql_variables(/*copy_only_critical_param=*/true); 
+
 	if (options.init_connect) {
 		free(options.init_connect);
 		options.init_connect = NULL;
@@ -2202,3 +2198,68 @@ void PgSQL_Connection::set_error_from_PQerrorMessage() {
 	const std::string_view& full_msg = !primary_msg.empty() ? primary_msg : lib_errmsg;
 	PgSQL_Error_Helper::fill_error_info(error_info, sqlstate.data(), full_msg.data(), severity.data());
 }
+
+std::pair<const char*, uint32_t> PgSQL_Connection::get_startup_parameter_and_hash(enum pgsql_variable_name idx) {
+	// within valid range?
+	assert(idx >= 0 && idx < PGSQL_NAME_LAST_HIGH_WM);
+
+	// Attempt to retrieve value from default startup parameters
+	if (startup_parameters_hash[idx] != 0) {
+		assert(startup_parameters[idx]);
+		return { startup_parameters[idx], startup_parameters_hash[idx] };
+	}
+	// fall back to thread-specific default
+	return { pgsql_thread___default_variables[idx], 0 };
+}
+
+void PgSQL_Connection::copy_pgsql_variables_to_startup_parameters(bool copy_only_critical_param) {
+
+	//memcpy(startup_parameters_hash, var_hash, sizeof(uint32_t) * PGSQL_NAME_LAST_LOW_WM);
+	for (int i = 0; i < PGSQL_NAME_LAST_LOW_WM; ++i) {
+		assert(var_hash[i]);
+		assert(variables[i].value);
+		startup_parameters_hash[i] = var_hash[i];
+		free(startup_parameters[i]);
+		startup_parameters[i] = strdup(variables[i].value);
+	}
+
+	if (copy_only_critical_param) return;
+
+	for (int i = PGSQL_NAME_LAST_LOW_WM + 1; i < PGSQL_NAME_LAST_HIGH_WM; i++) {
+		if (var_hash[i] != 0) {
+			startup_parameters_hash[i] = var_hash[i];
+			free(startup_parameters[i]);
+			startup_parameters[i] = strdup(variables[i].value);
+		} else {
+			startup_parameters_hash[i] = 0;
+			free(startup_parameters[i]);
+			startup_parameters[i] = nullptr;
+		}
+	}
+}
+
+void PgSQL_Connection::copy_startup_parameters_to_pgsql_variables(bool copy_only_critical_param) {
+
+	//memcpy(var_hash, startup_parameters_hash, sizeof(uint32_t) * PGSQL_NAME_LAST_LOW_WM);
+	for (int i = 0; i < PGSQL_NAME_LAST_LOW_WM; i++) {
+		assert(startup_parameters_hash[i]);
+		assert(startup_parameters[i]);
+		var_hash[i] = startup_parameters_hash[i];
+		free(variables[i].value);
+		variables[i].value = strdup(startup_parameters[i]);
+	}
+
+	if (copy_only_critical_param) return;
+
+	for (int i = PGSQL_NAME_LAST_LOW_WM + 1; i < PGSQL_NAME_LAST_HIGH_WM; i++) {
+		if (startup_parameters_hash[i]) {
+			var_hash[i] = startup_parameters_hash[i];
+			free(variables[i].value);
+			variables[i].value = strdup(startup_parameters[i]);
+		} else {
+			var_hash[i] = 0;
+			free(variables[i].value);
+			variables[i].value = nullptr;
+		}
+	}
+}

lib/PgSQL_HostGroups_Manager.cpp

@@ -2298,7 +2298,7 @@ PgSQL_Connection * PgSQL_SrvConnList::get_random_MyConn(PgSQL_Session *sess, boo
 						}
 
 						// we must create a new connection
-						conn = new PgSQL_Connection();
+						conn = new PgSQL_Connection(false);
 						conn->parent=mysrvc;
 						// if attributes.multiplex == true , STATUS_MYSQL_CONNECTION_NO_MULTIPLEX_HG is set to false. And vice-versa
 						conn->set_status(!conn->parent->myhgc->attributes.multiplex, STATUS_MYSQL_CONNECTION_NO_MULTIPLEX_HG);
@@ -2312,7 +2312,7 @@ PgSQL_Connection * PgSQL_SrvConnList::get_random_MyConn(PgSQL_Session *sess, boo
 					unsigned int conns_free = mysrvc->ConnectionsFree->conns_length();
 					unsigned int conns_used = mysrvc->ConnectionsUsed->conns_length();
 					if ((conns_used > conns_free) && (mysrvc->max_connections > (conns_free/2 + conns_used/2)) ) {
-						conn = new PgSQL_Connection();
+						conn = new PgSQL_Connection(false);
 						conn->parent=mysrvc;
 						// if attributes.multiplex == true , STATUS_MYSQL_CONNECTION_NO_MULTIPLEX_HG is set to false. And vice-versa
 						conn->set_status(!conn->parent->myhgc->attributes.multiplex, STATUS_MYSQL_CONNECTION_NO_MULTIPLEX_HG);
@@ -2357,7 +2357,7 @@ PgSQL_Connection * PgSQL_SrvConnList::get_random_MyConn(PgSQL_Session *sess, boo
 			__sync_fetch_and_add(&PgHGM->status.server_connections_delayed, 1);
 			return NULL;
 		} else {
-			conn = new PgSQL_Connection();
+			conn = new PgSQL_Connection(false);
 			conn->parent=mysrvc;
 			// if attributes.multiplex == true , STATUS_MYSQL_CONNECTION_NO_MULTIPLEX_HG is set to false. And vice-versa
 			conn->set_status(!conn->parent->myhgc->attributes.multiplex, STATUS_MYSQL_CONNECTION_NO_MULTIPLEX_HG);

lib/PgSQL_Protocol.cpp

@@ -1211,11 +1211,9 @@ EXECUTION_STATE PgSQL_Protocol::process_handshake_response_packet(unsigned char*
 					if (pgsql_variables.client_set_value(sess, PGSQL_DATESTYLE, datestyle.c_str(), false)) {
 						// change current datestyle
 						sess->current_datestyle = PgSQL_DateStyle_Util::parse_datestyle(datestyle);
-						sess->set_default_session_variable(PGSQL_DATESTYLE, datestyle.c_str());
 					}
 				} else {
 					pgsql_variables.client_set_value(sess, idx, value_copy.c_str(), false);
-					sess->set_default_session_variable((enum pgsql_variable_name)idx, value_copy.c_str());
 				}
 			} else {
 				// parameter provided is not part of the tracked variables. Will lock on hostgroup on next query.
@@ -1234,10 +1232,10 @@ EXECUTION_STATE PgSQL_Protocol::process_handshake_response_packet(unsigned char*
 				continue;
 			const char* val = pgsql_thread___default_variables[i];
 			pgsql_variables.client_set_value(sess, i, val, false);
-			sess->set_default_session_variable((pgsql_variable_name)i, val);
 		}
 
 		sess->client_myds->myconn->reorder_dynamic_variables_idx();
+		sess->client_myds->myconn->copy_pgsql_variables_to_startup_parameters(false);
 	}
 	else {
 		// we always duplicate username and password, or crashes happen