Refactor ac_info switch

Summary: In preparation for using `thd->ac_node->ac_info` for connection accounting, this change refactors most of the logic of `set_session_db_helper()` into the class `Ac_switch_guard`. The main motivation is that connection accounting first obtains connection for the new db, and if successful, releases connection for the old db. So the existing `ac_info` needs to be saved, and either released at the end if the db switch is successful, or put back info `ac_node` and the new one released if switch failed.

Reviewed By: lth

Differential Revision: D24319962

fbshipit-source-id: 80dca3a0b71

sql/sql_multi_tenancy.cc

@@ -295,12 +295,10 @@ int multi_tenancy_admit_query(THD *thd,
   return 0;
 }
 
-
-/*
+/**
  * Exit a query based on database entity
  *
  * @param thd THD structure
- * @param attrs Resource attributes
  *
  * @return 0 if successful, 1 if otherwise
  */
@@ -328,7 +326,7 @@ int multi_tenancy_exit_query(THD *thd)
  * @param type Resource type
  * @param attrs Resource attributes
  *
- * @return resource entity name. Emtpy string if no plugin is installed.
+ * @return resource entity name. Empty string if no plugin is installed.
  */
 std::string multi_tenancy_get_entity(
     THD *thd, MT_RESOURCE_TYPE type, const MT_RESOURCE_ATTRS *attrs)
@@ -1027,6 +1025,35 @@ void AC::remove(const char* entity) {
   mysql_rwlock_unlock(&LOCK_ac);
 }
 
+/**
+  @brief Constructor.
+*/
+Ac_switch_guard::Ac_switch_guard(THD *_thd, const char *_new_db) {
+  thd = _thd;
+  if (_new_db)
+    new_db = _new_db;
+  if (thd->db)
+    old_db = thd->db;
+
+  // Release old ac_info.
+  if (thd->is_in_ac) {
+    // Current admission control is per database.
+    DBUG_ASSERT(!old_db.empty());
+    multi_tenancy_exit_query(thd);
+  }
+}
+
+/**
+  @brief Destructor commits or rolls back the switch.
+*/
+Ac_switch_guard::~Ac_switch_guard() {
+  std::string &db = committed ? old_db : new_db;
+  if (!db.empty()) {
+    MT_RESOURCE_ATTRS attrs = { nullptr, nullptr, db.c_str() };
+    multi_tenancy_close_connection(thd, &attrs);
+  }
+}
+
 ST_FIELD_INFO admission_control_queue_fields_info[]=
 {
   {"SCHEMA_NAME", NAME_LEN, MYSQL_TYPE_STRING, 0, 0, 0, SKIP_OPEN_TABLE},

sql/sql_multi_tenancy.h

@@ -352,6 +352,20 @@ class AC {
   }
 };
 
+/**
+  @brief Class temporarily holding existing resources during ac_info switch.
+*/
+class Ac_switch_guard {
+  THD *thd;
+  std::string old_db;
+  std::string new_db;
+  bool committed = false;
+
+public:
+  Ac_switch_guard(THD *, const char *);
+  ~Ac_switch_guard();
+  void commit() { committed = true; }
+};
 
 extern AC *db_ac;
 

sql/sql_parse.cc

@@ -1494,58 +1494,29 @@ static int process_noncurrent_db_rw (THD *thd, TABLE_LIST *all_tables)
  */
 static bool set_session_db_helper(THD *thd, const LEX_STRING *new_db)
 {
-  // initialize empty attributes. only database may be set later
-  MT_RESOURCE_ATTRS attrs = {nullptr, nullptr, nullptr};
-  // save existing session db
-  std::string old_db;
+  bool error = false;
+  Ac_switch_guard switch_guard(thd, new_db->str);
 
   if (!thd->db || strcmp(thd->db, new_db->str))
   {
-    if (thd->db)
-    {
-      old_db = thd->db;
-    }
-    // since we are changing the session db and the query is part of
-    // multi-statement packet, we need to reset admission control so the new
-    // queries won't bypass the new db's admission control.
-    if (thd->is_in_ac)
-    {
-      // current admission control is per database
-      DBUG_ASSERT(thd->db);
-      // if the thd is already in admission control for the previous
-      // database, we need to release it before resetting the value.
-      attrs.database = thd->db;
-      multi_tenancy_exit_query(thd);
-    }
     // now check resource if we can switch the connection to another database.
     // Since other attributes are null, this doesn't have any effect if the
     // throttling is not db based.
-    attrs.database = new_db->str; // set the new db in attributes
+    MT_RESOURCE_ATTRS attrs = { nullptr, nullptr, new_db->str };
     if (attrs.database && multi_tenancy_add_connection(thd, &attrs))
     {
       my_error(ER_MULTI_TENANCY_MAX_CONNECTION, MYF(0), attrs.database);
-      return true;
+      error = true;
     }
   }
 
-  if (!mysql_change_db(thd, new_db, FALSE))
-  {
-    // release old connection resource in multi-tenancy plugin
-    if (!old_db.empty())
-    {
-      attrs.database = old_db.c_str();
-      multi_tenancy_close_connection(thd, &attrs);
-    }
-    return false;
-  }
-  else if (attrs.database)
-  {
-    // failed to change to the new database. release the connection resource
-    multi_tenancy_close_connection(thd, &attrs);
-  }
+  if (!error && !mysql_change_db(thd, new_db, false)) // Do not force switch.
+    switch_guard.commit();
+  else
+    error = true;
 
-  // mysql_change_db failed. error is set in side mysql_change_db
-  return true;
+  // Switch guard either commits or rolls back the switch operation.
+  return error;
 }
 
 #ifndef EMBEDDED_LIBRARY