Merge pull request #684 from cbandy/close-during-gc

sodabrew · sodabrew · commit 924d934df458 · 2016-02-16T08:54:43.000-08:00
Close connection during garbage collection
diff --git a/ext/mysql2/client.c b/ext/mysql2/client.c
@@ -234,7 +234,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,7 @@ 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);
+  wrapper->automatic_close = 1;
   wrapper->server_version = 0;
   wrapper->reconnect_enabled = 0;
   wrapper->connect_timeout = 0;
@@ -381,13 +382,12 @@ 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.
+ * Immediately disconnect from the server; normally the garbage collector
+ * will disconnect automatically when a connection is no longer needed.
+ * Explicitly closing this will free up server resources sooner than waiting
+ * for the garbage collector.
  *
- * @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);
@@ -1081,6 +1081,39 @@ 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 = false
+ *
+ * Set this to +false+ to leave the connection open after it is garbage
+ * collected. To avoid "Aborted connection" errors on the server, explicitly
+ * call +close+ when the connection is no longer needed.
+ *
+ * @see http://dev.mysql.com/doc/en/communication-errors.html
+ */
+static VALUE set_automatic_close(VALUE self, VALUE value) {
+  GET_CLIENT(self);
+  if (RTEST(value)) {
+    wrapper->automatic_close = 1;
+  } else {
+#ifndef _WIN32
+    wrapper->automatic_close = 0;
+#else
+    rb_warn("Connections are always closed by garbage collector on Windows");
+#endif
+  }
+  return value;
+}
+
 /* call-seq:
  *    client.reconnect = true
  *
@@ -1264,6 +1297,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);
diff --git a/ext/mysql2/client.h b/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;
diff --git a/lib/mysql2/client.rb b/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]) unless opts[key].nil?
diff --git a/spec/mysql2/client_spec.rb b/spec/mysql2/client_spec.rb
@@ -1,5 +1,6 @@
 # encoding: UTF-8
 require 'spec_helper'
+require 'stringio'
 
 RSpec.describe Mysql2::Client do
   context "using defaults file" do
@@ -172,48 +173,97 @@ def run_gc
     expect {
       Mysql2::Client.new(DatabaseCredentials['root']).close
     }.to_not change {
-      @client.query("SHOW STATUS LIKE 'Aborted_clients'").first['Value'].to_i
+      @client.query("SHOW STATUS LIKE 'Aborted_%'").to_a +
+        @client.query("SHOW STATUS LIKE 'Threads_connected'").to_a
     }
   end
 
   it "should not leave dangling connections after garbage collection" do
     run_gc
+    expect {
+      expect {
+        10.times do
+          Mysql2::Client.new(DatabaseCredentials['root']).query('SELECT 1')
+        end
+      }.to change {
+        @client.query("SHOW STATUS LIKE 'Threads_connected'").first['Value'].to_i
+      }.by(10)
 
-    client = Mysql2::Client.new(DatabaseCredentials['root'])
-    before_count = client.query("SHOW STATUS LIKE 'Threads_connected'").first['Value'].to_i
+      run_gc
+    }.to_not change {
+      @client.query("SHOW STATUS LIKE 'Aborted_%'").to_a +
+        @client.query("SHOW STATUS LIKE 'Threads_connected'").to_a
+    }
+  end
 
-    10.times do
-      Mysql2::Client.new(DatabaseCredentials['root']).query('SELECT 1')
+  context "#automatic_close" do
+    it "is enabled by default" do
+      client = Mysql2::Client.new(DatabaseCredentials['root'])
+      expect(client.automatic_close?).to be(true)
     end
-    after_count = client.query("SHOW STATUS LIKE 'Threads_connected'").first['Value'].to_i
-    expect(after_count).to eq(before_count + 10)
 
-    run_gc
-    final_count = client.query("SHOW STATUS LIKE 'Threads_connected'").first['Value'].to_i
-    expect(final_count).to eq(before_count)
-  end
+    if RUBY_PLATFORM =~ /mingw|mswin/
+      it "cannot be disabled" do
+        stderr, $stderr = $stderr, StringIO.new
 
-  it "should not close connections when running in a child process" do
-    pending("fork is not available on this platform") unless Process.respond_to?(:fork)
+        begin
+          Mysql2::Client.new(DatabaseCredentials['root'].merge(:automatic_close => false))
+          expect($stderr.string).to include('always closed by garbage collector')
+          $stderr.reopen
 
-    run_gc
-    client = Mysql2::Client.new(DatabaseCredentials['root'])
+          client = Mysql2::Client.new(DatabaseCredentials['root'])
+          client.automatic_close = false
+          expect($stderr.string).to include('always closed by garbage collector')
+          $stderr.reopen
 
-    # this empty `fork` call fixes this tests on RBX; without it, the next
-    # `fork` call hangs forever. WTF?
-    fork {}
+          expect { client.automatic_close = true }.to_not change { $stderr.string }
+        ensure
+          $stderr = stderr
+        end
+      end
+    else
+      it "can be configured" do
+        client = Mysql2::Client.new(DatabaseCredentials['root'].merge(:automatic_close => false))
+        expect(client.automatic_close?).to be(false)
+      end
 
-    fork do
-      client.query('SELECT 1')
-      client = nil
-      run_gc
-    end
+      it "can be assigned" do
+        client = Mysql2::Client.new(DatabaseCredentials['root'])
+        client.automatic_close = false
+        expect(client.automatic_close?).to be(false)
+
+        client.automatic_close = true
+        expect(client.automatic_close?).to be(true)
+
+        client.automatic_close = nil
+        expect(client.automatic_close?).to be(false)
+
+        client.automatic_close = 9
+        expect(client.automatic_close?).to be(true)
+      end
+
+      it "should not close connections when running in a child process" do
+        run_gc
+        client = Mysql2::Client.new(DatabaseCredentials['root'])
+        client.automatic_close = false
 
-    Process.wait
+        # this empty `fork` call fixes this tests on RBX; without it, the next
+        # `fork` call hangs forever. WTF?
+        fork {}
+
+        fork do
+          client.query('SELECT 1')
+          client = nil
+          run_gc
+        end
 
-    # this will throw an error if the underlying socket was shutdown by the
-    # child's GC
-    expect { client.query('SELECT 1') }.to_not raise_exception
+        Process.wait
+
+        # this will throw an error if the underlying socket was shutdown by the
+        # child's GC
+        expect { client.query('SELECT 1') }.to_not raise_exception
+      end
+    end
   end
 
   it "should be able to connect to database with numeric-only name" do
