Allow the garbage collector to close connections

Author: cbandy

ext/mysql2/client.c

@@ -233,7 +233,7 @@ void decr_mysql2_client(mysql_client_wrapper *wrapper)
 
   if (wrapper->refcount == 0) {
 #ifndef _WIN32
-    if (wrapper->connected) {
+    if (wrapper->connected && !wrapper->automatic_close) {
       /* The client is being garbage collected while connected. Prevent
        * mysql_close() from sending a mysql-QUIT or from calling shutdown() on
        * the socket by invalidating it. invalidate_fd() will drop this
@@ -260,6 +260,11 @@ static VALUE allocate(VALUE klass) {
   obj = Data_Make_Struct(klass, mysql_client_wrapper, rb_mysql_client_mark, rb_mysql_client_free, wrapper);
   wrapper->encoding = Qnil;
   MARK_CONN_INACTIVE(self);
+#ifndef _WIN32
+  wrapper->automatic_close = 0;
+#else
+  wrapper->automatic_close = 1;
+#endif
   wrapper->server_version = 0;
   wrapper->reconnect_enabled = 0;
   wrapper->connect_timeout = 0;
@@ -382,12 +387,9 @@ static VALUE rb_connect(VALUE self, VALUE user, VALUE pass, VALUE host, VALUE po
 
 /*
  * Terminate the connection; call this when the connection is no longer needed.
- * The garbage collector can close the connection, but doing so emits an
- * "Aborted connection" error on the server and increments the Aborted_clients
- * status variable.
+ * To have the garbage collector close the connection, enable +automatic_close+.
  *
- * @see http://dev.mysql.com/doc/en/communication-errors.html
- * @return [void]
+ * @return [nil]
  */
 static VALUE rb_mysql_client_close(VALUE self) {
   GET_CLIENT(self);
@@ -1085,6 +1087,37 @@ static VALUE rb_mysql_client_encoding(VALUE self) {
 }
 #endif
 
+/* call-seq:
+ *    client.automatic_close?
+ *
+ * @return [Boolean]
+ */
+static VALUE get_automatic_close(VALUE self) {
+  GET_CLIENT(self);
+  return wrapper->automatic_close ? Qtrue : Qfalse;
+}
+
+/* call-seq:
+ *    client.automatic_close = true
+ *
+ * Set this to +true+ to let the garbage collector close this connection.
+ */
+static VALUE set_automatic_close(VALUE self, VALUE value) {
+  GET_CLIENT(self);
+  if (RTEST(value)) {
+    value = Qtrue;
+    wrapper->automatic_close = 1;
+  } else {
+#ifndef _WIN32
+    value = Qfalse;
+    wrapper->automatic_close = 0;
+#else
+    rb_raise(cMysql2Error, "Connections are always closed by garbage collector on Windows");
+#endif
+  }
+  return value;
+}
+
 /* call-seq:
  *    client.reconnect = true
  *
@@ -1268,6 +1301,8 @@ void init_mysql2_client() {
   rb_define_method(cMysql2Client, "more_results?", rb_mysql_client_more_results, 0);
   rb_define_method(cMysql2Client, "next_result", rb_mysql_client_next_result, 0);
   rb_define_method(cMysql2Client, "store_result", rb_mysql_client_store_result, 0);
+  rb_define_method(cMysql2Client, "automatic_close?", get_automatic_close, 0);
+  rb_define_method(cMysql2Client, "automatic_close=", set_automatic_close, 1);
   rb_define_method(cMysql2Client, "reconnect=", set_reconnect, 1);
   rb_define_method(cMysql2Client, "warning_count", rb_mysql_client_warning_count, 0);
   rb_define_method(cMysql2Client, "query_info_string", rb_mysql_info, 0);

ext/mysql2/client.h

@@ -43,6 +43,7 @@ typedef struct {
   int reconnect_enabled;
   unsigned int connect_timeout;
   int active;
+  int automatic_close;
   int connected;
   int initialized;
   int refcount;

lib/mysql2/client.rb

@@ -30,10 +30,10 @@ def initialize(opts = {})
       opts[:connect_timeout] = 120 unless opts.key?(:connect_timeout)
 
       # TODO: stricter validation rather than silent massaging
-      [:reconnect, :connect_timeout, :local_infile, :read_timeout, :write_timeout, :default_file, :default_group, :secure_auth, :init_command].each do |key|
+      [:reconnect, :connect_timeout, :local_infile, :read_timeout, :write_timeout, :default_file, :default_group, :secure_auth, :init_command, :automatic_close].each do |key|
         next unless opts.key?(key)
         case key
-        when :reconnect, :local_infile, :secure_auth
+        when :reconnect, :local_infile, :secure_auth, :automatic_close
           send(:"#{key}=", !!opts[key]) # rubocop:disable Style/DoubleNegation
         when :connect_timeout, :read_timeout, :write_timeout
           send(:"#{key}=", opts[key].to_i)

spec/mysql2/client_spec.rb

@@ -210,6 +210,62 @@ def run_gc
     expect { client.query('SELECT 1') }.to_not raise_exception
   end
 
+  context "#automatic_close" do
+    if RUBY_PLATFORM =~ /mingw|mswin/
+      it "is enabled by default" do
+        client = Mysql2::Client.new(DatabaseCredentials['root'])
+        expect(client.automatic_close?).to be(true)
+      end
+
+      it "cannot be disabled" do
+        expect {
+          Mysql2::Client.new(DatabaseCredentials['root'].merge(:automatic_close => false))
+        }.to raise_error(Mysql2::Error)
+
+        client = Mysql2::Client.new(DatabaseCredentials['root'])
+
+        expect { client.automatic_close = false }.to raise_error(Mysql2::Error)
+        expect { client.automatic_close = true }.to_not raise_error
+      end
+    else
+      it "is disabled by default" do
+        client = Mysql2::Client.new(DatabaseCredentials['root'])
+        expect(client.automatic_close?).to be(false)
+      end
+
+      it "can be configured" do
+        client = Mysql2::Client.new(DatabaseCredentials['root'].merge(:automatic_close => true))
+        expect(client.automatic_close?).to be(true)
+      end
+
+      it "can be assigned" do
+        client = Mysql2::Client.new(DatabaseCredentials['root'])
+        client.automatic_close = true
+        expect(client.automatic_close?).to be(true)
+
+        client.automatic_close = false
+        expect(client.automatic_close?).to be(false)
+
+        client.automatic_close = 9
+        expect(client.automatic_close?).to be(true)
+
+        client.automatic_close = nil
+        expect(client.automatic_close?).to be(false)
+      end
+    end
+
+    it "should terminate connections during garbage collection" do
+      run_gc
+      expect {
+        Mysql2::Client.new(DatabaseCredentials['root'].merge(:automatic_close => true)).query('SELECT 1')
+        run_gc
+      }.to_not change {
+        @client.query("SHOW STATUS LIKE 'Aborted_%'").to_a +
+          @client.query("SHOW STATUS LIKE 'Threads_connected'").to_a
+      }
+    end
+  end
+
   it "should be able to connect to database with numeric-only name" do
     creds = DatabaseCredentials['numericuser']
     @client.query "CREATE DATABASE IF NOT EXISTS `#{creds['database']}`"

spec/spec_helper.rb

@@ -90,5 +90,6 @@ def with_internal_encoding(encoding)
         "test", "test", 'val1', 'val1,val2'
       )
     ]
+    client.close
   end
 end