diff --git a/contrib/ruby/ext/trilogy-ruby/cext.c b/contrib/ruby/ext/trilogy-ruby/cext.c index 2430de67..296730de 100644 --- a/contrib/ruby/ext/trilogy-ruby/cext.c +++ b/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; diff --git a/contrib/ruby/test/client_test.rb b/contrib/ruby/test/client_test.rb index 841db8c2..171d46af 100644 --- a/contrib/ruby/test/client_test.rb +++ b/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 diff --git a/inc/trilogy/socket.h b/inc/trilogy/socket.h index 1ef8b31d..010deb1f 100644 --- a/inc/trilogy/socket.h +++ b/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 diff --git a/src/socket.c b/src/socket.c index 274ecb05..46b2aa92 100644 --- a/src/socket.c +++ b/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; }