update MySQLImpl to support character_set_results (#77)

* update MySQLImpl to support character_set_results

* remove WFComplexClientTask::is_user_request().

Co-authored-by: XieHan <XieHan@sogou-inc.com>

Author: Barenboim

src/factory/MySQLTaskImpl.cc

@@ -14,6 +14,7 @@
   limitations under the License.
 
   Authors: Wu Jiaxu (wujiaxu@sogou-inc.com)
+           Xie Han (xiehan@sogou-inc.com)
            Li Yingxin (liyingxin@sogou-inc.com)
 */
 
@@ -54,14 +55,15 @@ class ComplexMySQLTask : public WFComplexClientTask<MySQLRequest, MySQLResponse>
 	std::string username_;
 	std::string password_;
 	std::string db_;
+	std::string res_charset_;
 	int character_set_;
-	bool succ_;
 #define NO_TRANSACTION          -1
 #define TRANSACTION_OUT         0
 #define TRANSACTION_IN          1
 #define TRANSACTION_CONN_RESET  -2
 	int transaction_state_;
 #define PREPARE_IN              2
+	bool succ_;
 	bool is_user_request_;
 
 public:
@@ -125,35 +127,44 @@ bool ComplexMySQLTask::check_request()
 CommMessageOut *ComplexMySQLTask::message_out()
 {
 	long long seqid = this->get_seq();
+	MySQLRequest *req;
 
 	if (seqid == 0)
-	{
-		succ_ = false;
-		is_user_request_ = false;
-		return new MySQLHandshakeRequest;
-	}
+		req = new MySQLHandshakeRequest;
 	else if (seqid == 1)
 	{
-		succ_ = false;
-		is_user_request_ = false;
 		auto *auth_req = new MySQLAuthRequest;
 		auto *conn = this->get_connection();
 		auto *ctx = static_cast<handshake_ctx *>(conn->get_context());
 
 		auth_req->set_seqid(ctx->mysql_seqid);
 		auth_req->set_challenge(ctx->challenge);
 		delete ctx;
-		conn->set_context(nullptr, nullptr);
+		conn->set_context(NULL, nullptr);
 		auth_req->set_auth(username_, password_, db_, character_set_);
-		return auth_req;
+		req = auth_req;
+	}
+	else if (seqid == 2 && res_charset_.size() != 0)
+	{
+		req = new MySQLRequest;
+		req->set_query("SET NAMES " + res_charset_);
+	}
+	else
+		req = NULL;
+
+	if (req)
+	{
+		succ_ = false;
+		is_user_request_ = false;
+		return req;
 	}
 
 	if (is_transaction())
 	{
 		auto *target = static_cast<RouteManager::RouteTarget *>(this->get_target());
 
-		if (seqid == 2 && ((target->state & TRANSACTION_IN) ||
-						   (target->state & PREPARE_IN))) // CONN RESET
+		if (seqid <= 3 && (seqid == 2 || res_charset_.size() != 0) &&
+			(target->state & (TRANSACTION_IN | PREPARE_IN)))
 		{
 			target->state = TRANSACTION_OUT;
 			transaction_state_ = TRANSACTION_CONN_RESET;
@@ -210,6 +221,8 @@ CommMessageIn *ComplexMySQLTask::message_in()
 		return new MySQLHandshakeResponse;
 	else if (seqid == 1)
 		return new MySQLAuthResponse;
+	else if (seqid == 2 && !is_user_request_)
+		return new MySQLResponse;
 
 	return this->WFClientTask::message_in();
 }
@@ -242,24 +255,21 @@ int ComplexMySQLTask::keep_alive_timeout()
 			succ_ = true;
 		}
 	}
-	else if (seqid == 1)
+	else if (!is_user_request_)
 	{
-		auto *resp = static_cast<MySQLAuthResponse *>(this->get_message_in());
+		auto *resp = static_cast<MySQLResponse *>(this->get_message_in());
 
 		succ_ = resp->is_ok_packet();
-		if (succ_)
-		{
-			return is_transaction() ? MYSQL_KEEPALIVE_TRANSACTION
-									: MYSQL_KEEPALIVE_DEFAULT;
-		}
-		else
+		if (!succ_)
 		{
-			this->resp = std::move(*static_cast<MySQLResponse *>(resp));
+			this->resp = std::move(*resp);
 			return 0;
 		}
 	}
+	else
+		return this->keep_alive_timeo;
 
-	return this->keep_alive_timeo;
+	return MYSQL_KEEPALIVE_DEFAULT;
 }
 
 /*
@@ -427,9 +437,11 @@ bool ComplexMySQLTask::init_success()
 	{
 		auto query_kv = URIParser::split_query(uri_.query);
 
-		for (const auto& kv : query_kv)
+		for (auto& kv : query_kv)
 		{
-			if (strcasecmp(kv.first.c_str(), "character_set") == 0)
+			if (strcasecmp(kv.first.c_str(), "transaction") == 0)
+				transaction = std::move(kv.second);
+			else if (strcasecmp(kv.first.c_str(), "character_set") == 0)
 			{
 				character_set_ = __mysql_get_character_set(kv.second);
 				if (character_set_ < 0)
@@ -439,16 +451,19 @@ bool ComplexMySQLTask::init_success()
 					return false;
 				}
 			}
-			else if (strcasecmp(kv.first.c_str(), "transaction") == 0)
-				transaction = kv.second;
+			else if (strcasecmp(kv.first.c_str(), "character_set_results") == 0)
+				res_charset_ = std::move(kv.second);
 		}
 	}
 
-	size_t info_len = username_.size() + password_.size() + db_.size() + 40;
+	size_t info_len = username_.size() + password_.size() + db_.size() +
+					  res_charset_.size() + 50;
 	char *info = new char[info_len];
 
-	snprintf(info, info_len, "mysql|user:%s|pass:%s|db:%s|charset:%d",
-			 username_.c_str(), password_.c_str(), db_.c_str(), character_set_);
+	snprintf(info, info_len, "mysql|user:%s|pass:%s|db:%s|"
+							 "charset:%d|rcharset:%s",
+			 username_.c_str(), password_.c_str(), db_.c_str(),
+			 character_set_, res_charset_.c_str());
 	this->WFComplexClientTask::set_type(type);
 
 	if (!transaction.empty())
@@ -478,10 +493,17 @@ bool ComplexMySQLTask::finish_once()
 
 		if (this->state == WFT_STATE_SUCCESS && !succ_)
 		{
+			long long seqid = this->get_seq();
+
+			if (seqid == 0)
+				this->error = WFT_ERR_MYSQL_HOST_NOT_ALLOWED;
+			else if (seqid == 1)
+				this->error = WFT_ERR_MYSQL_ACCESS_DENIED;
+			else
+				this->error = WFT_ERR_MYSQL_INVALID_CHARACTER_SET;
+
 			this->disable_retry();
 			this->state = WFT_STATE_TASK_ERROR;
-			this->error = this->get_seq() == 0 ? WFT_ERR_MYSQL_HOST_NOT_ALLOWED
-											   : WFT_ERR_MYSQL_ACCESS_DENIED;
 		}
 
 		return false;

src/factory/RedisTaskImpl.cc

@@ -118,7 +118,7 @@ CommMessageOut *ComplexRedisTask::message_out()
 
 int ComplexRedisTask::keep_alive_timeout()
 {
-	if (this->is_user_request())
+	if (this->is_user_request_)
 		return this->keep_alive_timeo;
 
 	RedisResponse *resp = this->get_resp();

src/factory/WFTaskFactory.inl

@@ -317,11 +317,6 @@ protected:
 		this->resp.set_size_limit(size);
 	}
 
-	bool is_user_request() const
-	{
-		return this->get_message_out() == &this->req;
-	}
-
 	void disable_retry()
 	{
 		retry_times_ = retry_max_;

src/protocol/MySQLMessage.inl

@@ -24,7 +24,7 @@
 namespace protocol
 {
 
-using MySQLHandshakeRequest = MySQLMessage;
+using MySQLHandshakeRequest = MySQLRequest;
 
 class MySQLHandshakeResponse : public MySQLResponse
 {
@@ -51,7 +51,7 @@ public:
 		memcpy(auth_plugin_data_part_2_, auth2, 12);
 	}
 
-	bool host_disallowed() { return this->disallowed_; }
+	bool host_disallowed() const { return this->disallowed_; }
 
 private:
 	virtual int decode_packet(const char *buf, size_t buflen);
@@ -73,7 +73,7 @@ public:
 	MySQLHandshakeResponse& operator= (MySQLHandshakeResponse&& move) = default;
 };
 
-class MySQLAuthRequest : public MySQLMessage
+class MySQLAuthRequest : public MySQLRequest
 {
 public:
 	void set_auth(const std::string& username,