Prevent concurrent uses of #ping and #close

Fix: brianmario#1433

Apply the same locking mecanism `#query` uses to these two methods.

Author: byroot

ext/mysql2/client.c

@@ -420,6 +420,24 @@ static VALUE allocate(VALUE klass) {
   return obj;
 }
 
+static void rb_mysql_client_set_active_fiber(VALUE self) {
+  VALUE fiber_current = rb_fiber_current();
+  GET_CLIENT(self);
+
+  // see if this connection is still waiting on a result from a previous query
+  if (NIL_P(wrapper->active_fiber)) {
+    // mark this connection active
+    wrapper->active_fiber = fiber_current;
+  } else if (wrapper->active_fiber == fiber_current) {
+    rb_raise(cMysql2Error, "This connection is still waiting for a result, try again once you have the result");
+  } else {
+    VALUE inspect = rb_inspect(wrapper->active_fiber);
+    const char *thr = StringValueCStr(inspect);
+
+    rb_raise(cMysql2Error, "This connection is in use by: %s", thr);
+  }
+}
+
 /* call-seq:
  *    Mysql2::Client.escape(string)
  *
@@ -571,11 +589,14 @@ static VALUE rb_mysql_connect(VALUE self, VALUE user, VALUE pass, VALUE host, VA
  */
 static VALUE rb_mysql_client_close(VALUE self) {
   GET_CLIENT(self);
+  rb_mysql_client_set_active_fiber(self);
 
   if (wrapper->client) {
     rb_thread_call_without_gvl(nogvl_close, wrapper, RUBY_UBF_IO, 0);
   }
 
+  wrapper->active_fiber = Qnil;
+
   return Qnil;
 }
 
@@ -798,24 +819,6 @@ static VALUE disconnect_and_mark_inactive(VALUE self) {
   return Qnil;
 }
 
-static void rb_mysql_client_set_active_fiber(VALUE self) {
-  VALUE fiber_current = rb_fiber_current();
-  GET_CLIENT(self);
-
-  // see if this connection is still waiting on a result from a previous query
-  if (NIL_P(wrapper->active_fiber)) {
-    // mark this connection active
-    wrapper->active_fiber = fiber_current;
-  } else if (wrapper->active_fiber == fiber_current) {
-    rb_raise(cMysql2Error, "This connection is still waiting for a result, try again once you have the result");
-  } else {
-    VALUE inspect = rb_inspect(wrapper->active_fiber);
-    const char *thr = StringValueCStr(inspect);
-
-    rb_raise(cMysql2Error, "This connection is in use by: %s", thr);
-  }
-}
-
 /* call-seq:
  *    client.abandon_results!
  *
@@ -1233,12 +1236,16 @@ static void *nogvl_ping(void *ptr) {
  */
 static VALUE rb_mysql_client_ping(VALUE self) {
   GET_CLIENT(self);
+  rb_mysql_client_set_active_fiber(self);
 
+  VALUE result = Qnil;
   if (!CONNECTED(wrapper)) {
-    return Qfalse;
+    result = Qfalse;
   } else {
-    return (VALUE)rb_thread_call_without_gvl(nogvl_ping, wrapper->client, RUBY_UBF_IO, 0);
+    result = (VALUE)rb_thread_call_without_gvl(nogvl_ping, wrapper->client, RUBY_UBF_IO, 0);
   }
+  wrapper->active_fiber = Qnil;
+  return result;
 }
 
 /* call-seq:

spec/em/em_spec.rb

@@ -69,7 +69,13 @@
       let(:client) { Mysql2::EM::Client.new DatabaseCredentials['root'] }
       let(:error) { StandardError.new('some error') }
       before { allow(client).to receive(:async_result).and_raise(error) }
-      after { client.close }
+      after do
+        begin
+          client.close
+        rescue Mysql2::Error
+          # Client is still in use, can't close it.
+        end
+      end
 
       it "should swallow exceptions raised in by the client" do
         errors = []

spec/mysql2/client_spec.rb

@@ -1279,4 +1279,30 @@ def connect(*args); end
     expect(client.inspect).not_to include("pass")
     expect(client.inspect).not_to include("secretsecret")
   end
+
+  it "should not allow concurrent use of #ping" do
+    @client.ping
+    thread = Thread.new { @client.query("SELECT SLEEP(1)") }
+    thread.join(0.1)
+    10.times do
+      expect do
+        @client.ping
+      end.to raise_error(Mysql2::Error, /This connection is in use by/)
+    end
+    expect(thread.value.to_a).to eq([{ "SLEEP(1)" => 0 }])
+    expect(@client.ping).to eq(true)
+  end
+
+  it "should not allow concurrent use of #close" do
+    @client.ping
+    thread = Thread.new { @client.query("SELECT SLEEP(1)") }
+    thread.join(0.1)
+    10.times do
+      expect do
+        @client.close
+      end.to raise_error(Mysql2::Error, /This connection is in use by/)
+    end
+    expect(thread.value.to_a).to eq([{ "SLEEP(1)" => 0 }])
+    expect(@client.close).to be_nil
+  end
 end

spec/spec_helper.rb

@@ -178,6 +178,12 @@ def ssl_cert_host
   end
 
   config.after(:example) do
-    @clients.each(&:close)
+    @clients.each do |client|
+      begin
+        client.close
+      rescue Mysql2::Error
+        # Client is still in use, can't close it.
+      end
+    end
   end
 end