Add connection check to Trilogy

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);

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)
 

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

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;
+    }
+}

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)                                                                                       \

test/socket_test.c

@@ -0,0 +1,68 @@
+#include <errno.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#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;
+}