From aff4f4b99d62d630c54c8c6125eea66378a1a1b3 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. The implementation here matches what is also done in the Go driver, see also https://github.blog/2020-05-20-three-bugs-in-the-go-mysql-driver/ where this is decribed. 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. --- contrib/ruby/ext/trilogy-ruby/cext.c | 22 +++++++++ contrib/ruby/test/client_test.rb | 12 +++++ inc/trilogy/socket.h | 16 +++++++ src/socket.c | 22 +++++++++ test/runner.c | 1 + test/socket_test.c | 68 ++++++++++++++++++++++++++++ 6 files changed, 141 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..1fbadd0a 100644 --- a/contrib/ruby/ext/trilogy-ruby/cext.c +++ b/contrib/ruby/ext/trilogy-ruby/cext.c @@ -1020,6 +1020,27 @@ static VALUE rb_trilogy_closed(VALUE self) } } +static VALUE rb_trilogy_check(VALUE self) +{ + struct trilogy_ctx *ctx = get_ctx(self); + + if (ctx->conn.socket == NULL) { + return Qfalse; + } + + int rc = trilogy_sock_check(ctx->conn.socket); + switch (rc) { + case TRILOGY_OK: + return Qtrue; + case TRILOGY_CLOSED_CONNECTION: + return Qfalse; + case TRILOGY_SYSERR: + trilogy_syserr_fail_str(errno, rb_str_new_cstr("Failed to check connection")); + UNREACHABLE_RETURN(Qfalse); + } + return Qfalse; +} + static VALUE rb_trilogy_discard(VALUE self) { struct trilogy_ctx *ctx = get_ctx(self); @@ -1111,6 +1132,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..889fc22b 100644 --- a/contrib/ruby/test/client_test.rb +++ b/contrib/ruby/test/client_test.rb @@ -477,6 +477,18 @@ def test_trilogy_closed? ensure_closed client end + def test_trilogy_check? + client = new_tcp_client + + assert_predicate client, :check? + + client.close + + refute_predicate client, :check? + 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; +}