RUBY-3658 Ensure partial writes are retried (#2927)

* HELP-74319 Ensure partial reads are retried

* simplify the socket wait logic

* Ubuntu 20.04 has been retired on GitHub Actions

* reduce the number of tests run on GH

focus on modern servers, recent rubies

* well, let's try 22.04, then

* respect the reported number of bytes written

write() is blocking, but apparently may still not write the
entire requested block (for various reasons). It is safest to
NOT assume it wrote everything, and instead check the return
value for the number of bytes actually written.

* different errors being raised

Author: jamis

.github/workflows/test.yml

@@ -10,40 +10,15 @@ jobs:
     env:
       CI: true
       TESTOPTS: "-v"
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     continue-on-error: true
     strategy:
       fail-fast: false
       matrix:
-        os: [ ubuntu-20.04 ]
-        ruby: ["2.7", "3.0", "3.1", "3.2", "3.3"]
-        mongodb: ["4.4", "5.0", "6.0", "7.0", "8.0"]
-        topology: [replica_set, sharded_cluster]
-        include:
-          - os: macos
-            ruby: "2.7"
-            mongodb: "7.0"
-            topology: server
-          - os: macos
-            ruby: "3.0"
-            mongodb: "7.0"
-            topology: server
-          - os: ubuntu-latest
-            ruby: "2.7"
-            mongodb: "7.0"
-            topology: server
-          - os: ubuntu-latest
-            ruby: "3.1"
-            mongodb: "7.0"
-            topology: server
-          - os: ubuntu-latest
-            ruby: "3.2"
-            mongodb: "7.0"
-            topology: server
-          - os: ubuntu-latest
-            ruby: "3.2"
-            mongodb: "8.0"
-            topology: replica_set
+        os: [ ubuntu-22.04 ]
+        ruby: [ "3.2" ]
+        mongodb: [ "7.0", "8.0" ]
+        topology: [ replica_set, sharded_cluster ]
     steps:
     - name: repo checkout
       uses: actions/checkout@v2

lib/mongo/socket.rb

@@ -491,9 +491,8 @@ def write_without_timeout(*args)
         buf = buf.to_s
         i = 0
         while i < buf.length
-          chunk = buf[i...i+WRITE_CHUNK_SIZE]
-          @socket.write(chunk)
-          i += WRITE_CHUNK_SIZE
+          chunk = buf[i, WRITE_CHUNK_SIZE]
+          i += @socket.write(chunk)
         end
       end
     end
@@ -523,32 +522,40 @@ 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
+      end
+
+      !rv.nil?
+    end
+
     def unix_socket?(sock)
       defined?(UNIXSocket) && sock.is_a?(UNIXSocket)
     end

spec/integration/client_side_encryption/custom_endpoint_spec.rb

@@ -94,11 +94,7 @@
       end
 
       let(:error_regex) do
-        if BSON::Environment.jruby?
-          /SocketError/
-        else
-          /Connection refused/
-        end
+        /Connection refused|SocketError|SocketTimeoutError/
       end
 
       it_behaves_like 'raising a KMS error'

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,61 @@
       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
+            expect(socket.write(data, timeout: timeout)).
+              to be === data.length
+          end
+        end
+      end
+    end
+  end
 end