HELP-74319 Ensure partial reads are retried

Author: jamis

lib/mongo/socket.rb

@@ -523,32 +523,41 @@ def write_with_timeout(*args, timeout:)
 
     def write_chunk(chunk, timeout)
       deadline = Utils.monotonic_time + timeout
+
       written = 0
-      begin
-        written += @socket.write_nonblock(chunk[written..-1])
-      rescue IO::WaitWritable, Errno::EINTR
-        select_timeout = deadline - Utils.monotonic_time
-        rv = Kernel.select(nil, [@socket], nil, select_timeout)
-        if BSON::Environment.jruby?
-          # Ignore the return value of Kernel.select.
-          # On JRuby, select appears to return nil prior to timeout expiration
-          # (apparently due to a EAGAIN) which then causes us to fail the read
-          # even though we could have retried it.
-          # Check the deadline ourselves.
-          if deadline
-            select_timeout = deadline - Utils.monotonic_time
-            if select_timeout <= 0
-              raise_timeout_error!("Took more than #{timeout} seconds to receive data", true)
-            end
+      while written < chunk.length
+        begin
+          written += @socket.write_nonblock(chunk[written..-1])
+        rescue IO::WaitWritable, Errno::EINTR
+          if !wait_for_socket_to_be_writable(deadline)
+            raise_timeout_error!("Took more than #{timeout} seconds to receive data", true)
           end
-        elsif rv.nil?
-          raise_timeout_error!("Took more than #{timeout} seconds to receive data (select call timed out)", true)
+
+          retry
         end
-        retry
       end
+
       written
     end
 
+    def wait_for_socket_to_be_writable(deadline)
+      select_timeout = deadline - Utils.monotonic_time
+      rv = Kernel.select(nil, [@socket], nil, select_timeout)
+      if BSON::Environment.jruby?
+        # Ignore the return value of Kernel.select.
+        # On JRuby, select appears to return nil prior to timeout expiration
+        # (apparently due to a EAGAIN) which then causes us to fail the read
+        # even though we could have retried it.
+        # Check the deadline ourselves.
+        select_timeout = deadline - Utils.monotonic_time
+        return select_timeout > 0
+      elsif rv.nil?
+        return false
+      end
+
+      true
+    end
+
     def unix_socket?(sock)
       defined?(UNIXSocket) && sock.is_a?(UNIXSocket)
     end

spec/mongo/socket_spec.rb

@@ -68,12 +68,6 @@
 
     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
 
@@ -116,4 +110,62 @@
       end
     end
   end
+
+  describe '#write' 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') }
+
+    context 'with timeout' do
+      let(:timeout) { 5_000 }
+
+      context 'data is less than WRITE_CHUNK_SIZE' do
+        let(:data) { "a" * 1024 }
+
+        context 'when a partial write occurs' do
+          before do
+            expect(raw_socket)
+              .to receive(:write_nonblock)
+              .twice
+              .and_return(data.length / 2)
+          end
+
+          it 'eventually writes everything' do
+            expect(socket.write(data, timeout: timeout)).
+              to be === data.length
+          end
+        end
+      end
+
+      context 'data is greater than WRITE_CHUNK_SIZE' do
+        let(:data) { "a" * (2 * Mongo::Socket::WRITE_CHUNK_SIZE + 256) }
+
+        context 'when a partial write occurs' do
+          before do
+            expect(raw_socket)
+              .to receive(:write_nonblock)
+              .exactly(4).times
+              .and_return(Mongo::Socket::WRITE_CHUNK_SIZE,
+                          128,
+                          Mongo::Socket::WRITE_CHUNK_SIZE - 128,
+                          256)
+          end
+
+          it 'eventually writes everything' do
+puts "=== spec begins"
+            expect(socket.write(data, timeout: timeout)).
+              to be === data.length
+          end
+        end
+      end
+    end
+  end
 end