Fix RUBY-1898 Socket reads may raise Timeout::Error (#1482)

* RUBY-1898 Correctly map select timeouts to Mongo::Error::TimeoutError

* RUBY-1898 Handle IO::WaitWritable when reading

* RUBY-1898 Raise more informative exceptions

* RUBY-1898 Rewrite the code to avoid weird failures

* Separate contexts to avoid unintentional connection reuse

* RUBY-1898 Handle IO::WaitWritable when reading

* Close all clients before

* Remove default 5 legacy read retries in the test suite since it interferes with assertions

* Remove all read retries for the connection timeout test

* Remove unneeded line

* Remove obsolete IO::WaitWritable rescue

Author: p-mongo

lib/mongo/socket.rb

@@ -49,6 +49,9 @@ class Socket
     # @return [ Hash ] The options.
     attr_reader :options
 
+    # @return [ Float ] timeout The socket timeout.
+    attr_reader :timeout
+
     # Is the socket connection alive?
     #
     # @example Is the socket alive?
@@ -123,10 +126,14 @@ def initialize(family)
     def read(length)
       handle_errors do
         data = read_from_socket(length)
-        raise IOError unless (data.length > 0 || length == 0)
+        unless (data.length > 0 || length == 0)
+          raise IOError, "Expected to read > 0 bytes but read 0 bytes"
+        end
         while data.length < length
           chunk = read_from_socket(length - data.length)
-          raise IOError unless (chunk.length > 0 || length == 0)
+          unless (chunk.length > 0 || length == 0)
+            raise IOError, "Expected to read > 0 bytes but read 0 bytes"
+          end
           data << chunk
         end
         data
@@ -185,7 +192,9 @@ def read_from_socket(length)
         return ''.force_encoding('BINARY')
       end
 
-      deadline = (Time.now + timeout) if timeout
+      if _timeout = self.timeout
+        deadline = Time.now + _timeout
+      end
 
       # We want to have a fixed and reasonably small size buffer for reads
       # because, for example, OpenSSL reads in 16 kb chunks max.
@@ -232,10 +241,22 @@ def read_from_socket(length)
           data[retrieved, chunk.length] = chunk
           retrieved += chunk.length
         end
-      rescue IO::WaitReadable
-        select_timeout = (deadline - Time.now) if deadline
-        if (select_timeout && select_timeout <= 0) || !Kernel::select([@socket], nil, [@socket], select_timeout)
-          raise Timeout::Error.new("Took more than #{timeout} seconds to receive data.")
+      # As explained in https://ruby-doc.com/core-trunk/IO.html#method-c-select,
+      # reading from a TLS socket may require writing which may raise WaitWritable
+      rescue IO::WaitReadable, IO::WaitWritable => exc
+        if deadline
+          select_timeout = deadline - Time.now
+          if select_timeout <= 0
+            raise Errno::ETIMEDOUT, "Took more than #{_timeout} seconds to receive data"
+          end
+        end
+        if exc.is_a?(IO::WaitReadable)
+          select_args = [[@socket], nil, [@socket], select_timeout]
+        else
+          select_args = [nil, [@socket], [@socket], select_timeout]
+        end
+        unless Kernel::select(*select_args)
+          raise Errno::ETIMEDOUT, "Took more than #{_timeout} seconds to receive data"
         end
         retry
       end

lib/mongo/socket/ssl.rb

@@ -33,9 +33,6 @@ class SSL < Socket
       # @return [ Integer ] port The port to connect to.
       attr_reader :port
 
-      # @return [ Float ] timeout The socket timeout.
-      attr_reader :timeout
-
       # Establishes a socket connection.
       #
       # @example Connect the socket.

lib/mongo/socket/tcp.rb

@@ -26,9 +26,6 @@ class TCP < Socket
       # @return [ Integer ] port The port to connect to.
       attr_reader :port
 
-      # @return [ Float ] timeout The socket timeout.
-      attr_reader :timeout
-
       # Establishes a socket connection.
       #
       # @example Connect the socket.

lib/mongo/socket/unix.rb

@@ -23,9 +23,6 @@ class Unix < Socket
       # @return [ String ] path The path to connect to.
       attr_reader :path
 
-      # @return [ Float ] timeout The socket timeout.
-      attr_reader :timeout
-
       # Initializes a new Unix socket.
       #
       # @example Create the Unix socket.

spec/mongo/server/connection_spec.rb

@@ -757,38 +757,42 @@ class ConnectionSpecTestException < Exception; end
       end
     end
 
-    context 'when a socket timeout is set' do
+    context 'when a socket timeout is set on client' do
 
       let(:connection) do
         described_class.new(server, socket_timeout: 10)
       end
 
-      it 'sets the timeout' do
+      it 'is propagated to connection timeout' do
         expect(connection.timeout).to eq(10)
       end
+    end
 
+    context 'when an operation never completes' do
       let(:client) do
-        authorized_client.with(socket_timeout: 1.5)
+        authorized_client.with(socket_timeout: 1.5,
+          # Read retries would cause the reads to be attempted twice,
+          # thus making the find take twice as long to time out.
+          retry_reads: false, max_read_retries: 0)
       end
 
       before do
-        authorized_collection.delete_many
-        authorized_collection.insert_one(a: 1)
+        client.cluster.next_primary
       end
 
-      it 'raises a timeout when it expires' do
+      it 'times out and raises SocketTimeoutError' do
         start = Time.now
         begin
           Timeout::timeout(1.5 + 2) do
             client[authorized_collection.name].find("$where" => "sleep(2000) || true").first
           end
         rescue => ex
           end_time = Time.now
-          expect(ex).to be_a(Timeout::Error)
-          expect(ex.message).to eq("Took more than 1.5 seconds to receive data.")
+          expect(ex).to be_a(Mongo::Error::SocketTimeoutError)
+          expect(ex.message).to match(/Took more than 1.5 seconds to receive data/)
         end
-        # Account for wait queue timeout (2s) and rescue
-        expect(end_time - start).to be_within(2.5).of(1.5)
+        # allow 1.5 seconds +- 0.5 seconds
+        expect(end_time - start).to be_within(1).of(2)
       end
 
       context 'when the socket_timeout is negative' do
@@ -816,7 +820,7 @@ class ConnectionSpecTestException < Exception; end
         it 'raises a timeout error' do
           expect {
             reply
-          }.to raise_exception(Timeout::Error)
+          }.to raise_exception(Mongo::Error::SocketTimeoutError)
         end
       end
     end

spec/mongo/socket_spec.rb

@@ -1,4 +1,4 @@
-require 'lite_spec_helper'
+require 'spec_helper'
 
 describe Mongo::Socket do
 
@@ -51,4 +51,66 @@
       end.to raise_error(Mongo::Error::SocketError, 'OpenSSL::SSL::SSLError: Test error (for fake-address) (MongoDB may not be configured with SSL support)')
     end
   end
+
+  describe '#read' do
+    let(:target_host) do
+      host = ClusterConfig.instance.primary_address_host
+      # Take ipv4 address
+      Socket.getaddrinfo(host, 0).detect { |ai| ai.first == 'AF_INET' }[3]
+    end
+
+    let(:socket) do
+      Mongo::Socket::TCP.new(target_host, ClusterConfig.instance.primary_address_port, 1, Socket::PF_INET)
+    end
+
+    let(:raw_socket) { socket.instance_variable_get('@socket') }
+
+    let(:wait_readable_class) do
+      Class.new(Exception) do
+        include IO::WaitReadable
+      end
+    end
+
+    context 'timeout' do
+      clean_slate_for_all
+
+      shared_examples_for 'times out' do
+        it 'times out' do
+          expect(socket).to receive(:timeout).at_least(:once).and_return(0.2)
+          # When we raise WaitWritable, the socket object is ready for
+          # writing which makes the read method invoke read_nonblock many times
+          expect(raw_socket).to receive(:read_nonblock).at_least(:once) do |len, buf|
+            sleep 0.01
+            raise exception_class
+          end
+
+          expect do
+            socket.read(10)
+          end.to raise_error(Mongo::Error::SocketTimeoutError, /Took more than .* seconds to receive data \(for /)
+        end
+      end
+
+      context 'with WaitReadable' do
+
+        let(:exception_class) do
+          Class.new(Exception) do
+            include IO::WaitReadable
+          end
+        end
+
+        it_behaves_like 'times out'
+      end
+
+      context 'with WaitWritable' do
+
+        let(:exception_class) do
+          Class.new(Exception) do
+            include IO::WaitWritable
+          end
+        end
+
+        it_behaves_like 'times out'
+      end
+    end
+  end
 end

spec/support/cluster_config.rb

@@ -74,6 +74,20 @@ def primary_address
     end
   end
 
+  def primary_address_str
+    primary_address
+  end
+
+  def primary_address_host
+    both = primary_address_str
+    both.split(':').first
+  end
+
+  def primary_address_port
+    both = primary_address_str
+    both.split(':')[1] || 27017
+  end
+
   # Try running a command on the admin database to see if the mongod was
   # started with auth.
   def auth_enabled?

spec/support/spec_config.rb

@@ -328,7 +328,6 @@ def base_test_options
     {
       max_pool_size: 1,
       heartbeat_frequency: 20,
-      max_read_retries: 5,
       # The test suite seems to perform a number of operations
       # requiring server selection. Hence a timeout of 1 here,
       # together with e.g. a misconfigured replica set,