From bc403f061f062e111d5be9f06635636a7ac8824c Mon Sep 17 00:00:00 2001 From: Dirkjan Bussink Date: Sat, 28 Oct 2023 14:42:05 +0200 Subject: [PATCH] Add connection check to Trilogy This adds a very cheap connection check to Trilogy. This cheap check can be used before a connection for example is checked out from a connection pool. It allows for making the client more robust in the case of the server having disconnected (like a failover, an intermediate proxy like ProxySQL or VTGate restarted etc) but it can be reconnected safely. It works by doing a non-blocking `recv` to peek if there's any data. If there is, or if the function would otherwise block it means the connection is still safe. If the TCP connection has been closed, `recv` immediately returns with zero bytes read. This indicates the connection is no longer safe to use. We do this directly on the file descriptor to be consistent regardless of things like TLS being used. Co-authored-by: Daniel Colson Co-authored-by: John Hawthorn --- contrib/ruby/ext/trilogy-ruby/cext.c | 12 +++++ contrib/ruby/test/client_test.rb | 14 ++++++ inc/trilogy/socket.h | 16 +++++++ src/socket.c | 22 +++++++++ test/runner.c | 1 + test/socket_test.c | 68 ++++++++++++++++++++++++++++ 6 files changed, 133 insertions(+) create mode 100644 test/socket_test.c diff --git a/contrib/ruby/ext/trilogy-ruby/cext.c b/contrib/ruby/ext/trilogy-ruby/cext.c index 3a8bddcd..88926441 100644 --- a/contrib/ruby/ext/trilogy-ruby/cext.c +++ b/contrib/ruby/ext/trilogy-ruby/cext.c @@ -1020,6 +1020,17 @@ static VALUE rb_trilogy_closed(VALUE self) } } +static VALUE rb_trilogy_check(VALUE self) +{ + struct trilogy_ctx *ctx = get_open_ctx(self); + + int rc = trilogy_sock_check(ctx->conn.socket); + if (rc != TRILOGY_OK && rc != TRILOGY_AGAIN) { + handle_trilogy_error(ctx, rc, "trilogy_sock_check"); + } + return Qtrue; +} + static VALUE rb_trilogy_discard(VALUE self) { struct trilogy_ctx *ctx = get_ctx(self); @@ -1111,6 +1122,7 @@ RUBY_FUNC_EXPORTED void Init_cext() rb_define_method(Trilogy, "escape", rb_trilogy_escape, 1); rb_define_method(Trilogy, "close", rb_trilogy_close, 0); rb_define_method(Trilogy, "closed?", rb_trilogy_closed, 0); + rb_define_method(Trilogy, "check", rb_trilogy_check, 0); rb_define_method(Trilogy, "discard!", rb_trilogy_discard, 0); rb_define_method(Trilogy, "last_insert_id", rb_trilogy_last_insert_id, 0); rb_define_method(Trilogy, "affected_rows", rb_trilogy_affected_rows, 0); diff --git a/contrib/ruby/test/client_test.rb b/contrib/ruby/test/client_test.rb index e25811cd..439fdd64 100644 --- a/contrib/ruby/test/client_test.rb +++ b/contrib/ruby/test/client_test.rb @@ -477,6 +477,20 @@ def test_trilogy_closed? ensure_closed client end + def test_trilogy_check + client = new_tcp_client + + assert_equal true, client.check + + client.close + + assert_raises Trilogy::ConnectionClosed do + client.check + end + ensure + ensure_closed client + end + def test_read_timeout client = new_tcp_client(read_timeout: 0.1) diff --git a/inc/trilogy/socket.h b/inc/trilogy/socket.h index 967b5615..1ef8b31d 100644 --- a/inc/trilogy/socket.h +++ b/inc/trilogy/socket.h @@ -111,4 +111,20 @@ trilogy_sock_t *trilogy_sock_new(const trilogy_sockopt_t *opts); int trilogy_sock_resolve(trilogy_sock_t *raw); int trilogy_sock_upgrade_ssl(trilogy_sock_t *raw); +/* trilogy_sock_check - Verify if the socket is still alive and not disconnected. + * + * This check is very cheap to do and reduces the number of errors when for + * example the server has restarted since the connection was opened. In connection + * pooling implementations, this check can be done before the connection is + * returned. + * + * raw - A connected trilogy_sock_t pointer. Using a disconnected trilogy_sock_t is undefined. + * + * Return values: + * TRILOGY_OK - The connection is alive on the client side and can be. + * TRILOGY_CLOSED_CONNECTION - The connection is closed. + * TRILOGY_SYSERR - A system error occurred, check errno. + */ +int trilogy_sock_check(trilogy_sock_t *raw); + #endif diff --git a/src/socket.c b/src/socket.c index 9806aa74..3128d449 100644 --- a/src/socket.c +++ b/src/socket.c @@ -714,3 +714,25 @@ int trilogy_sock_upgrade_ssl(trilogy_sock_t *_sock) sock->ssl = NULL; return TRILOGY_OPENSSL_ERR; } + +int trilogy_sock_check(trilogy_sock_t *_sock) +{ + struct trilogy_sock *sock = (struct trilogy_sock *)_sock; + char buf[1]; + while (1) { + ssize_t data_read = recv(sock->fd, buf, 1, MSG_PEEK); + if (data_read > 0) { + return TRILOGY_OK; + } + if (data_read == 0) { + return TRILOGY_CLOSED_CONNECTION; + } + if (errno == EINTR) { + continue; + } + if (errno == EAGAIN || errno == EWOULDBLOCK) { + return TRILOGY_OK; + } + return TRILOGY_SYSERR; + } +} diff --git a/test/runner.c b/test/runner.c index f2400a36..b10f8f13 100644 --- a/test/runner.c +++ b/test/runner.c @@ -20,6 +20,7 @@ const trilogy_sockopt_t *get_connopt(void) { return &connopt; } SUITE(packet_parser_test) \ SUITE(charset_test) \ SUITE(blocking_test) \ + SUITE(socket_test) \ SUITE(parse_handshake_test) \ SUITE(parse_ok_packet_test) \ SUITE(parse_eof_packet_test) \ diff --git a/test/socket_test.c b/test/socket_test.c new file mode 100644 index 00000000..c72a6bb7 --- /dev/null +++ b/test/socket_test.c @@ -0,0 +1,68 @@ +#include +#include +#include +#include + +#include "test.h" + +#include "trilogy/client.h" +#include "trilogy/error.h" + +#define do_connect(CONN) \ + do { \ + int err = trilogy_init(CONN); \ + ASSERT_OK(err); \ + err = trilogy_connect(CONN, get_connopt()); \ + ASSERT_OK(err); \ + } while (0) + +TEST test_check_connected() +{ + trilogy_conn_t conn; + + do_connect(&conn); + + int err = trilogy_sock_check(conn.socket); + ASSERT_OK(err); + + trilogy_free(&conn); + PASS(); +} + + +TEST test_check_disconnected() +{ + trilogy_conn_t conn; + + do_connect(&conn); + shutdown(trilogy_sock_fd(conn.socket), SHUT_RD); + + int err = trilogy_sock_check(conn.socket); + ASSERT_ERR(TRILOGY_CLOSED_CONNECTION, err); + + trilogy_free(&conn); + PASS(); +} + +TEST test_check_closed() +{ + trilogy_conn_t conn; + + do_connect(&conn); + close_socket(&conn); + + int err = trilogy_sock_check(conn.socket); + ASSERT_ERR(TRILOGY_SYSERR, err); + + trilogy_free(&conn); + PASS(); +} + +int socket_test() +{ + RUN_TEST(test_check_connected); + RUN_TEST(test_check_disconnected); + RUN_TEST(test_check_closed); + + return 0; +}