From a0105faf49ade2c695523b88c927c402aae6eb53 Mon Sep 17 00:00:00 2001 From: Thomas Munro Date: Tue, 23 Mar 2021 22:15:23 +1300 Subject: [PATCH v6 2/2] WIP -- use secure_read() to peek XXX Just a sketch, probably has bugs --- doc/src/sgml/config.sgml | 4 +- src/backend/libpq/pqcomm.c | 84 ++++++++++++++++++------------------- src/backend/tcop/postgres.c | 11 ++++- src/include/libpq/libpq.h | 2 +- 4 files changed, 53 insertions(+), 48 deletions(-) diff --git a/doc/src/sgml/config.sgml b/doc/src/sgml/config.sgml index 5cd0d38dbf..e522be460c 100644 --- a/doc/src/sgml/config.sgml +++ b/doc/src/sgml/config.sgml @@ -1008,8 +1008,8 @@ include_dir 'conf.d' Sets the time interval between checks that the client is still connected, while running queries. The check is performed by testing - whether one byte could be read from the socket with - MSG_PEEK. If the kernel reports that the connection + whether a part of the next message can be read from the client. + If the kernel reports that the connection has been closed or lost, a long running query can abort immediately, rather than discovering the problem when it eventually tries to send the response. diff --git a/src/backend/libpq/pqcomm.c b/src/backend/libpq/pqcomm.c index 74b309a1ef..a763a1b535 100644 --- a/src/backend/libpq/pqcomm.c +++ b/src/backend/libpq/pqcomm.c @@ -921,13 +921,13 @@ socket_set_nonblocking(bool nonblocking) } /* -------------------------------- - * pq_recvbuf - load some bytes into the input buffer + * pq_recvbuf_ext - load some bytes into the input buffer * * returns 0 if OK, EOF if trouble * -------------------------------- */ static int -pq_recvbuf(void) +pq_recvbuf_ext(bool nonblocking) { if (PqRecvPointer > 0) { @@ -943,8 +943,7 @@ pq_recvbuf(void) PqRecvLength = PqRecvPointer = 0; } - /* Ensure that we're in blocking mode */ - socket_set_nonblocking(false); + socket_set_nonblocking(nonblocking); /* Can fill buffer from PqRecvLength and upwards */ for (;;) @@ -956,6 +955,9 @@ pq_recvbuf(void) if (r < 0) { + if (nonblocking && (errno == EAGAIN || errno == EWOULDBLOCK)) + return 0; + if (errno == EINTR) continue; /* Ok if interrupted */ @@ -983,6 +985,13 @@ pq_recvbuf(void) } } +static int +pq_recvbuf(void) +{ + return pq_recvbuf_ext(false); +} + + /* -------------------------------- * pq_getbyte - get a single byte from connection, or return EOF * -------------------------------- @@ -1924,52 +1933,39 @@ pq_settcpusertimeout(int timeout, Port *port) return STATUS_OK; } -/* -------------------------------- - * pq_check_client_connection - is the client still connected? - * -------------------------------- +/* + * Proactively check if the client connection has gone away. Return 0 if still + * connected or currently reading a message already, EOF if disconnected, and 1 + * if at least one byte is available to read. If 1 is returned, the first + * byte of the next message is written to *c without consuming it, but we can't + * find out if the client has disconnected until we consume more data. */ -bool -pq_check_client_connection(void) +int +pq_check_client_connection(unsigned char *c) { - bool connected = true; - char nextbyte; - int r; - -retry: -#ifdef WIN32 - pgwin32_noblock = 1; -#endif - r = recv(MyProcPort->sock, &nextbyte, 1, MSG_PEEK); -#ifdef WIN32 - pgwin32_noblock = 0; -#endif + /* We're already in the middle of a message, so assume we are connected. */ + if (PqCommReadingMsg) + return 0; - if (r == 0) - { - /* EOF detected. */ - connected = false; - } - else if (r > 0) - { - /* Data available to read. Connection looks good. */ - } - else if (errno == EINTR) - { - /* Interrupted by a signal, so retry. */ - goto retry; - } - else if (errno == EAGAIN || errno == EWOULDBLOCK) + /* Do we have a byte already in our read buffer? */ + if (PqRecvPointer < PqRecvLength) { - /* No data available to read. Connection looks good. */ + *c = PqRecvBuffer[PqRecvPointer]; + return 1; } - else + + /* Try to read at least one byte from secure_read() without blocking. */ + if (pq_recvbuf_ext(true)) + return EOF; + + /* Now do we have a byte in our read buffer? */ + if (PqRecvPointer < PqRecvLength) { - /* Got some other error. We'd better log the reason. */ - ereport(COMMERROR, - (errcode(ERRCODE_CONNECTION_EXCEPTION), - errmsg("could not check client connection: %m"))); - connected = false; + *c = PqRecvBuffer[PqRecvPointer]; + return 1; } - return connected; + /* No data, but we're still connected. */ + + return 0; } diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c index 44fb7b7c30..23d1ee84da 100644 --- a/src/backend/tcop/postgres.c +++ b/src/backend/tcop/postgres.c @@ -3170,7 +3170,16 @@ ProcessInterrupts(void) */ if (!DoingCommandRead && client_connection_check_interval > 0) { - if (!pq_check_client_connection()) + unsigned char peekbyte; + int r; + + /* + * Try to peek ahead to see if we've been disconnected. If we see + * a pipelined 'X' message, we'll treat that as a disconnection + * too, but we don't try to look further ahead than that. + */ + r = pq_check_client_connection(&peekbyte); + if (r == EOF || (r == 1 && peekbyte == 'X')) ClientConnectionLost = true; enable_timeout_after(CLIENT_CONNECTION_CHECK_TIMEOUT, client_connection_check_interval); diff --git a/src/include/libpq/libpq.h b/src/include/libpq/libpq.h index 3bd97c4e93..f6cd3ee067 100644 --- a/src/include/libpq/libpq.h +++ b/src/include/libpq/libpq.h @@ -71,7 +71,7 @@ extern int pq_getbyte(void); extern int pq_peekbyte(void); extern int pq_getbyte_if_available(unsigned char *c); extern int pq_putmessage_v2(char msgtype, const char *s, size_t len); -extern bool pq_check_client_connection(void); +extern int pq_check_client_connection(unsigned char *c); /* * prototypes for functions in be-secure.c -- 2.30.1