Establish socket connection in Ruby

This moves the responsability for establishing the initial TCP or unix
socket connection to Ruby. This allows us to lean on Ruby's
implementation and configuration of DNS timeouts as well as happy
eyeballs/fast fallback.

In addition to improved timeouts, this should also deliver more detailed
and specific error messages.

Co-authored-by: Adam Hess <[email protected]>
Co-authored-by: Jean Boussier <[email protected]>

contrib/ruby/ext/trilogy-ruby/cext.c

@@ -295,42 +295,29 @@ static int _cb_ruby_wait(trilogy_sock_t *sock, trilogy_wait_t wait)
     return TRILOGY_OK;
 }
 
-struct nogvl_sock_args {
-    int rc;
-    trilogy_sock_t *sock;
-};
-
-static void *no_gvl_resolve(void *data)
+static int try_connect(struct trilogy_ctx *ctx, trilogy_handshake_t *handshake, const trilogy_sockopt_t *opts, int fd)
 {
-    struct nogvl_sock_args *args = data;
-    args->rc = trilogy_sock_resolve(args->sock);
-    return NULL;
-}
+    if (fd < 0) {
+        return TRILOGY_ERR;
+    }
 
-static int try_connect(struct trilogy_ctx *ctx, trilogy_handshake_t *handshake, const trilogy_sockopt_t *opts)
-{
     trilogy_sock_t *sock = trilogy_sock_new(opts);
     if (sock == NULL) {
         return TRILOGY_ERR;
     }
 
-    struct nogvl_sock_args args = {.rc = 0, .sock = sock};
-
-    // Do the DNS resolving with the GVL unlocked. At this point all
-    // configuration data is copied and available to the trilogy socket.
-    rb_thread_call_without_gvl(no_gvl_resolve, (void *)&args, RUBY_UBF_IO, NULL);
-
-    int rc = args.rc;
-
-    if (rc != TRILOGY_OK) {
-        trilogy_sock_close(sock);
-        return rc;
-    }
+    int rc;
 
     /* replace the default wait callback with our GVL-aware callback so we can
 escape the GVL on each wait operation without going through call_without_gvl */
     sock->wait_cb = _cb_ruby_wait;
-    rc = trilogy_connect_send_socket(&ctx->conn, sock);
+
+    int newfd = dup(fd);
+    if (newfd < 0) {
+        return TRILOGY_ERR;
+    }
+
+    rc = trilogy_connect_set_fd(&ctx->conn, sock, newfd);
     if (rc < 0) {
         trilogy_sock_close(sock);
         return rc;
@@ -447,7 +434,17 @@ static void authenticate(struct trilogy_ctx *ctx, trilogy_handshake_t *handshake
     }
 }
 
-static VALUE rb_trilogy_connect(VALUE self, VALUE encoding, VALUE charset, VALUE opts)
+#ifndef HAVE_RB_IO_DESCRIPTOR /* Ruby < 3.1 */
+static int rb_io_descriptor(VALUE io)
+{
+    rb_io_t *fptr;
+    GetOpenFile(io, fptr);
+    rb_io_check_closed(fptr);
+    return fptr->fd;
+}
+#endif
+
+static VALUE rb_trilogy_connect(VALUE self, VALUE raw_socket, VALUE encoding, VALUE charset, VALUE opts)
 {
     struct trilogy_ctx *ctx = get_ctx(self);
     trilogy_sockopt_t connopt = {0};
@@ -602,7 +599,16 @@ static VALUE rb_trilogy_connect(VALUE self, VALUE encoding, VALUE charset, VALUE
         connopt.tls_max_version = NUM2INT(val);
     }
 
-    int rc = try_connect(ctx, &handshake, &connopt);
+    VALUE io = rb_io_get_io(raw_socket);
+
+    rb_io_t *fptr;
+    GetOpenFile(io, fptr);
+    rb_io_check_readable(fptr);
+    rb_io_check_writable(fptr);
+
+    int fd = rb_io_descriptor(io);
+
+    int rc = try_connect(ctx, &handshake, &connopt, fd);
     if (rc != TRILOGY_OK) {
         if (connopt.path) {
             handle_trilogy_error(ctx, rc, "trilogy_connect - unable to connect to %s", connopt.path);
@@ -1141,7 +1147,7 @@ RUBY_FUNC_EXPORTED void Init_cext(void)
     VALUE Trilogy = rb_const_get(rb_cObject, rb_intern("Trilogy"));
     rb_define_alloc_func(Trilogy, allocate_trilogy);
 
-    rb_define_private_method(Trilogy, "_connect", rb_trilogy_connect, 3);
+    rb_define_private_method(Trilogy, "_connect", rb_trilogy_connect, 4);
     rb_define_method(Trilogy, "change_db", rb_trilogy_change_db, 1);
     rb_define_alias(Trilogy, "select_db", "change_db");
     rb_define_method(Trilogy, "query", rb_trilogy_query, 1);

contrib/ruby/ext/trilogy-ruby/extconf.rb

@@ -17,5 +17,6 @@
 have_library("crypto", "CRYPTO_malloc")
 have_library("ssl", "SSL_new")
 have_func("rb_interned_str", "ruby.h")
+have_func("rb_io_descriptor", "ruby.h") # Ruby 3.1+
 
 create_makefile "trilogy/cext"

contrib/ruby/lib/trilogy.rb

@@ -1,12 +1,21 @@
 # frozen_string_literal: true
 
+require "socket"
 require "trilogy/version"
 require "trilogy/error"
 require "trilogy/result"
 require "trilogy/cext"
 require "trilogy/encoding"
 
 class Trilogy
+  IO_TIMEOUT_ERROR =
+    if defined?(IO::TimeoutError)
+      IO::TimeoutError
+    else
+      Class.new
+    end
+  private_constant :IO_TIMEOUT_ERROR
+
   def initialize(options = {})
     options[:port] = options[:port].to_i if options[:port]
     mysql_encoding = options[:encoding] || "utf8mb4"
@@ -15,7 +24,54 @@ def initialize(options = {})
     @connection_options = options
     @connected_host = nil
 
-    _connect(encoding, charset, options)
+    socket = nil
+    begin
+      if host = options[:host]
+        port = options[:port] || 3306
+        connect_timeout = options[:connect_timeout] || options[:write_timeout]
+
+        socket = TCPSocket.new(host, port, connect_timeout: connect_timeout)
+
+        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+        if keepalive_enabled = options[:keepalive_enabled]
+          keepalive_idle = options[:keepalive_idle]
+          keepalive_interval = options[:keepalive_interval]
+          keepalive_count = options[:keepalive_count]
+
+          socket.setsockopt(Socket::SOL_SOCKET, Socket::SO_KEEPALIVE, true)
+
+          if keepalive_idle > 0 && defined?(Socket::TCP_KEEPIDLE)
+            socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_KEEPIDLE, keepalive_idle)
+          end
+          if keepalive_interval > 0 && defined?(Socket::TCP_KEEPINTVL)
+            socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_KEEPINTVL, keepalive_interval)
+          end
+          if keepalive_count > 0 && defined?(Socket::TCP_KEEPCNT)
+            socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_KEEPCNT, keepalive_count)
+          end
+        end
+      else
+        path = options[:socket] ||= "/tmp/mysql.sock"
+        socket = UNIXSocket.new(path)
+      end
+    rescue Errno::ETIMEDOUT, IO_TIMEOUT_ERROR => e
+      raise Trilogy::TimeoutError, e.message
+    rescue SocketError => e
+      connection_str = host ? "#{host}:#{port}" : path
+      raise Trilogy::BaseConnectionError, "unable to connect to \"#{connection_str}\": #{e.message}"
+    rescue => e
+      if e.respond_to?(:errno)
+        raise Trilogy::SyscallError.from_errno(e.errno, e.message)
+      else
+        raise
+      end
+    end
+
+    _connect(socket, encoding, charset, options)
+  ensure
+    # Socket's fd will be dup'd in C
+    socket&.close
   end
 
   def connection_options

contrib/ruby/test/client_test.rb

@@ -27,7 +27,9 @@ def test_trilogy_connect_tcp_to_wrong_port
     e = assert_raises Trilogy::ConnectionError do
       new_tcp_client port: 13307
     end
-    assert_equal "Connection refused - trilogy_connect - unable to connect to #{DEFAULT_HOST}:13307", e.message
+    assert_includes e.message, "Connection refused - Connection refused - connect"
+    assert_includes e.message, DEFAULT_HOST
+    assert_includes e.message, "13307"
   end
 
   def test_trilogy_connect_unix_socket
@@ -977,7 +979,7 @@ def test_connection_invalid_dns
     ex = assert_raises Trilogy::ConnectionError do
       new_tcp_client(host: "mysql.invalid", port: 3306)
     end
-    assert_equal "trilogy_connect - unable to connect to mysql.invalid:3306: TRILOGY_DNS_ERROR", ex.message
+    assert_includes ex.message, "unable to connect to \"mysql.invalid:3306\": getaddrinfo:"
   end
 
   def test_memsize

contrib/ruby/trilogy.gemspec

@@ -25,7 +25,6 @@ Gem::Specification.new do |s|
   s.required_ruby_version = ">= 3.0"
 
   s.add_dependency "bigdecimal"
-
   s.add_development_dependency "rake-compiler", "~> 1.0"
   s.add_development_dependency "minitest", "~> 5.5"
 end

inc/trilogy/client.h

@@ -150,6 +150,8 @@ int trilogy_connect_send(trilogy_conn_t *conn, const trilogy_sockopt_t *opts);
  */
 int trilogy_connect_send_socket(trilogy_conn_t *conn, trilogy_sock_t *sock);
 
+int trilogy_connect_set_fd(trilogy_conn_t *conn, trilogy_sock_t *sock, int fd);
+
 /* trilogy_connect_recv - Read the initial handshake from the server.
  *
  * This should be called after trilogy_connect_send returns TRILOGY_OK. Calling

inc/trilogy/socket.h

@@ -88,6 +88,8 @@ typedef struct trilogy_sock_t {
 
 static inline int trilogy_sock_connect(trilogy_sock_t *sock) { return sock->connect_cb(sock); }
 
+void trilogy_sock_set_fd(trilogy_sock_t *sock, int fd);
+
 static inline ssize_t trilogy_sock_read(trilogy_sock_t *sock, void *buf, size_t n)
 {
     return sock->read_cb(sock, buf, n);

src/client.c

@@ -311,6 +311,16 @@ int trilogy_connect_send_socket(trilogy_conn_t *conn, trilogy_sock_t *sock)
     return TRILOGY_OK;
 }
 
+int trilogy_connect_set_fd(trilogy_conn_t *conn, trilogy_sock_t *sock, int fd)
+{
+    trilogy_sock_set_fd(sock, fd);
+
+    conn->socket = sock;
+    conn->packet_parser.sequence_number = 0;
+
+    return TRILOGY_OK;
+}
+
 int trilogy_connect_recv(trilogy_conn_t *conn, trilogy_handshake_t *handshake_out)
 {
     int rc = read_packet(conn);

src/socket.c

@@ -25,6 +25,12 @@ struct trilogy_sock {
     SSL *ssl;
 };
 
+void trilogy_sock_set_fd(trilogy_sock_t *_sock, int fd)
+{
+    struct trilogy_sock *sock = (struct trilogy_sock *)_sock;
+    sock->fd = fd;
+}
+
 static int _cb_raw_fd(trilogy_sock_t *_sock)
 {
     struct trilogy_sock *sock = (struct trilogy_sock *)_sock;