Track the Statement refcount before freeing it

Author: sodabrew

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

@@ -92,8 +92,8 @@ static void rb_mysql_result_free_result(mysql2_result_wrapper * wrapper) {
   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);
 
       if (wrapper->result_buffers) {
         unsigned int i;
@@ -125,6 +125,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);
 }
 
@@ -341,23 +345,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
@@ -873,11 +877,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");
   }
 
@@ -903,7 +907,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;
@@ -921,7 +925,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;
@@ -943,8 +947,8 @@ 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));
     }
@@ -977,9 +981,10 @@ VALUE rb_mysql_result_to_obj(VALUE client, VALUE encoding, VALUE options, MYSQL_
   /* Keep a handle to the Statement to ensure it doesn't get garbage collected first */
   wrapper->statement = statement;
   if (statement != Qnil) {
-    mysql_stmt_wrapper *stmt_wrapper = DATA_PTR(statement);
-    wrapper->stmt = stmt_wrapper->stmt;
-    stmt_wrapper->refcount++;
+    wrapper->stmt_wrapper = DATA_PTR(statement);
+    wrapper->stmt_wrapper->refcount++;
+  } else {
+    wrapper->stmt_wrapper = NULL;
   }
 
   rb_obj_call_init(obj, 0, NULL);

ext/mysql2/result.h

@@ -17,7 +17,7 @@ typedef struct {
   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

@@ -19,10 +19,16 @@ static void rb_mysql_stmt_mark(void * ptr) {
 
 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 = 0;
     stmt_wrapper->stmt = NULL;
   }
 

ext/mysql2/statement.h

@@ -3,14 +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