Adding new api to query the stage of the mysql_real_connect* state machine

Summary:
For nonblocking connect we need a way to query what stage of the mysql_real_connect_nonblocking state machine is at.
This will allow users to monitor and take appropriate actions based on the progress of the connection.

Reviewed By: lloyd

Differential Revision: D23387197

Author: aditya-jalan

include/mysql.h

@@ -528,6 +528,8 @@ int STDCALL mysql_session_track_get_next(MYSQL *mysql,
                                          const char **data, size_t *length);
 int STDCALL mysql_resp_attr_find(MYSQL *mysql, const char *lookup,
                                  const char **data, size_t *length);
+
+enum connect_stage STDCALL mysql_get_connect_stage(MYSQL *mysql);
 /* local infile support */
 
 #define LOCAL_INFILE_ERROR_LEN 512

include/mysql.h.pp

@@ -187,6 +187,21 @@
   MYSQL_OPTION_MULTI_STATEMENTS_ON,
   MYSQL_OPTION_MULTI_STATEMENTS_OFF
 };
+enum connect_stage {
+  CONNECT_STAGE_INVALID = 0,
+  CONNECT_STAGE_NOT_STARTED,
+  CONNECT_STAGE_NET_BEGIN_CONNECT,
+  CONNECT_STAGE_NET_WAIT_CONNECT,
+  CONNECT_STAGE_NET_COMPLETE_CONNECT,
+  CONNECT_STAGE_READ_GREETING,
+  CONNECT_STAGE_PARSE_HANDSHAKE,
+  CONNECT_STAGE_ESTABLISH_SSL,
+  CONNECT_STAGE_AUTHENTICATE,
+  CONNECT_STAGE_PREP_SELECT_DATABASE,
+  CONNECT_STAGE_PREP_INIT_COMMANDS,
+  CONNECT_STAGE_SEND_ONE_INIT_COMMAND,
+  CONNECT_STAGE_COMPLETE
+};
 enum enum_session_state_type {
   SESSION_TRACK_SYSTEM_VARIABLES,
   SESSION_TRACK_SCHEMA,
@@ -681,6 +696,7 @@
                                          const char **data, size_t *length);
 int mysql_resp_attr_find(MYSQL *mysql, const char *lookup,
                                  const char **data, size_t *length);
+enum connect_stage mysql_get_connect_stage(MYSQL *mysql);
 void mysql_set_local_infile_handler(
     MYSQL *mysql, int (*local_infile_init)(void **, const char *, void *),
     int (*local_infile_read)(void *, char *, unsigned int),

include/mysql_com.h

@@ -1070,6 +1070,25 @@ enum enum_mysql_set_option {
   MYSQL_OPTION_MULTI_STATEMENTS_OFF
 };
 
+/*
+  Describes the current state of Asynchronous connection phase state machine
+*/
+enum connect_stage {
+  CONNECT_STAGE_INVALID = 0,
+  CONNECT_STAGE_NOT_STARTED,
+  CONNECT_STAGE_NET_BEGIN_CONNECT,
+  CONNECT_STAGE_NET_WAIT_CONNECT,
+  CONNECT_STAGE_NET_COMPLETE_CONNECT,
+  CONNECT_STAGE_READ_GREETING,
+  CONNECT_STAGE_PARSE_HANDSHAKE,
+  CONNECT_STAGE_ESTABLISH_SSL,
+  CONNECT_STAGE_AUTHENTICATE,
+  CONNECT_STAGE_PREP_SELECT_DATABASE,
+  CONNECT_STAGE_PREP_INIT_COMMANDS,
+  CONNECT_STAGE_SEND_ONE_INIT_COMMAND,
+  CONNECT_STAGE_COMPLETE
+};
+
 /**
   Type of state change information that the server can include in the Ok
   packet.

libmysql/CMakeLists.txt

@@ -189,7 +189,8 @@ SET(CLIENT_API_NONBLOCKING_FUNCTIONS
   mysql_send_query_nonblocking
   mysql_store_result_nonblocking
   mysql_get_socket_descriptor
-
+  mysql_get_connect_stage
+  mysql_reset_connection_nonblocking
   CACHE INTERNAL "Nonblocking functions exported by client API"
 )
 

sql-common/client.cc

@@ -4970,6 +4970,45 @@ int STDCALL mysql_get_socket_descriptor(MYSQL *mysql) {
   return -1;
 }
 
+static const std::vector<std::pair<csm_function, connect_stage>> stages = {
+    {csm_begin_connect, CONNECT_STAGE_NET_BEGIN_CONNECT},
+    {csm_wait_connect, CONNECT_STAGE_NET_WAIT_CONNECT},
+    {csm_complete_connect, CONNECT_STAGE_NET_COMPLETE_CONNECT},
+    {csm_read_greeting, CONNECT_STAGE_READ_GREETING},
+    {csm_parse_handshake, CONNECT_STAGE_PARSE_HANDSHAKE},
+    {csm_establish_ssl, CONNECT_STAGE_ESTABLISH_SSL},
+    {csm_authenticate, CONNECT_STAGE_AUTHENTICATE},
+    {csm_prep_select_database, CONNECT_STAGE_PREP_SELECT_DATABASE},
+#if !defined(MYSQL_SERVER)
+    {csm_prep_init_commands, CONNECT_STAGE_PREP_INIT_COMMANDS},
+    {csm_send_one_init_command, CONNECT_STAGE_SEND_ONE_INIT_COMMAND},
+#endif
+};
+
+connect_stage STDCALL mysql_get_connect_stage(MYSQL *mysql) {
+  if (mysql) {
+    NET *net = &mysql->net;
+    if (!net->vio) {
+      /* Seems the first stage hasn't started yet or it was unsuccessful, and
+         needs to be restarted so return the 1st stage */
+      return CONNECT_STAGE_NOT_STARTED;
+    }
+
+    mysql_async_connect *ctx = ASYNC_DATA(mysql)->connect_context;
+    if (!ctx) {
+      // If context is null and vio->net is set, means connection is complete
+      return CONNECT_STAGE_COMPLETE;
+    }
+
+    for (const auto &stage_pair : stages) {
+      if (ctx->state_function == stage_pair.first) {
+        return stage_pair.second;
+      }
+    }
+  }
+  return CONNECT_STAGE_INVALID;
+}
+
 /* clang-format off */
 /**
   @page page_protocol_connection_phase_packets_protocol_handshake_response Protocol::HandshakeResponse:

testclients/mysql_client_test.cc

@@ -23466,6 +23466,106 @@ static void test_bug25584097() {
   thd.join();
 }
 
+static void test_get_connect_stage() {
+  MYSQL *mysql_async = NULL, *mysql_sync = NULL;
+  net_async_status status;
+
+  myheader("test_get_connect_stage");
+
+  enum connect_stage cs = mysql_get_connect_stage(mysql_async);
+  if (cs != CONNECT_STAGE_INVALID) {
+    fprintf(stdout,
+            "\n Expected connect_stage to be "
+            "CONNECT_STAGE_INVALID for uninitialized mysql");
+    exit(1);
+  }
+
+  /* test stages for a nonblocking conneciton */
+  if (!(mysql_async = mysql_client_init(NULL))) {
+    myerror("mysql_client_init() failed");
+    exit(1);
+  }
+  cs = mysql_get_connect_stage(mysql_async);
+  if (cs != CONNECT_STAGE_NOT_STARTED) {
+    fprintf(stdout,
+            "\n Expected connect_stage to be "
+            "CONNECT_STAGE_NET_BEGIN_CONNECT for just initialized connection"
+            "but its %d",
+            cs);
+    exit(1);
+  }
+
+  enum connect_stage cs_prev = cs;
+  do {
+    status = mysql_real_connect_nonblocking(
+        mysql_async, opt_host, opt_user, opt_password, current_db, opt_port,
+        opt_unix_socket, CLIENT_MULTI_STATEMENTS);
+    cs = mysql_get_connect_stage(mysql_async);
+    // Always expect state machine to make forward progress
+    if (cs < cs_prev) {
+      fprintf(stdout,
+              "\n Did not expect connect_stage to be %d, as previous was %d",
+              cs, cs_prev);
+      exit(1);
+    } else if (cs != cs_prev) {
+      if (!opt_silent)
+        fprintf(stdout,
+                "\n Nonblocking connect made transition from stage %d to %d",
+                cs_prev, cs);
+    }
+    cs_prev = cs;
+  } while (status == NET_ASYNC_NOT_READY);
+  if (status == NET_ASYNC_ERROR) {
+    fprintf(stdout, "\n mysql_real_connect_nonblocking() failed. stage:%d", cs);
+    exit(1);
+  }
+
+  cs = mysql_get_connect_stage(mysql_async);
+  if (cs != CONNECT_STAGE_COMPLETE) {
+    fprintf(stdout,
+            "\n Expected connect_stage as CONNECT_STAGE_COMPLETE, its %d", cs);
+    exit(1);
+  }
+  if (!opt_silent)
+    fprintf(stdout, "\n Nonblocking connect successful. Final stage %d", cs);
+
+  mysql_close(mysql_async);
+
+  /* test stages for a blocking conneciton */
+  if (!(mysql_sync = mysql_client_init(NULL))) {
+    myerror("mysql_client_init() failed");
+    exit(1);
+  }
+  cs = mysql_get_connect_stage(mysql_sync);
+  if (cs != CONNECT_STAGE_NOT_STARTED) {
+    fprintf(stdout,
+            "\n Sync:Expected connect_stage to be "
+            "CONNECT_STAGE_NOT_STARTED for just initialized connection"
+            "but its %d",
+            cs);
+    exit(1);
+  }
+
+  if (!opt_silent)
+    fprintf(stdout, "\n Starting blocking connect. Starting stage %d", cs);
+  if (!(mysql_real_connect(mysql_sync, opt_host, opt_user, opt_password,
+                           current_db, opt_port, opt_unix_socket, 0))) {
+    myerror("connection failed");
+    exit(1);
+  }
+  cs = mysql_get_connect_stage(mysql_sync);
+  if (cs != CONNECT_STAGE_COMPLETE) {
+    fprintf(stdout,
+            "\n Sync:Expected connect_stage as CONNECT_STAGE_COMPLETE, its %d",
+            cs);
+    exit(1);
+  }
+  if (!opt_silent)
+    fprintf(stdout, "\n Blocking connect successful. Final stage %d", cs);
+
+  mysql_close(mysql_sync);
+}
+
 static struct my_tests_st my_tests[] = {
     {"test_bug5194", test_bug5194},
     {"disable_query_logs", disable_query_logs},
@@ -23782,6 +23882,7 @@ static struct my_tests_st my_tests[] = {
     {"test_wl13128", test_wl13128},
     {"test_bug25584097", test_bug25584097},
     {"test_34556764", test_34556764},
+    {"test_get_connect_stage", test_get_connect_stage},
     {nullptr, nullptr}};
 
 static struct my_tests_st *get_my_tests() { return my_tests; }