Support connect_timeout keyword arg in TCPSocket

Co-authored-by: Manef Zahra <manef.zahra@shopify.com>
Co-authored-by: Patrick Lin <patrick.lin@shopify.com>
Co-authored-by: Kevin Menard <kevin.menard@shopify.com>
Co-authored-by: Randy Stauner <randy.stauner@shopify.com>

closes #3421

Author: manefz

CHANGELOG.md

@@ -43,6 +43,7 @@ Compatibility:
 * Implement `Time#deconstruct_keys` from Ruby 3.2 (#3039, @rwstauner).
 * Do not autosplat a proc that accepts a single positional argument and keywords (#3039, @andrykonchin).
 * Support passing anonymous * and ** parameters as method call arguments (#3039, @andrykonchin).
+* Support `connect_timeout` keyword argument to `TCPSocket.{new,open}` (#3421, @manefz, @patricklinpl, @nirvdrum, @rwstauner).
 
 Performance:
 

lib/truffle/socket/tcp_socket.rb

@@ -44,7 +44,7 @@ def self.gethostbyname(hostname)
     [hostname, alternatives, family, *addresses]
   end
 
-  def initialize(host, service, local_host = nil, local_service = nil)
+  def initialize(host, service, local_host = nil, local_service = nil, connect_timeout: nil)
     @no_reverse_lookup = Primitive.class(self).do_not_reverse_lookup
 
     if host
@@ -108,7 +108,7 @@ def initialize(host, service, local_host = nil, local_service = nil)
       end
 
       connect_status = Truffle::Socket::Foreign
-        .connect(descriptor, Socket.sockaddr_in(port, address))
+        .connect(descriptor, Socket.sockaddr_in(port, address), connect_timeout)
 
       break if connect_status >= 0
     end

lib/truffle/socket/truffle/foreign.rb

@@ -27,6 +27,8 @@
 # NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
 # EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+require 'timeout'
+
 module Truffle
   module Socket
     module Foreign
@@ -107,14 +109,23 @@ def self.bind(descriptor, sockaddr)
         end
       end
 
-      def self.connect(descriptor, sockaddr)
+      def self.connect(descriptor, sockaddr, connect_timeout = nil)
         sockaddr = Socket.coerce_to_string(sockaddr)
 
         sockaddr_p = Primitive.io_thread_buffer_allocate(sockaddr.bytesize)
         begin
           sockaddr_p.write_bytes(sockaddr)
 
-          _connect(descriptor, sockaddr_p, sockaddr.bytesize)
+          if Primitive.nil?(connect_timeout)
+            _connect(descriptor, sockaddr_p, sockaddr.bytesize)
+          else
+            # In CRuby a connect_timeout of 0 fires immediately whereas a 0 for
+            # Timeout.timeout means "no timeout", so we change 0 to 1ns
+            # (the smallest time interval for clock_gettime).
+            Timeout.timeout(connect_timeout == 0 ? 0.000000001 : connect_timeout, IO::TimeoutError) do
+              _connect(descriptor, sockaddr_p, sockaddr.bytesize)
+            end
+          end
         ensure
           Primitive.io_thread_buffer_free(sockaddr_p)
         end

spec/tags/library/socket/tcpsocket/initialize_tags.txt

@@ -1,8 +1,7 @@
-fails:TCPSocket#initialize raises Errno::ETIMEDOUT with :connect_timeout when no server is listening on the given address
-fails:TCPSocket#initialize with a running server connects to a server when passed connect_timeout argument
-fails:TCPSocket#initialize raises IO::TimeoutError with :connect_timeout when no server is listening on the given address
 fails:TCPSocket#initialize with a running server does not use the given block and warns to use TCPSocket::open
 fails:TCPSocket#initialize using IPv4 when a server is listening on the given address creates a socket which is set to nonblocking
 fails:TCPSocket#initialize using IPv4 when a server is listening on the given address creates a socket which is set to close on exec
 fails:TCPSocket#initialize using IPv6 when a server is listening on the given address creates a socket which is set to nonblocking
 fails:TCPSocket#initialize using IPv6 when a server is listening on the given address creates a socket which is set to close on exec
+slow:TCPSocket#initialize raises IO::TimeoutError with :connect_timeout when no server is listening on the given address
+slow:TCPSocket#initialize with a running server connects to a server when passed connect_timeout argument

spec/tags/library/socket/tcpsocket/open_tags.txt

@@ -1,3 +1,2 @@
-fails:TCPSocket.open raises Errno::ETIMEDOUT with :connect_timeout when no server is listening on the given address
-fails:TCPSocket.open with a running server connects to a server when passed connect_timeout argument
-fails:TCPSocket.open raises IO::TimeoutError with :connect_timeout when no server is listening on the given address
+slow:TCPSocket.open raises IO::TimeoutError with :connect_timeout when no server is listening on the given address
+slow:TCPSocket.open with a running server connects to a server when passed connect_timeout argument

src/main/ruby/truffleruby/core/io.rb

@@ -38,6 +38,8 @@ class IO
 
   include Enumerable
 
+  class TimeoutError < IOError; end
+
   module WaitReadable; end
   module WaitWritable; end
 

test/mri/excludes/TestSocket_TCPSocket.rb

@@ -1,5 +1,4 @@
 exclude :test_initialize_failure, "needs investigation"
 exclude :test_inspect, "needs investigation"
 exclude :test_recvfrom, "needs investigation"
-exclude :test_initialize_connect_timeout, "needs investigation"
 exclude :test_initialize_resolv_timeout, "needs investigation"