Fix Timeout interrupt handling on Ruby 2.3 and protect Mysql2::Statement#execute

ext/mysql2/statement.c

@@ -470,7 +470,7 @@ void init_mysql2_statement() {
 
   rb_define_method(cMysql2Statement, "param_count", param_count, 0);
   rb_define_method(cMysql2Statement, "field_count", field_count, 0);
-  rb_define_method(cMysql2Statement, "execute", execute, -1);
+  rb_define_method(cMysql2Statement, "_execute", execute, -1);
   rb_define_method(cMysql2Statement, "fields", fields, 0);
   rb_define_method(cMysql2Statement, "last_id", rb_mysql_stmt_last_id, 0);
   rb_define_method(cMysql2Statement, "affected_rows", rb_mysql_stmt_affected_rows, 0);

lib/mysql2.rb

@@ -62,5 +62,23 @@ def self.key_hash_as_symbols(hash)
       return nil unless hash
       Hash[hash.map { |k, v| [k.to_sym, v] }]
     end
+
+    #
+    # In Mysql2::Client#query and Mysql2::Statement#execute,
+    # Thread#handle_interrupt is used to prevent Timeout#timeout
+    # from interrupting query execution.
+    #
+    # Timeout::ExitException was removed in Ruby 2.3.0, 2.2.3, and 2.1.8,
+    # but is present in earlier 2.1.x and 2.2.x, so we provide a shim.
+    #
+    if Thread.respond_to?(:handle_interrupt)
+      require 'timeout'
+      # rubocop:disable Style/ConstantName
+      TimeoutError = if defined?(::Timeout::ExitException)
+        ::Timeout::ExitException
+      else
+        ::Timeout::Error
+      end
+    end
   end
 end

lib/mysql2/client.rb

@@ -80,10 +80,8 @@ def initialize(opts = {})
     end
 
     if Thread.respond_to?(:handle_interrupt)
-      require 'timeout'
-
       def query(sql, options = {})
-        Thread.handle_interrupt(::Timeout::ExitException => :never) do
+        Thread.handle_interrupt(::Mysql2::Util::TimeoutError => :never) do
           _query(sql, @query_options.merge(options))
         end
       end

lib/mysql2/statement.rb

@@ -1,5 +1,17 @@
 module Mysql2
   class Statement
     include Enumerable
+
+    if Thread.respond_to?(:handle_interrupt)
+      def execute(*args)
+        Thread.handle_interrupt(::Mysql2::Util::TimeoutError => :never) do
+          _execute(*args)
+        end
+      end
+    else
+      def execute(*args)
+        _execute(*args)
+      end
+    end
   end
 end

spec/mysql2/client_spec.rb

@@ -467,7 +467,7 @@ def run_gc
         }.to raise_error(Mysql2::Error)
       end
 
-      it 'should be impervious to connection-corrupting timeouts ' do
+      it 'should be impervious to connection-corrupting timeouts in #query' 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(0.2)') } }.to raise_error(Timeout::Error)
@@ -476,6 +476,20 @@ def run_gc
         expect { @client.query('SELECT 1') }.to_not raise_error
       end
 
+      it 'should be impervious to connection-corrupting timeouts in #execute' do
+        # the statement handle gets corrupted and will segfault the tests if interrupted,
+        # so we can't even use pending on this test, really have to skip it on older Rubies.
+        skip('`Thread.handle_interrupt` is not defined') unless Thread.respond_to?(:handle_interrupt)
+
+        # attempt to break the connection
+        stmt = @client.prepare('SELECT SLEEP(?)')
+        expect { Timeout.timeout(0.1) { stmt.execute(0.2) } }.to raise_error(Timeout::Error)
+        stmt.close
+
+        # expect the connection to not be broken
+        expect { @client.query('SELECT 1') }.to_not raise_error
+      end
+
       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)