Merge pull request #5060 from sysown/v3.0_session_param_reset_connection

Fix Backend Connection Reset and Parameter Handling for RESET/SET TO DEFAULT

Author: renecannao

include/PgSQL_Connection.h

@@ -216,7 +216,7 @@ class PgSQL_Conn_Param {
 
 class PgSQL_Variable {
 public:
-	char *value = (char*)"";
+	char *value = nullptr;
 	void fill_server_internal_session(nlohmann::json &j, int conn_num, int idx);
 	void fill_client_internal_session(nlohmann::json &j, int idx);
 };
@@ -244,7 +244,7 @@ class PgSQL_Connection_userinfo {
 
 class PgSQL_Connection {
 public:
-	PgSQL_Connection();
+	explicit 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;
@@ -464,13 +500,16 @@ class PgSQL_Connection {
 		unsigned long long pgconnpoll_put;
 	} statuses;
 
-	PgSQL_Variable variables[PGSQL_NAME_LAST_HIGH_WM];
-	uint32_t var_hash[PGSQL_NAME_LAST_HIGH_WM];
+	std::array<PgSQL_Variable, PGSQL_NAME_LAST_HIGH_WM> variables = {};
+	std::array<uint32_t, PGSQL_NAME_LAST_HIGH_WM> var_hash = {};
 	// for now we store possibly missing variables in the lower range
 	// we may need to fix that, but this will cost performance
-	bool var_absent[PGSQL_NAME_LAST_HIGH_WM] = { false };
+	std::array<bool, PGSQL_NAME_LAST_HIGH_WM> var_absent = {};
 	std::vector<uint32_t> dynamic_variables_idx;
 
+	std::array<uint32_t, PGSQL_NAME_LAST_HIGH_WM> startup_parameters_hash = {};
+	std::array<char*, PGSQL_NAME_LAST_HIGH_WM> startup_parameters = {};
+
 	/**
 	 * @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

include/PgSQL_Thread.h

@@ -957,7 +957,7 @@ class PgSQL_Threads_Handler
 		char* ldap_user_variable;
 		char* add_ldap_user_comment;
 		char* default_session_track_gtids;
-		char* default_variables[PGSQL_NAME_LAST_HIGH_WM];
+		char* default_variables[PGSQL_NAME_LAST_LOW_WM];
 		char* firewall_whitelist_errormsg;
 #ifdef DEBUG
 		bool session_debug;

include/proxysql_structs.h

@@ -1192,7 +1192,7 @@ __thread int pgsql_thread___query_cache_soft_ttl_pct;
 __thread int pgsql_thread___query_cache_handle_warnings;
 
 __thread bool pgsql_thread___session_idle_show_processlist;
-__thread char* pgsql_thread___default_variables[PGSQL_NAME_LAST_HIGH_WM];
+__thread char* pgsql_thread___default_variables[PGSQL_NAME_LAST_LOW_WM];
 __thread int pgsql_thread___handle_unknown_charset;
 //---------------------------
 
@@ -1497,7 +1497,7 @@ extern __thread int pgsql_thread___query_cache_soft_ttl_pct;
 extern __thread int pgsql_thread___query_cache_handle_warnings;
 
 extern __thread bool pgsql_thread___session_idle_show_processlist;
-extern __thread char* pgsql_thread___default_variables[PGSQL_NAME_LAST_HIGH_WM];
+extern __thread char* pgsql_thread___default_variables[PGSQL_NAME_LAST_LOW_WM];
 extern __thread int pgsql_thread___handle_unknown_charset;
 //---------------------------
 

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;
@@ -179,10 +180,10 @@ PgSQL_Connection::PgSQL_Connection() {
 	options.init_connect_sent = false;
 	userinfo = new PgSQL_Connection_userinfo();
 
-	for (int i = 0; i < PGSQL_NAME_LAST_HIGH_WM; i++) {
-		variables[i].value = NULL;
-		var_hash[i] = 0;
-	}
+	//for (int i = 0; i < PGSQL_NAME_LAST_HIGH_WM; i++) {
+	//	variables[i].value = NULL;
+	//	var_hash[i] = 0;
+	//}
 
 	new_result = true;
 	is_copy_out = false;
@@ -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,22 @@ 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[i] = nullptr;
+			startup_parameters_hash[i] = 0;
+		}
+	}
 }
 
 void PgSQL_Connection::next_event(PG_ASYNC_ST new_st) {
@@ -756,44 +766,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);
-		}
-
-		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);
+		myds->sess->mybe->server_myds->myconn->copy_pgsql_variables_to_startup_parameters(true);
 
-				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 +1875,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 +2199,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] };
+	}
+	assert(!(idx < PGSQL_NAME_LAST_LOW_WM));
+	return { "", 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;
+		}
+	}
+}