Track connection refcount to prevent GC races

In certain cases a Mysql2::Client and a Mysql2::Result may be freed
in the same GC run. If the Client is freed before the Result a segv
may occur when the Result tries to use the socket to make sure there
isn't any more data to read before freeing itself.

Huge thanks to @evanphx for helping me track this one down and @rkh
for helping with TravisCI since that's where the issue was finally
reproducable.

Author: brianmario

ext/mysql2/client.c

@@ -185,10 +185,12 @@ static VALUE nogvl_close(void *ptr) {
 static void rb_mysql_client_free(void * ptr) {
   mysql_client_wrapper *wrapper = (mysql_client_wrapper *)ptr;
 
-  nogvl_close(wrapper);
+  if (wrapper->refcount == 0) {
+    nogvl_close(wrapper);
 
-  xfree(wrapper->client);
-  xfree(ptr);
+    xfree(wrapper->client);
+    xfree(ptr);
+  }
 }
 
 static VALUE allocate(VALUE klass) {
@@ -200,6 +202,7 @@ static VALUE allocate(VALUE klass) {
   wrapper->reconnect_enabled = 0;
   wrapper->connected = 0; /* means that a database connection is open */
   wrapper->initialized = 0; /* means that that the wrapper is initialized */
+  wrapper->refcount = 0;
   wrapper->client = (MYSQL*)xmalloc(sizeof(MYSQL));
   return obj;
 }
@@ -393,7 +396,7 @@ static VALUE rb_mysql_client_async_result(VALUE self) {
     return Qnil;
   }
 
-  resultObj = rb_mysql_result_to_obj(result);
+  resultObj = rb_mysql_result_to_obj(wrapper, result);
   /* pass-through query options for result construction later */
   rb_iv_set(resultObj, "@query_options", rb_hash_dup(rb_iv_get(self, "@current_query_options")));
 
@@ -942,7 +945,7 @@ static VALUE rb_mysql_client_store_result(VALUE self)
     return Qnil;
   }
 
-  resultObj = rb_mysql_result_to_obj(result);
+  resultObj = rb_mysql_result_to_obj(wrapper, result);
   /* pass-through query options for result construction later */
   rb_iv_set(resultObj, "@query_options", rb_hash_dup(rb_iv_get(self, "@current_query_options")));
 

ext/mysql2/client.h

@@ -38,6 +38,7 @@ typedef struct {
   int active;
   int connected;
   int initialized;
+  int refcount;
   MYSQL *client;
 } mysql_client_wrapper;
 

ext/mysql2/result.c

@@ -72,6 +72,7 @@ static void rb_mysql_result_free_result(mysql2_result_wrapper * wrapper) {
   if (wrapper && wrapper->resultFreed != 1) {
     mysql_free_result(wrapper->result);
     wrapper->resultFreed = 1;
+    wrapper->client_wrapper->refcount--;
   }
 }
 
@@ -562,7 +563,7 @@ static VALUE rb_mysql_result_count(VALUE self) {
 }
 
 /* Mysql2::Result */
-VALUE rb_mysql_result_to_obj(MYSQL_RES * r) {
+VALUE rb_mysql_result_to_obj(mysql_client_wrapper *client_wrapper, MYSQL_RES *r) {
   VALUE obj;
   mysql2_result_wrapper * wrapper;
   obj = Data_Make_Struct(cMysql2Result, mysql2_result_wrapper, rb_mysql_result_mark, rb_mysql_result_free, wrapper);
@@ -575,6 +576,8 @@ VALUE rb_mysql_result_to_obj(MYSQL_RES * r) {
   wrapper->rows = Qnil;
   wrapper->encoding = Qnil;
   wrapper->streamingComplete = 0;
+  wrapper->client_wrapper = client_wrapper;
+  wrapper->client_wrapper->refcount++;
   rb_obj_call_init(obj, 0, NULL);
   return obj;
 }

ext/mysql2/result.h

@@ -2,7 +2,7 @@
 #define MYSQL2_RESULT_H
 
 void init_mysql2_result();
-VALUE rb_mysql_result_to_obj(MYSQL_RES * r);
+VALUE rb_mysql_result_to_obj(mysql_client_wrapper *client_wrapper, MYSQL_RES * r);
 
 typedef struct {
   VALUE fields;
@@ -14,6 +14,7 @@ typedef struct {
   char streamingComplete;
   char resultFreed;
   MYSQL_RES *result;
+  mysql_client_wrapper *client_wrapper;
 } mysql2_result_wrapper;
 
 #define GetMysql2Result(obj, sval) (sval = (mysql2_result_wrapper*)DATA_PTR(obj));