Add PqreadPending to check for buffered rx data in SSL/ZPQ

Author: usernamedt

src/bin/pgbench/pgbench.c

@@ -6299,6 +6299,9 @@ threadRun(void *arg)
 		int			nsocks;		/* number of sockets to be waited for */
 		int64		min_usec;
 		int64		now_usec = 0;	/* set this only if needed */
+		bool		buffered_rx = false;	/* true if some of the clients has
+											 * data left in SSL/ZPQ read
+											 * buffers */
 
 		/*
 		 * identify which client sockets should be checked for input, and
@@ -6339,6 +6342,9 @@ threadRun(void *arg)
 				 */
 				int			sock = PQsocket(st->con);
 
+				/* check if conn has buffered SSL / ZPQ read data */
+				buffered_rx = buffered_rx || PQreadPending(st->con);
+
 				if (sock < 0)
 				{
 					pg_log_error("invalid socket: %s", PQerrorMessage(st->con));
@@ -6389,7 +6395,7 @@ threadRun(void *arg)
 			{
 				if (nsocks > 0)
 				{
-					rc = wait_on_socket_set(sockets, min_usec);
+					rc = buffered_rx ? 1 : wait_on_socket_set(sockets, min_usec);
 				}
 				else			/* nothing active, simple sleep */
 				{
@@ -6398,7 +6404,7 @@ threadRun(void *arg)
 			}
 			else				/* no explicit delay, wait without timeout */
 			{
-				rc = wait_on_socket_set(sockets, 0);
+				rc = buffered_rx ? 1 : wait_on_socket_set(sockets, 0);
 			}
 
 			if (rc < 0)
@@ -6437,8 +6443,11 @@ threadRun(void *arg)
 					pg_log_error("invalid socket: %s", PQerrorMessage(st->con));
 					goto done;
 				}
-
-				if (!socket_has_input(sockets, sock, nsocks++))
+				if (PQreadPending(st->con))
+				{
+					nsocks++;
+				}
+				else if (!socket_has_input(sockets, sock, nsocks++))
 					continue;
 			}
 			else if (st->state == CSTATE_FINISHED ||

src/interfaces/libpq/exports.txt

@@ -179,3 +179,4 @@ PQgetgssctx               176
 PQsetSSLKeyPassHook_OpenSSL         177
 PQgetSSLKeyPassHook_OpenSSL         178
 PQdefaultSSLKeyPassHook_OpenSSL     179
+PQreadPending                       180

src/interfaces/libpq/fe-connect.c

@@ -2157,12 +2157,6 @@ connectDBComplete(PGconn *conn)
 				return 1;		/* success! */
 
 			case PGRES_POLLING_READING:
-			    /* if there is some buffered RX data in ZpqStream
-			     * then don't proceed to pqWaitTimed */
-			    if (zpq_buffered_rx(conn->zstream)) {
-			        break;
-			    }
-
 				ret = pqWaitTimed(1, 0, conn, finish_time);
 				if (ret == -1)
 				{

src/interfaces/libpq/fe-exec.c

@@ -1810,9 +1810,7 @@ PQgetResult(PGconn *conn)
 		 * EOF indication.  We expect therefore that this won't result in any
 		 * undue delay in reporting a previous write failure.)
 		 */
-		if (flushResult || (zpq_buffered_rx(conn->zstream) == 0 &&
-			pqWait(true, false, conn)) ||
-			pqReadData(conn) < 0)
+		if (flushResult || pqWait(true, false, conn) || pqReadData(conn) < 0)
 		{
 			/*
 			 * conn->errorMessage has been set by pqWait or pqReadData. We
@@ -3285,6 +3283,12 @@ PQflush(PGconn *conn)
 	return pqFlush(conn);
 }
 
+int
+PQreadPending(PGconn *conn)
+{
+	return pqReadPending(conn);
+}
+
 
 /*
  *		PQfreemem - safely frees memory allocated

src/interfaces/libpq/fe-misc.c

@@ -1128,6 +1128,9 @@ pqWriteReady(PGconn *conn)
  *
  * If SSL is in use, the SSL buffer is checked prior to checking the socket
  * for read data directly.
+ *
+ * If ZPQ stream is in use, the ZPQ buffer is checked prior to checking
+ * the socket for read data directly.
  */
 static int
 pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
@@ -1143,14 +1146,10 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
 		return -1;
 	}
 
-#ifdef USE_SSL
-	/* Check for SSL library buffering read bytes */
-	if (forRead && conn->ssl_in_use && pgtls_read_pending(conn))
+	if (forRead && (pqReadPending(conn) > 0))
 	{
-		/* short-circuit the select */
 		return 1;
 	}
-#endif
 
 	/* We will retry as long as we get EINTR */
 	do
@@ -1169,6 +1168,33 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
 	return result;
 }
 
+/*
+ * Check if there is some data pending in ZPQ / SSL read buffers.
+ * Returns -1 on failure, 0 if no, 1 if yes.
+ */
+int
+pqReadPending(PGconn *conn)
+{
+	if (!conn)
+		return -1;
+
+	/* check for ZPQ stream buffered read bytes */
+	if (zpq_buffered_rx(conn->zpqStream))
+	{
+		/* short-circuit the select */
+		return 1;
+	}
+
+#ifdef USE_SSL
+	/* Check for SSL library buffering read bytes */
+	if (conn->ssl_in_use && pgtls_read_pending(conn))
+	{
+		/* short-circuit the select */
+		return 1;
+	}
+#endif
+	return 0;
+}
 
 /*
  * Check a file descriptor for read and/or write data, possibly waiting.

src/interfaces/libpq/fe-protocol3.c

@@ -1679,7 +1679,7 @@ pqGetCopyData3(PGconn *conn, char **buffer, int async)
 			if (async)
 				return 0;
 			/* Need to load more data */
-			if ((zpq_buffered_rx(conn->zstream) == 0 && pqWait(true, false, conn)) ||
+			if (pqWait(true, false, conn) ||
 				pqReadData(conn) < 0)
 				return -2;
 			continue;
@@ -1737,7 +1737,7 @@ pqGetline3(PGconn *conn, char *s, int maxlen)
 	while ((status = PQgetlineAsync(conn, s, maxlen - 1)) == 0)
 	{
 		/* need to load more data */
-		if ((zpq_buffered_rx(conn->zstream) == 0 && pqWait(true, false, conn)) ||
+		if (pqWait(true, false, conn) ||
 			pqReadData(conn) < 0)
 		{
 			*s = '\0';
@@ -1975,7 +1975,7 @@ pqFunctionCall3(PGconn *conn, Oid fnid,
 		if (needInput)
 		{
 			/* Wait for some data to arrive (or for the channel to close) */
-			if ((zpq_buffered_rx(conn->zstream) == 0 && pqWait(true, false, conn)) ||
+			if (pqWait(true, false, conn) ||
 				pqReadData(conn) < 0)
 				break;
 		}

src/interfaces/libpq/libpq-fe.h

@@ -461,6 +461,9 @@ extern PGPing PQpingParams(const char *const *keywords,
 /* Force the write buffer to be written (or at least try) */
 extern int	PQflush(PGconn *conn);
 
+extern int
+			PQreadPending(PGconn *conn);
+
 /*
  * "Fast path" interface --- not really recommended for application
  * use

src/interfaces/libpq/libpq-int.h

@@ -695,6 +695,7 @@ extern int	pqWaitTimed(int forRead, int forWrite, PGconn *conn,
 						time_t finish_time);
 extern int	pqReadReady(PGconn *conn);
 extern int	pqWriteReady(PGconn *conn);
+extern int	pqReadPending(PGconn *conn);
 
 /* === in fe-secure.c === */