Use Thread.handle_interrupt to protect query

When available, we prevent `Timeout.timeout` from corrupting
connections using `Thread.handle_interrupt`. Fixes brianmario#542.

Revert "Update specs for Ruby 2.1 Timeout behavior by specifying 'Timeout::Error' as the klass parameter."

This reverts commit d0a5199.

Author: tamird

ext/mysql2/client.c

@@ -1259,7 +1259,6 @@ void init_mysql2_client() {
   rb_define_singleton_method(cMysql2Client, "info", rb_mysql_client_info, 0);
 
   rb_define_method(cMysql2Client, "close", rb_mysql_client_close, 0);
-  rb_define_method(cMysql2Client, "query", rb_mysql_client_query, -1);
   rb_define_method(cMysql2Client, "abandon_results!", rb_mysql_client_abandon_results, 0);
   rb_define_method(cMysql2Client, "escape", rb_mysql_client_real_escape, 1);
   rb_define_method(cMysql2Client, "server_info", rb_mysql_client_server_info, 0);
@@ -1280,6 +1279,7 @@ void init_mysql2_client() {
   rb_define_method(cMysql2Client, "encoding", rb_mysql_client_encoding, 0);
 #endif
 
+  rb_define_private_method(cMysql2Client, "_query", rb_mysql_client_query, -1);
   rb_define_private_method(cMysql2Client, "connect_timeout=", set_connect_timeout, 1);
   rb_define_private_method(cMysql2Client, "read_timeout=", set_read_timeout, 1);
   rb_define_private_method(cMysql2Client, "write_timeout=", set_write_timeout, 1);

lib/mysql2/client.rb

@@ -74,6 +74,17 @@ def self.default_query_options
       @@default_query_options
     end
 
+    if Thread.respond_to?(:handle_interrupt)
+      def query(*args, &block)
+        Thread.handle_interrupt(Timeout::ExitException => :never) do
+          _query(*args, &block)
+        end
+      end
+    else
+      alias_method :query, :_query
+      public :query
+    end
+
     def query_info
       info = query_info_string
       return {} unless info

spec/mysql2/client_spec.rb

@@ -462,59 +462,40 @@ def connect *args
         }.should raise_error(Mysql2::Error)
       end
 
-      it "should close the connection when an exception is raised" do
-        begin
-          Timeout.timeout(1, Timeout::Error) do
-            @client.query("SELECT sleep(2)")
-          end
-        rescue Timeout::Error
-        end
 
-        lambda {
-          @client.query("SELECT 1")
-        }.should raise_error(Mysql2::Error, 'closed MySQL connection')
+      it 'should be impervious to connection-corrupting timeouts ' do
+        pending('`Thread.handle_interrupt` is not defined') unless Thread.respond_to?(:handle_interrupt)
+        # attempt to break the connection
+        expect { Timeout.timeout(0.1) { @client.query('SELECT SLEEP(1)') } }.to raise_error(Timeout::Error)
+
+        # expect the connection to not be broken
+        expect { @client.query('SELECT 1') }.to_not raise_error
       end
 
-      it "should handle Timeouts without leaving the connection hanging if reconnect is true" do
-        client = Mysql2::Client.new(DatabaseCredentials['root'].merge(:reconnect => true))
-        begin
-          Timeout.timeout(1, Timeout::Error) do
-            client.query("SELECT sleep(2)")
-          end
-        rescue Timeout::Error
+      context 'when a non-standard exception class is raised' do
+        it "should close the connection when an exception is raised" do
+          expect { Timeout.timeout(0.1, ArgumentError) { @client.query('SELECT SLEEP(1)') } }.to raise_error(ArgumentError)
+          expect { @client.query('SELECT 1') }.to raise_error(Mysql2::Error, 'closed MySQL connection')
         end
 
-        lambda {
-          client.query("SELECT 1")
-        }.should_not raise_error(Mysql2::Error)
-      end
+        it "should handle Timeouts without leaving the connection hanging if reconnect is true" do
+          client = Mysql2::Client.new(DatabaseCredentials['root'].merge(:reconnect => true))
 
-      it "should handle Timeouts without leaving the connection hanging if reconnect is set to true after construction true" do
-        client = Mysql2::Client.new(DatabaseCredentials['root'])
-        begin
-          Timeout.timeout(1, Timeout::Error) do
-            client.query("SELECT sleep(2)")
-          end
-        rescue Timeout::Error
+          expect { Timeout.timeout(0.1, ArgumentError) { client.query('SELECT SLEEP(1)') } }.to raise_error(ArgumentError)
+          expect { client.query('SELECT 1') }.to_not raise_error
         end
 
-        lambda {
-          client.query("SELECT 1")
-        }.should raise_error(Mysql2::Error)
+        it "should handle Timeouts without leaving the connection hanging if reconnect is set to true after construction true" do
+          client = Mysql2::Client.new(DatabaseCredentials['root'])
 
-        client.reconnect = true
+          expect { Timeout.timeout(0.1, ArgumentError) { client.query('SELECT SLEEP(1)') } }.to raise_error(ArgumentError)
+          expect { client.query('SELECT 1') }.to raise_error(Mysql2::Error)
 
-        begin
-          Timeout.timeout(1, Timeout::Error) do
-            client.query("SELECT sleep(2)")
-          end
-        rescue Timeout::Error
-        end
-
-        lambda {
-          client.query("SELECT 1")
-        }.should_not raise_error(Mysql2::Error)
+          client.reconnect = true
 
+          expect { Timeout.timeout(0.1, ArgumentError) { client.query('SELECT SLEEP(1)') } }.to raise_error(ArgumentError)
+          expect { client.query('SELECT 1') }.to_not raise_error
+        end
       end
 
       it "threaded queries should be supported" do