Merge pull request #378 from brianmario/client-refcount

Track connection refcount to prevent GC races

Author: brianmario

ext/mysql2/client.c

@@ -193,13 +193,15 @@ static VALUE nogvl_close(void *ptr) {
   return Qnil;
 }
 
-static void rb_mysql_client_free(void * ptr) {
+static void rb_mysql_client_free(void *ptr) {
   mysql_client_wrapper *wrapper = (mysql_client_wrapper *)ptr;
 
-  nogvl_close(wrapper);
-
-  xfree(wrapper->client);
-  xfree(ptr);
+  wrapper->refcount--;
+  if (wrapper->refcount == 0) {
+    nogvl_close(wrapper);
+    xfree(wrapper->client);
+    xfree(wrapper);
+  }
 }
 
 static VALUE allocate(VALUE klass) {
@@ -211,6 +213,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 = 1;
   wrapper->client = (MYSQL*)xmalloc(sizeof(MYSQL));
   return obj;
 }
@@ -372,9 +375,6 @@ static VALUE nogvl_use_result(void *ptr) {
 static VALUE rb_mysql_client_async_result(VALUE self) {
   MYSQL_RES * result;
   VALUE resultObj;
-#ifdef HAVE_RUBY_ENCODING_H
-  mysql2_result_wrapper * result_wrapper;
-#endif
   GET_CLIENT(self);
 
   /* if we're not waiting on a result, do nothing */
@@ -404,14 +404,7 @@ static VALUE rb_mysql_client_async_result(VALUE self) {
     return Qnil;
   }
 
-  resultObj = rb_mysql_result_to_obj(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")));
-
-#ifdef HAVE_RUBY_ENCODING_H
-  GetMysql2Result(resultObj, result_wrapper);
-  result_wrapper->encoding = wrapper->encoding;
-#endif
+  resultObj = rb_mysql_result_to_obj(self, wrapper->encoding, rb_hash_dup(rb_iv_get(self, "@current_query_options")), result);
   return resultObj;
 }
 
@@ -937,10 +930,6 @@ static VALUE rb_mysql_client_store_result(VALUE self)
 {
   MYSQL_RES * result;
   VALUE resultObj;
-#ifdef HAVE_RUBY_ENCODING_H
-  mysql2_result_wrapper * result_wrapper;
-#endif
-
   GET_CLIENT(self);
 
   result = (MYSQL_RES *)rb_thread_blocking_region(nogvl_store_result, wrapper, RUBY_UBF_IO, 0);
@@ -953,14 +942,7 @@ static VALUE rb_mysql_client_store_result(VALUE self)
     return Qnil;
   }
 
-  resultObj = rb_mysql_result_to_obj(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")));
-
-#ifdef HAVE_RUBY_ENCODING_H
-  GetMysql2Result(resultObj, result_wrapper);
-  result_wrapper->encoding = wrapper->encoding;
-#endif
+  resultObj = rb_mysql_result_to_obj(self, wrapper->encoding, rb_hash_dup(rb_iv_get(self, "@current_query_options")), result);
   return resultObj;
 
 }

ext/mysql2/client.h

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

ext/mysql2/result.c

@@ -64,22 +64,33 @@ static void rb_mysql_result_mark(void * wrapper) {
     rb_gc_mark(w->fields);
     rb_gc_mark(w->rows);
     rb_gc_mark(w->encoding);
+    rb_gc_mark(w->client);
   }
 }
 
 /* this may be called manually or during GC */
 static void rb_mysql_result_free_result(mysql2_result_wrapper * wrapper) {
   if (wrapper && wrapper->resultFreed != 1) {
+    /* FIXME: this may call flush_use_result, which can hit the socket */
     mysql_free_result(wrapper->result);
     wrapper->resultFreed = 1;
   }
 }
 
 /* this is called during GC */
-static void rb_mysql_result_free(void * wrapper) {
-  mysql2_result_wrapper * w = wrapper;
-  /* FIXME: this may call flush_use_result, which can hit the socket */
-  rb_mysql_result_free_result(w);
+static void rb_mysql_result_free(void *ptr) {
+  mysql2_result_wrapper * wrapper = ptr;
+  rb_mysql_result_free_result(wrapper);
+
+  // If the GC gets to client first it will be nil
+  if (wrapper->client != Qnil) {
+    wrapper->client_wrapper->refcount--;
+    if (wrapper->client_wrapper->refcount == 0) {
+      xfree(wrapper->client_wrapper->client);
+      xfree(wrapper->client_wrapper);
+    }
+  }
+
   xfree(wrapper);
 }
 
@@ -565,7 +576,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(VALUE client, VALUE encoding, VALUE options, 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);
@@ -576,9 +587,16 @@ VALUE rb_mysql_result_to_obj(MYSQL_RES * r) {
   wrapper->result = r;
   wrapper->fields = Qnil;
   wrapper->rows = Qnil;
-  wrapper->encoding = Qnil;
+  wrapper->encoding = encoding;
   wrapper->streamingComplete = 0;
+  wrapper->client = client;
+  wrapper->client_wrapper = DATA_PTR(client);
+  wrapper->client_wrapper->refcount++;
+
   rb_obj_call_init(obj, 0, NULL);
+
+  rb_iv_set(obj, "@query_options", options);
+
   return obj;
 }
 

ext/mysql2/result.h

@@ -2,18 +2,20 @@
 #define MYSQL2_RESULT_H
 
 void init_mysql2_result();
-VALUE rb_mysql_result_to_obj(MYSQL_RES * r);
+VALUE rb_mysql_result_to_obj(VALUE client, VALUE encoding, VALUE options, MYSQL_RES *r);
 
 typedef struct {
   VALUE fields;
   VALUE rows;
+  VALUE client;
   VALUE encoding;
   unsigned int numberOfFields;
   unsigned long numberOfRows;
   unsigned long lastRowProcessed;
   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));

spec/mysql2/client_spec.rb

@@ -257,10 +257,10 @@ def connect *args
           mark[:END] = Time.now
           mark.include?(:USR1).should be_true
           (mark[:USR1] - mark[:START]).should >= 1
-          (mark[:USR1] - mark[:START]).should < 1.1
+          (mark[:USR1] - mark[:START]).should < 1.2
           (mark[:END] - mark[:USR1]).should > 0.9
           (mark[:END] - mark[:START]).should >= 2
-          (mark[:END] - mark[:START]).should < 2.1
+          (mark[:END] - mark[:START]).should < 2.2
           Process.kill(:TERM, pid)
           Process.waitpid2(pid)
         ensure