Use connect_timeout for initial connection

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

@@ -239,7 +239,18 @@ static int _cb_ruby_wait(trilogy_sock_t *sock, trilogy_wait_t wait)
         wait_flag = RB_WAITFD_OUT;
         break;
 
+    case TRILOGY_WAIT_CONNECT:
+        // wait for connection to be writable
+        timeout = &sock->opts.connect_timeout;
+        if (timeout->tv_sec == 0 && timeout->tv_usec == 0) {
+            // We used to use the write timeout for this, so if a connect timeout isn't configured, default to that.
+            timeout = &sock->opts.write_timeout;
+        }
+        wait_flag = RB_WAITFD_OUT;
+        break;
+
     case TRILOGY_WAIT_HANDSHAKE:
+        // wait for handshake packet on initial connection
         timeout = &sock->opts.connect_timeout;
         wait_flag = RB_WAITFD_IN;
         break;

contrib/ruby/test/client_test.rb

@@ -552,7 +552,7 @@ def test_write_timeout_closed_connection
     end
   end
 
-  def test_connect_timeout
+  def test_handshake_timeout
     serv = TCPServer.new(0)
     port = serv.addr[1]
 
@@ -563,6 +563,20 @@ def test_connect_timeout
     ensure_closed serv
   end
 
+  def test_connect_timeout
+    assert_raises Trilogy::TimeoutError do
+      # 192.0.2.0/24 is TEST-NET-1 which should only be for docs/examples
+      new_tcp_client(host: "192.0.2.1", connect_timeout: 0.1)
+    end
+  end
+
+  def test_connect_timeout_with_only_write_timeout
+    assert_raises Trilogy::TimeoutError do
+      # 192.0.2.0/24 is TEST-NET-1 which should only be for docs/examples
+      new_tcp_client(host: "192.0.2.1", write_timeout: 0.1)
+    end
+  end
+
   def test_large_query
     client = new_tcp_client
 

inc/trilogy/socket.h

@@ -13,6 +13,7 @@ typedef enum {
     TRILOGY_WAIT_READ = 0,
     TRILOGY_WAIT_WRITE = 1,
     TRILOGY_WAIT_HANDSHAKE = 2,
+    TRILOGY_WAIT_CONNECT = 3,
 } trilogy_wait_t;
 
 // We use the most strict mode as value 1 so if anyone ever

src/socket.c

@@ -40,6 +40,7 @@ static int _cb_wait(trilogy_sock_t *_sock, trilogy_wait_t wait)
     case TRILOGY_WAIT_READ:
         pfd.events = POLLIN;
         break;
+    case TRILOGY_WAIT_CONNECT:
     case TRILOGY_WAIT_WRITE:
         pfd.events = POLLOUT;
         break;
@@ -252,7 +253,7 @@ static int raw_connect_internal(struct trilogy_sock *sock, const struct addrinfo
         }
     }
 
-    if ((rc = trilogy_sock_wait_write((trilogy_sock_t *)sock)) < 0) {
+    if ((rc = trilogy_sock_wait((trilogy_sock_t *)sock, TRILOGY_WAIT_CONNECT)) < 0) {
         goto failrc;
     }