Enable session tracking based on server capabilities

Signed-off-by: Wazir Ahmed <wazir@proxysql.com>

include/MySQL_HostGroups_Manager.h

@@ -227,6 +227,15 @@ class MySrvC {	// MySQL Server Container
 	bool shunned_and_kill_all_connections; // if a serious failure is detected, this will cause all connections to die even if the server is just shunned
 	int32_t use_ssl;
 	char *comment;
+
+	// 'server_backoff_time' stores a timestamp that prevents the server from being
+	// considered for random selection ('MyHGC::get_random_MySrvC') until that time passes.
+	//
+	// This is primarily used when `session_track_variables::ENFORCED` mode is active.
+	// If a server lacks the required capabilities in this mode, it is temporarily
+	// excluded from selection for a specified duration.
+	unsigned long long server_backoff_time;
+
 	MySrvConnList *ConnectionsUsed;
 	MySrvConnList *ConnectionsFree;
 	/**

include/MySQL_Session.h

@@ -554,6 +554,7 @@ class MySQL_Session: public Base_Session<MySQL_Session, MySQL_Data_Stream, MySQL
 	void reset_warning_hostgroup_flag_and_release_connection();
 	void set_previous_status_mode3(bool allow_execute=true);
 	char* get_current_query(int max_length = -1);
+	bool handle_session_track_capabilities();
 
 	friend void SQLite3_Server_session_handler(MySQL_Session*, void *_pa, PtrSize_t *pkt);
 

include/MySQL_Thread.h

@@ -369,8 +369,13 @@ struct th_metrics_map_idx {
  */
 struct session_track_variables {
 	enum mode {
-		DISABLED = 0,	///< Session variable tracking is disabled (default)
-		ENABLED			///< Session variable tracking is enabled
+		// Disabled; default mode
+		DISABLED = 0,
+		// Enable session tracking if backend supports it
+		OPTIONAL,
+		// Enforce session tracking; connection fails if backend does
+		// not support CLIENT_DEPRECATE_EOF and CLIENT_SESSION_TRACKING
+		ENFORCED
 	};
 };
 

lib/MyHGC.cpp

@@ -38,6 +38,10 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_
 		for (j=0; j<l; j++) {
 			mysrvc=mysrvs->idx(j);
 			if (mysrvc->get_status() == MYSQL_SERVER_STATUS_ONLINE) { // consider this server only if ONLINE
+				// skip servers that are in backoff period
+				if (mysrvc->server_backoff_time > sess->thread->curtime)
+					continue;
+
 				if (mysrvc->myhgc->num_online_servers.load(std::memory_order_relaxed) <= mysrvc->myhgc->attributes.max_num_online_servers) { // number of online servers in HG is within configured range
 					if (mysrvc->ConnectionsUsed->conns_length() < mysrvc->max_connections) { // consider this server only if didn't reach max_connections
 						if (mysrvc->current_latency_us < (mysrvc->max_latency_us ? mysrvc->max_latency_us : mysql_thread___default_max_latency_ms*1000)) { // consider the host only if not too far

lib/MySQL_Protocol.cpp

@@ -1111,9 +1111,13 @@ bool MySQL_Protocol::generate_pkt_initial_handshake(bool send, void **ptr, unsig
 	// variable. This is the first step of ensuring that client connections doesn't
 	// enable 'CLIENT_DEPRECATE_EOF' unless explicitly stated by 'mysql-enable_client_deprecate_eof'.
 	// Second step occurs during client handshake response (process_pkt_handshake_response).
-	if (deprecate_eof_active && mysql_thread___enable_client_deprecate_eof) {
-		extended_capabilities |= CLIENT_DEPRECATE_EOF;
+	if (deprecate_eof_active) {
+		if (mysql_thread___enable_client_deprecate_eof
+			|| mysql_thread___session_track_variables == session_track_variables::ENFORCED) {
+			extended_capabilities |= CLIENT_DEPRECATE_EOF;
+		}
 	}
+
 	// Copy the 'capability_flags_2'
 	uint16_t upper_word = static_cast<uint16_t>(extended_capabilities >> 16);
 	memcpy(_ptr+l, static_cast<void*>(&upper_word), sizeof(upper_word)); l += sizeof(upper_word);

lib/MySQL_Session.cpp

@@ -1915,10 +1915,12 @@ bool MySQL_Session::handler_again___verify_backend_session_track_gtids() {
 	proxy_debug(PROXY_DEBUG_MYSQL_CONNECTION, 5, "Session %p , client: %s , backend: %s\n", this, client_myds->myconn->options.session_track_gtids, mybe->server_myds->myconn->options.session_track_gtids);
 	// we first verify that the backend supports it
 	// if backend is old (or if it is not mysql) ignore this setting
-	if ((mybe->server_myds->myconn->mysql->server_capabilities & CLIENT_SESSION_TRACKING) == 0) {
-		// the backend doesn't support CLIENT_SESSION_TRACKING
-		return ret; // exit immediately
+	if ((mybe->server_myds->myconn->mysql->server_capabilities & CLIENT_SESSION_TRACKING) == 0
+		|| (mybe->server_myds->myconn->mysql->server_capabilities & CLIENT_DEPRECATE_EOF) == 0
+		|| mysql_thread___enable_server_deprecate_eof == false) {
+		return ret;
 	}
+
 	uint32_t b_int = mybe->server_myds->myconn->options.session_track_gtids_int;
 	uint32_t f_int = client_myds->myconn->options.session_track_gtids_int;
 
@@ -1965,16 +1967,26 @@ bool MySQL_Session::handler_again___verify_backend_session_track_gtids() {
 }
 
 bool MySQL_Session::handler_again___verify_backend_session_track_variables() {
-	if (mysql_thread___session_track_variables == session_track_variables::DISABLED) {
+	int mode = mysql_thread___session_track_variables;
+
+	// skip enabling session variable tracking in the following cases
+	if (mode == session_track_variables::DISABLED) {
+		return false;
+	}
+	if ((mybe->server_myds->myconn->mysql->server_capabilities & CLIENT_SESSION_TRACKING) == 0
+		|| (mybe->server_myds->myconn->mysql->server_capabilities & CLIENT_DEPRECATE_EOF) == 0) {
+		return false;
+	}
+	if (!mysql_thread___enable_server_deprecate_eof && mode != session_track_variables::ENFORCED) {
 		return false;
 	}
 
+	// enable session tracking
 	if (mybe->server_myds->myconn->options.session_track_variables_sent == false) {
 		mybe->server_myds->myconn->options.session_track_variables_sent = true;
 		set_previous_status_mode3();
 		NEXT_IMMEDIATE_NEW(SETTING_SESSION_TRACK_VARIABLES);
 	}
-
 	if (mybe->server_myds->myconn->options.session_track_state_sent == false) {
 		mybe->server_myds->myconn->options.session_track_state_sent = true;
 		set_previous_status_mode3();
@@ -2944,6 +2956,11 @@ bool MySQL_Session::handler_again___status_CONNECTING_SERVER(int *_rc) {
 		}
 		enum session_status st=status;
 		if (mybe->server_myds->myconn->async_state_machine==ASYNC_IDLE) {
+			if (handle_session_track_capabilities() == false) {
+				pause_until = thread->curtime + mysql_thread___connect_retries_delay*1000;
+				return false;
+			}
+
 			st=previous_status.top();
 			previous_status.pop();
 			NEXT_IMMEDIATE_NEW(st);
@@ -2973,6 +2990,13 @@ bool MySQL_Session::handler_again___status_CONNECTING_SERVER(int *_rc) {
 				previous_status.pop();
 				myds->wait_until=0;
 
+				if (handle_session_track_capabilities() == false) {
+					previous_status.push(st);
+					pause_until = thread->curtime + mysql_thread___connect_retries_delay * 1000;
+					set_status(CONNECTING_SERVER);
+					return false;
+				}
+
 				// NOTE: Even if a connection has correctly been created, since the CLIENT_DEPRECATE_EOF
 				// capability isn't always enforced to match for backend conns (no direct propagation), a
 				// mismatch can take place after the creation. Right now this is only true for
@@ -8432,3 +8456,57 @@ char* MySQL_Session::get_current_query(int max_length) {
 
 	return res;
 }
+
+/**
+ * @brief Handle session track capabilities validation.
+ *
+ * This function validates whether the backend connection has capabilities such as 'CLIENT_DEPRECATE_EOF'
+ * and 'CLIENT_SESSION_TRACKING' which are required to enable 'session_track_system_variables' in a MySQL session.
+ *
+ * If the connection lacks the capabilities and ProxySQL configuration is set in 'ENFORCED' mode, it returns the
+ * connection to the pool and set a backoff time for the backend server. This backoff time prevents the server from
+ * being selected again during connection pooling.
+ *
+ * @return 'true' if backend connection has required capabilities, otherwise returns 'false'.
+ */
+bool MySQL_Session::handle_session_track_capabilities() {
+	if (mysql_thread___session_track_variables != session_track_variables::ENFORCED) {
+		return true;
+	}
+
+	// this function should not be called in these states
+	if (client_myds == NULL
+		|| client_myds->myconn == NULL
+		|| mybe == NULL
+		|| mybe->server_myds == NULL
+		|| mybe->server_myds->myconn == NULL
+		|| mybe->server_myds->myconn->mysql == NULL) {
+		return true;
+	}
+
+	MySQL_Connection *be_conn = mybe->server_myds->myconn;
+	unsigned long srv_cap = be_conn->mysql->server_capabilities;
+	uint32_t client_flag = client_myds->myconn->options.client_flag;
+
+	bool client_support_session_track = ((client_flag & CLIENT_SESSION_TRACKING) != 0 && (client_flag & CLIENT_DEPRECATE_EOF) != 0);
+	bool server_support_session_track = ((srv_cap & CLIENT_SESSION_TRACKING) != 0 && (srv_cap & CLIENT_DEPRECATE_EOF) != 0);
+
+	// In fast forward, if client does not support session tracking,
+	// then ProxySQL do not have to enforce session track on backend connections
+	if ((session_fast_forward == SESSION_FORWARD_TYPE_PERMANENT) && !client_support_session_track) {
+		return true;
+	}
+
+	if (!server_support_session_track) {
+		// be_conn->parent->server_backoff_time = thread->curtime + (600 * 1000000); // 10 minutes
+		be_conn->parent->server_backoff_time = thread->curtime + (30 * 1000000); // 30 seconds
+		if (session_fast_forward) {
+			mybe->server_myds->destroy_MySQL_Connection_From_Pool(false);
+		} else {
+			mybe->server_myds->return_MySQL_Connection_To_Pool();
+		}
+		return false;
+	}
+
+	return true;
+}

lib/MySQL_Thread.cpp

@@ -2361,7 +2361,7 @@ char ** MySQL_Threads_Handler::get_variables_list() {
 		VariablesPointers_int["eventslog_format"] = make_tuple(&variables.eventslog_format, 0, 0, true);
 		VariablesPointers_int["wait_timeout"]     = make_tuple(&variables.wait_timeout,     0, 0, true);
 		VariablesPointers_int["data_packets_history_size"] = make_tuple(&variables.data_packets_history_size, 0, 0, true);
-		VariablesPointers_int["session_track_variables"]   = make_tuple(&variables.session_track_variables,   0, 1, false);
+		VariablesPointers_int["session_track_variables"]   = make_tuple(&variables.session_track_variables,   0, 2, false);
 	}
 
 
@@ -5643,7 +5643,11 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local(unsigned int _hid, MySQL_Sessi
 	std::vector<MySrvC *> parents; // this is a vector of srvers that needs to be excluded in case gtid_uuid is used
 	MySQL_Connection *c=NULL;
 	for (i=0; i<cached_connections->len; i++) {
-		c=(MySQL_Connection *)cached_connections->index(i);
+		c = (MySQL_Connection *) cached_connections->index(i);
+		// skip servers that are in backoff period
+		if (c->parent->server_backoff_time > curtime)
+			continue;
+
 		if (c->parent->myhgc->hid==_hid && sess->client_myds->myconn->match_tracked_options(c)) { // options are all identical
 			if (
 				(gtid_uuid == NULL) || // gtid_uuid is not used

lib/MySrvC.cpp

@@ -31,6 +31,7 @@ MySrvC::MySrvC(
 	bytes_recv=0;
 	max_connections_used=0;
 	queries_gtid_sync=0;
+	server_backoff_time = 0;
 	time_last_detected_error=0;
 	connect_ERR_at_time_last_detected_error=0;
 	shunned_automatic=false;

lib/mysql_connection.cpp

@@ -889,13 +889,18 @@ void MySQL_Connection::connect_start_SetClientFlag(unsigned long& client_flags)
 		}
 	}
 
-	// set 'CLIENT_DEPRECATE_EOF' flag if explicitly stated by 'mysql-enable_server_deprecate_eof'.
-	// Capability is disabled by default in 'mariadb_client', so setting this option is not optional
-	// for having 'CLIENT_DEPRECATE_EOF' in the connection to be stablished.
+	// 'CLIENT_DEPRECATE_EOF' capability is disabled by default in mariadb_client.
+	// Based on the value of 'mysql-enable_server_deprecate_eof', enable this
+	// capability in a new connection.
 	if (mysql_thread___enable_server_deprecate_eof) {
 		mysql->options.client_flag |= CLIENT_DEPRECATE_EOF;
 	}
 
+	// override 'mysql-enable_server_deprecate_eof' behavior if 'session_track_variables' is set to 'ENFORCED'
+	if (mysql_thread___session_track_variables == session_track_variables::ENFORCED) {
+		mysql->options.client_flag |= CLIENT_DEPRECATE_EOF;
+	}
+
 	if (myds != NULL) {
 		if (myds->sess != NULL) {
 			if (myds->sess->session_fast_forward) { // this is a fast_forward connection