THRIFT-6127: Close Ruby sockets after I/O timeouts

lib/rb/lib/thrift/transport/socket.rb

@@ -66,12 +66,11 @@ def write(str)
 
           len
         end
-      rescue TransportException => e
-        # pass this on
-        raise e
+      rescue TransportException
+        close
+        raise
       rescue StandardError => e
-        close_socket(@handle)
-        @handle = nil
+        close
         raise TransportException.new(TransportException::NOT_OPEN, e.message)
       end
     end
@@ -95,12 +94,11 @@ def read(sz)
             end
           end
         end
-      rescue TransportException => e
-        # don't let this get caught by the StandardError handler
-        raise e
+      rescue TransportException
+        close
+        raise
       rescue StandardError => e
-        close_socket(@handle)
-        @handle = nil
+        close
         raise TransportException.new(TransportException::NOT_OPEN, e.message)
       end
       if (data.nil? || data.length == 0)

lib/rb/spec/server_socket_spec.rb

@@ -80,12 +80,14 @@
       expect(transport.timeout).to eq(Thrift::BaseServerTransport::DEFAULT_CLIENT_TIMEOUT)
 
       expect(sock).to receive(:write_nonblock).with("test data").and_raise(IO::EAGAINWaitWritable)
+      expect(sock).to receive(:close)
       expect(Process).to receive(:clock_gettime).with(Process::CLOCK_MONOTONIC).and_return(100.0, 100.0, 105.1)
       expect(IO).to receive(:select).with(nil, [sock], nil, 5.0).and_return(nil)
 
       expect { transport.write("test data") }.to raise_error(Thrift::TransportException) do |e|
         expect(e.type).to eq(Thrift::TransportException::TIMED_OUT)
       end
+      expect(transport).not_to be_open
     end
 
     it "should close the handle when closed" do

lib/rb/spec/socket_spec.rb

@@ -194,5 +194,26 @@
         expect(e.message).to eq("Could not connect to localhost:9090")
       end
     end
+
+    it "should close the descriptor and notify its peer after a partial read timeout" do
+      local, peer = ::Socket.pair(:UNIX, :STREAM, 0)
+      socket = Thrift::Socket.new
+      socket.handle = local
+      socket.timeout = 1
+      peer.write("A")
+
+      expect(Process).to receive(:clock_gettime).with(Process::CLOCK_MONOTONIC).and_return(100.0, 100.0, 101.0)
+
+      expect { socket.read_all(2) }.to raise_error(Thrift::TransportException) { |error|
+        expect(error.type).to eq(Thrift::TransportException::TIMED_OUT)
+      }
+      expect(socket).not_to be_open
+      expect(IO.select([peer], nil, nil, 1)).not_to be_nil
+      expect(peer.read).to eq("")
+    ensure
+      socket&.close
+      local&.close unless local&.closed?
+      peer&.close
+    end
   end
 end

lib/rb/spec/socket_spec_shared.rb

@@ -133,6 +133,50 @@
     expect { @socket.write("test data") }.to raise_error(Thrift::TransportException) { |e| expect(e.type).to eq(Thrift::TransportException::TIMED_OUT) }
   end
 
+  it "should close after a read timeout follows partial progress" do
+    @socket.open
+    @socket.timeout = 1
+
+    expect(Process).to receive(:clock_gettime).with(Process::CLOCK_MONOTONIC).and_return(100.0, 100.0, 101.0)
+    expect(@handle).to receive(:read_nonblock).with(4).ordered.and_return("A")
+    expect(@handle).to receive(:read_nonblock).with(3).ordered.and_raise(IO::EAGAINWaitReadable)
+    expect(@handle).to receive(:close).once
+
+    expect { @socket.read_all(4) }.to raise_error(
+      Thrift::TransportException,
+      /Timed out reading 3 bytes/
+    ) { |error|
+      expect(error.type).to eq(Thrift::TransportException::TIMED_OUT)
+    }
+    expect(@socket.handle).to be_nil
+    expect(@socket).not_to be_open
+    expect { @socket.read(1) }.to raise_error(Thrift::TransportException, "closed stream") { |error|
+      expect(error.type).to eq(Thrift::TransportException::NOT_OPEN)
+    }
+  end
+
+  it "should close after a write timeout follows partial progress" do
+    @socket.open
+    @socket.timeout = 1
+
+    expect(Process).to receive(:clock_gettime).with(Process::CLOCK_MONOTONIC).and_return(100.0, 101.0)
+    expect(@handle).to receive(:write_nonblock).with("test data").ordered.and_return(1)
+    expect(@handle).to receive(:write_nonblock).with("est data").ordered.and_raise(IO::EAGAINWaitWritable)
+    expect(@handle).to receive(:close).once
+
+    expect { @socket.write("test data") }.to raise_error(
+      Thrift::TransportException,
+      /Timed out writing 9 bytes/
+    ) { |error|
+      expect(error.type).to eq(Thrift::TransportException::TIMED_OUT)
+    }
+    expect(@socket.handle).to be_nil
+    expect(@socket).not_to be_open
+    expect { @socket.write("retry") }.to raise_error(Thrift::TransportException, "closed stream") { |error|
+      expect(error.type).to eq(Thrift::TransportException::NOT_OPEN)
+    }
+  end
+
   it "should read buffered SSL data without waiting on the raw socket again" do
     @socket.timeout = 1
     @socket.open