Merge pull request #636 from sodabrew/stmt_gc_fix

Statement GC fixes

Author: sodabrew

ext/mysql2/client.c

@@ -490,7 +490,7 @@ static VALUE rb_mysql_client_async_result(VALUE self) {
   current = rb_hash_dup(rb_iv_get(self, "@current_query_options"));
   (void)RB_GC_GUARD(current);
   Check_Type(current, T_HASH);
-  resultObj = rb_mysql_result_to_obj(self, wrapper->encoding, current, result, NULL);
+  resultObj = rb_mysql_result_to_obj(self, wrapper->encoding, current, result, Qnil);
 
   return resultObj;
 }
@@ -1050,7 +1050,7 @@ static VALUE rb_mysql_client_store_result(VALUE self)
   current = rb_hash_dup(rb_iv_get(self, "@current_query_options"));
   (void)RB_GC_GUARD(current);
   Check_Type(current, T_HASH);
-  resultObj = rb_mysql_result_to_obj(self, wrapper->encoding, current, result, NULL);
+  resultObj = rb_mysql_result_to_obj(self, wrapper->encoding, current, result, Qnil);
 
   return resultObj;
 }

ext/mysql2/mysql2_ext.h

@@ -39,8 +39,8 @@ typedef unsigned int    uint;
 #endif
 
 #include <client.h>
-#include <result.h>
 #include <statement.h>
+#include <result.h>
 #include <infile.h>
 
 #endif

ext/mysql2/result.c

@@ -75,26 +75,37 @@ static VALUE sym_symbolize_keys, sym_as, sym_array, sym_database_timezone, sym_a
           sym_local, sym_utc, sym_cast_booleans, sym_cache_rows, sym_cast, sym_stream, sym_name;
 static ID intern_merge;
 
+/* Mark any VALUEs that are only referenced in C, so the GC won't get them. */
 static void rb_mysql_result_mark(void * wrapper) {
   mysql2_result_wrapper * w = wrapper;
   if (w) {
     rb_gc_mark(w->fields);
     rb_gc_mark(w->rows);
     rb_gc_mark(w->encoding);
     rb_gc_mark(w->client);
+    rb_gc_mark(w->statement);
   }
 }
 
 /* this may be called manually or during GC */
 static void rb_mysql_result_free_result(mysql2_result_wrapper * wrapper) {
-  unsigned int i;
   if (!wrapper) return;
 
   if (wrapper->resultFreed != 1) {
-    if (wrapper->stmt) {
-      mysql_stmt_free_result(wrapper->stmt);
+    if (wrapper->stmt_wrapper) {
+      mysql_stmt_free_result(wrapper->stmt_wrapper->stmt);
+
+      /* MySQL BUG? If the statement handle was previously used, and so
+       * mysql_stmt_bind_result was called, and if that result set and bind buffers were freed,
+       * MySQL still thinks the result set buffer is available and will prefetch the
+       * first result in mysql_stmt_execute. This will corrupt or crash the program.
+       * By setting bind_result_done back to 0, we make MySQL think that a result set
+       * has never been bound to this statement handle before to prevent the prefetch.
+       */
+      wrapper->stmt_wrapper->stmt->bind_result_done = 0;
 
       if (wrapper->result_buffers) {
+        unsigned int i;
         for (i = 0; i < wrapper->numberOfFields; i++) {
           if (wrapper->result_buffers[i].buffer) {
             xfree(wrapper->result_buffers[i].buffer);
@@ -105,8 +116,11 @@ static void rb_mysql_result_free_result(mysql2_result_wrapper * wrapper) {
         xfree(wrapper->error);
         xfree(wrapper->length);
       }
+      /* Clue that the next statement execute will need to allocate a new result buffer. */
+      wrapper->result_buffers = NULL;
     }
     /* FIXME: this may call flush_use_result, which can hit the socket */
+    /* For prepared statements, wrapper->result is the result metadata */
     mysql_free_result(wrapper->result);
     wrapper->resultFreed = 1;
   }
@@ -122,6 +136,10 @@ static void rb_mysql_result_free(void *ptr) {
     decr_mysql2_client(wrapper->client_wrapper);
   }
 
+  if (wrapper->statement != Qnil) {
+    decr_mysql2_stmt(wrapper->stmt_wrapper);
+  }
+
   xfree(wrapper);
 }
 
@@ -338,23 +356,23 @@ static VALUE rb_mysql_result_fetch_row_stmt(VALUE self, MYSQL_FIELD * fields, co
     rb_mysql_result_alloc_result_buffers(self, fields);
   }
 
-  if(mysql_stmt_bind_result(wrapper->stmt, wrapper->result_buffers)) {
-    rb_raise_mysql2_stmt_error2(wrapper->stmt
+  if (mysql_stmt_bind_result(wrapper->stmt_wrapper->stmt, wrapper->result_buffers)) {
+    rb_raise_mysql2_stmt_error2(wrapper->stmt_wrapper->stmt
 #ifdef HAVE_RUBY_ENCODING_H
       , conn_enc
 #endif
       );
   }
 
   {
-    switch((uintptr_t)rb_thread_call_without_gvl(nogvl_stmt_fetch, wrapper->stmt, RUBY_UBF_IO, 0)) {
+    switch((uintptr_t)rb_thread_call_without_gvl(nogvl_stmt_fetch, wrapper->stmt_wrapper->stmt, RUBY_UBF_IO, 0)) {
       case 0:
         /* success */
         break;
 
       case 1:
         /* error */
-        rb_raise_mysql2_stmt_error2(wrapper->stmt
+        rb_raise_mysql2_stmt_error2(wrapper->stmt_wrapper->stmt
 #ifdef HAVE_RUBY_ENCODING_H
           , conn_enc
 #endif
@@ -870,11 +888,11 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
     rb_warn(":cache_rows is ignored if :stream is true");
   }
 
-  if (wrapper->stmt && !cacheRows && !wrapper->is_streaming) {
+  if (wrapper->stmt_wrapper && !cacheRows && !wrapper->is_streaming) {
     rb_warn(":cache_rows is forced for prepared statements (if not streaming)");
   }
 
-  if (wrapper->stmt && !cast) {
+  if (wrapper->stmt_wrapper && !cast) {
     rb_warn(":cast is forced for prepared statements");
   }
 
@@ -900,7 +918,7 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
   }
 
   if (wrapper->lastRowProcessed == 0 && !wrapper->is_streaming) {
-    wrapper->numberOfRows = wrapper->stmt ? mysql_stmt_num_rows(wrapper->stmt) : mysql_num_rows(wrapper->result);
+    wrapper->numberOfRows = wrapper->stmt_wrapper ? mysql_stmt_num_rows(wrapper->stmt_wrapper->stmt) : mysql_num_rows(wrapper->result);
     if (wrapper->numberOfRows == 0) {
       wrapper->rows = rb_ary_new();
       return wrapper->rows;
@@ -918,7 +936,7 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
   args.app_timezone = app_timezone;
   args.block_given = block;
 
-  if (wrapper->stmt) {
+  if (wrapper->stmt_wrapper) {
     fetch_row_func = rb_mysql_result_fetch_row_stmt;
   } else {
     fetch_row_func = rb_mysql_result_fetch_row;
@@ -940,16 +958,16 @@ static VALUE rb_mysql_result_count(VALUE self) {
     return LONG2NUM(RARRAY_LEN(wrapper->rows));
   } else {
     /* MySQL returns an unsigned 64-bit long here */
-    if(wrapper->stmt) {
-      return ULL2NUM(mysql_stmt_num_rows(wrapper->stmt));
+    if (wrapper->stmt_wrapper) {
+      return ULL2NUM(mysql_stmt_num_rows(wrapper->stmt_wrapper->stmt));
     } else {
       return ULL2NUM(mysql_num_rows(wrapper->result));
     }
   }
 }
 
 /* Mysql2::Result */
-VALUE rb_mysql_result_to_obj(VALUE client, VALUE encoding, VALUE options, MYSQL_RES *r, MYSQL_STMT * s) {
+VALUE rb_mysql_result_to_obj(VALUE client, VALUE encoding, VALUE options, MYSQL_RES *r, VALUE statement) {
   VALUE obj;
   mysql2_result_wrapper * wrapper;
 
@@ -966,12 +984,20 @@ VALUE rb_mysql_result_to_obj(VALUE client, VALUE encoding, VALUE options, MYSQL_
   wrapper->client = client;
   wrapper->client_wrapper = DATA_PTR(client);
   wrapper->client_wrapper->refcount++;
-  wrapper->stmt = s;
   wrapper->result_buffers = NULL;
   wrapper->is_null = NULL;
   wrapper->error = NULL;
   wrapper->length = NULL;
 
+  /* Keep a handle to the Statement to ensure it doesn't get garbage collected first */
+  wrapper->statement = statement;
+  if (statement != Qnil) {
+    wrapper->stmt_wrapper = DATA_PTR(statement);
+    wrapper->stmt_wrapper->refcount++;
+  } else {
+    wrapper->stmt_wrapper = NULL;
+  }
+
   rb_obj_call_init(obj, 0, NULL);
   rb_iv_set(obj, "@query_options", options);
 

ext/mysql2/result.h

@@ -2,21 +2,22 @@
 #define MYSQL2_RESULT_H
 
 void init_mysql2_result();
-VALUE rb_mysql_result_to_obj(VALUE client, VALUE encoding, VALUE options, MYSQL_RES *r, MYSQL_STMT * s);
+VALUE rb_mysql_result_to_obj(VALUE client, VALUE encoding, VALUE options, MYSQL_RES *r, VALUE statement);
 
 typedef struct {
   VALUE fields;
   VALUE rows;
   VALUE client;
   VALUE encoding;
+  VALUE statement;
   unsigned int numberOfFields;
   unsigned long numberOfRows;
   unsigned long lastRowProcessed;
   char is_streaming;
   char streamingComplete;
   char resultFreed;
   MYSQL_RES *result;
-  MYSQL_STMT *stmt;
+  mysql_stmt_wrapper *stmt_wrapper;
   mysql_client_wrapper *client_wrapper;
   /* statement result bind buffers */
   MYSQL_BIND *result_buffers;

ext/mysql2/statement.c

@@ -12,17 +12,23 @@ static VALUE intern_usec, intern_sec, intern_min, intern_hour, intern_day, inter
 
 static void rb_mysql_stmt_mark(void * ptr) {
   mysql_stmt_wrapper* stmt_wrapper = (mysql_stmt_wrapper *)ptr;
-  if(! stmt_wrapper) return;
+  if (!stmt_wrapper) return;
 
   rb_gc_mark(stmt_wrapper->client);
 }
 
 static void rb_mysql_stmt_free(void * ptr) {
   mysql_stmt_wrapper* stmt_wrapper = (mysql_stmt_wrapper *)ptr;
+  decr_mysql2_stmt(stmt_wrapper);
+}
 
-  mysql_stmt_close(stmt_wrapper->stmt);
+void decr_mysql2_stmt(mysql_stmt_wrapper *stmt_wrapper) {
+  stmt_wrapper->refcount--;
 
-  xfree(ptr);
+  if (stmt_wrapper->refcount == 0) {
+    mysql_stmt_close(stmt_wrapper->stmt);
+    xfree(stmt_wrapper);
+  }
 }
 
 VALUE rb_raise_mysql2_stmt_error2(MYSQL_STMT *stmt
@@ -103,6 +109,7 @@ VALUE rb_mysql_stmt_new(VALUE rb_client, VALUE sql) {
   rb_stmt = Data_Make_Struct(cMysql2Statement, mysql_stmt_wrapper, rb_mysql_stmt_mark, rb_mysql_stmt_free, stmt_wrapper);
   {
     stmt_wrapper->client = rb_client;
+    stmt_wrapper->refcount = 1;
     stmt_wrapper->stmt = NULL;
   }
 
@@ -169,7 +176,7 @@ static VALUE field_count(VALUE self) {
 static void *nogvl_execute(void *ptr) {
   MYSQL_STMT *stmt = ptr;
 
-  if(mysql_stmt_execute(stmt)) {
+  if (mysql_stmt_execute(stmt)) {
     return (void*)Qfalse;
   } else {
     return (void*)Qtrue;
@@ -179,7 +186,7 @@ static void *nogvl_execute(void *ptr) {
 static void *nogvl_stmt_store_result(void *ptr) {
   MYSQL_STMT *stmt = ptr;
 
-  if(mysql_stmt_store_result(stmt)) {
+  if (mysql_stmt_store_result(stmt)) {
     return (void *)Qfalse;
   } else {
     return (void *)Qtrue;
@@ -340,8 +347,8 @@ static VALUE execute(int argc, VALUE *argv, VALUE self) {
   FREE_BINDS;
 
   metadata = mysql_stmt_result_metadata(stmt);
-  if(metadata == NULL) {
-    if(mysql_stmt_errno(stmt) != 0) {
+  if (metadata == NULL) {
+    if (mysql_stmt_errno(stmt) != 0) {
       // either CR_OUT_OF_MEMORY or CR_UNKNOWN_ERROR. both fatal.
 
       MARK_CONN_INACTIVE(stmt_wrapper->client);
@@ -359,12 +366,13 @@ static VALUE execute(int argc, VALUE *argv, VALUE self) {
   if (!is_streaming) {
     // recieve the whole result set from the server
     if (rb_thread_call_without_gvl(nogvl_stmt_store_result, stmt, RUBY_UBF_IO, 0) == Qfalse) {
+      mysql_free_result(metadata);
       rb_raise_mysql2_stmt_error(self);
     }
     MARK_CONN_INACTIVE(stmt_wrapper->client);
   }
 
-  resultObj = rb_mysql_result_to_obj(stmt_wrapper->client, wrapper->encoding, current, metadata, stmt);
+  resultObj = rb_mysql_result_to_obj(stmt_wrapper->client, wrapper->encoding, current, metadata, self);
 
   if (!is_streaming) {
     // cache all result

ext/mysql2/statement.h

@@ -3,13 +3,15 @@
 
 extern VALUE cMysql2Statement;
 
-void init_mysql2_statement();
-
 typedef struct {
   VALUE client;
   MYSQL_STMT *stmt;
+  int refcount;
 } mysql_stmt_wrapper;
 
+void init_mysql2_statement();
+void decr_mysql2_stmt(mysql_stmt_wrapper *stmt_wrapper);
+
 VALUE rb_mysql_stmt_new(VALUE rb_client, VALUE sql);
 VALUE rb_raise_mysql2_stmt_error2(MYSQL_STMT *stmt
 #ifdef HAVE_RUBY_ENCODING_H

spec/mysql2/statement_spec.rb

@@ -70,6 +70,22 @@
     @client.query 'DROP TABLE IF EXISTS mysql2_stmt_q'
   end
 
+  it "should be reusable 1000 times" do
+    statement = @client.prepare 'SELECT 1'
+    1000.times do
+      result = statement.execute
+      expect(result.to_a.length).to eq(1)
+    end
+  end
+
+  it "should be reusable 10000 times" do
+    statement = @client.prepare 'SELECT 1'
+    10000.times do
+      result = statement.execute
+      expect(result.to_a.length).to eq(1)
+    end
+  end
+
   it "should select dates" do
     statement = @client.prepare 'SELECT NOW()'
     result = statement.execute