Refactor pg_receivexlog main loop code, for readability, take 2.

Previously the source codes for processing the received data and handling
the end of stream were included in pg_receivexlog main loop. This commit
splits out them as separate functions. This is useful for improving the
readability of main loop code and making the future pg_receivexlog-related
patch simpler.

src/bin/pg_basebackup/receivelog.c

@@ -31,12 +31,23 @@ static char current_walfile_name[MAXPGPATH] = "";
 static bool reportFlushPosition = false;
 static XLogRecPtr lastFlushPosition = InvalidXLogRecPtr;
 
+static bool still_sending = true;		/* feedback still needs to be sent? */
+
 static PGresult *HandleCopyStream(PGconn *conn, XLogRecPtr startpos,
 				 uint32 timeline, char *basedir,
 			   stream_stop_callback stream_stop, int standby_message_timeout,
 				 char *partial_suffix, XLogRecPtr *stoppos);
 static int CopyStreamPoll(PGconn *conn, long timeout_ms);
 static int CopyStreamReceive(PGconn *conn, long timeout, char **buffer);
+static bool ProcessKeepaliveMsg(PGconn *conn, char *copybuf, int len,
+								XLogRecPtr blockpos, int64 *last_status);
+static bool ProcessXLogDataMsg(PGconn *conn, char *copybuf, int len,
+							   XLogRecPtr *blockpos, uint32 timeline,
+							   char *basedir, stream_stop_callback stream_stop,
+							   char *partial_suffix);
+static PGresult *HandleEndOfCopyStream(PGconn *conn, char *copybuf,
+									   XLogRecPtr blockpos, char *basedir, char *partial_suffix,
+									   XLogRecPtr *stoppos);
 
 static bool ReadEndOfStreamingResult(PGresult *res, XLogRecPtr *startpos,
 						 uint32 *timeline);
@@ -740,16 +751,13 @@ HandleCopyStream(PGconn *conn, XLogRecPtr startpos, uint32 timeline,
 	char	   *copybuf = NULL;
 	int64		last_status = -1;
 	XLogRecPtr	blockpos = startpos;
-	bool		still_sending = true;
+
+	still_sending = true;
 
 	while (1)
 	{
 		int			r;
-		int			xlogoff;
-		int			bytes_left;
-		int			bytes_written;
 		int64		now;
-		int			hdr_len;
 		long		sleeptime;
 
 		/*
@@ -818,48 +826,162 @@ HandleCopyStream(PGconn *conn, XLogRecPtr startpos, uint32 timeline,
 			goto error;
 		if (r == -2)
 		{
-			PGresult   *res = PQgetResult(conn);
+			PGresult	*res = HandleEndOfCopyStream(conn, copybuf, blockpos,
+													 basedir, partial_suffix, stoppos);
+			if (res == NULL)
+				goto error;
+			else
+				return res;
+		}
 
-			/*
-			 * The server closed its end of the copy stream.  If we haven't
-			 * closed ours already, we need to do so now, unless the server
-			 * threw an error, in which case we don't.
-			 */
-			if (still_sending)
-			{
-				if (!close_walfile(basedir, partial_suffix, blockpos))
+		/* Check the message type. */
+		if (copybuf[0] == 'k')
 		{
-					/* Error message written in close_walfile() */
-					PQclear(res);
+			if (!ProcessKeepaliveMsg(conn, copybuf, r, blockpos,
+									 &last_status))
 				goto error;
 		}
-				if (PQresultStatus(res) == PGRES_COPY_IN)
-				{
-					if (PQputCopyEnd(conn, NULL) <= 0 || PQflush(conn))
+		else if (copybuf[0] == 'w')
 		{
-						fprintf(stderr,
-								_("%s: could not send copy-end packet: %s"),
-								progname, PQerrorMessage(conn));
-						PQclear(res);
+			if (!ProcessXLogDataMsg(conn, copybuf, r, &blockpos,
+									timeline, basedir, stream_stop, partial_suffix))
 				goto error;
 		}
-					PQclear(res);
-					res = PQgetResult(conn);
+		else
+		{
+			fprintf(stderr, _("%s: unrecognized streaming header: \"%c\"\n"),
+					progname, copybuf[0]);
+			goto error;
 		}
-				still_sending = false;
 	}
+
+error:
 	if (copybuf != NULL)
 		PQfreemem(copybuf);
-			copybuf = NULL;
-			*stoppos = blockpos;
-			return res;
+	return NULL;
+}
+
+/*
+ * Wait until we can read CopyData message, or timeout.
+ *
+ * Returns 1 if data has become available for reading, 0 if timed out
+ * or interrupted by signal, and -1 on an error.
+ */
+static int
+CopyStreamPoll(PGconn *conn, long timeout_ms)
+{
+	int			ret;
+	fd_set		input_mask;
+	struct timeval timeout;
+	struct timeval *timeoutptr;
+
+	if (PQsocket(conn) < 0)
+	{
+		fprintf(stderr, _("%s: socket not open"), progname);
+		return -1;
 	}
 
-		/* Check the message type. */
-		if (copybuf[0] == 'k')
+	FD_ZERO(&input_mask);
+	FD_SET(PQsocket(conn), &input_mask);
+
+	if (timeout_ms < 0)
+		timeoutptr = NULL;
+	else
+	{
+		timeout.tv_sec = timeout_ms / 1000L;
+		timeout.tv_usec = (timeout_ms % 1000L) * 1000L;
+		timeoutptr = &timeout;
+	}
+
+	ret = select(PQsocket(conn) + 1, &input_mask, NULL, NULL, timeoutptr);
+	if (ret == 0 || (ret < 0 && errno == EINTR))
+		return 0;		/* Got a timeout or signal */
+	else if (ret < 0)
 	{
+		fprintf(stderr, _("%s: select() failed: %s\n"),
+				progname, strerror(errno));
+		return -1;
+	}
+
+	return 1;
+}
+
+/*
+ * Receive CopyData message available from XLOG stream, blocking for
+ * maximum of 'timeout' ms.
+ *
+ * If data was received, returns the length of the data. *buffer is set to
+ * point to a buffer holding the received message. The buffer is only valid
+ * until the next CopyStreamReceive call.
+ *
+ * 0 if no data was available within timeout, or wait was interrupted
+ * by signal. -1 on error. -2 if the server ended the COPY.
+ */
+static int
+CopyStreamReceive(PGconn *conn, long timeout, char **buffer)
+{
+	char	   *copybuf = NULL;
+	int			rawlen;
+
+	if (*buffer != NULL)
+		PQfreemem(*buffer);
+	*buffer = NULL;
+
+	/* Try to receive a CopyData message */
+	rawlen = PQgetCopyData(conn, &copybuf, 1);
+	if (rawlen == 0)
+	{
+		/*
+		 * No data available. Wait for some to appear, but not longer than
+		 * the specified timeout, so that we can ping the server.
+		 */
+		if (timeout != 0)
+		{
+			int		ret;
+
+			ret = CopyStreamPoll(conn, timeout);
+			if (ret <= 0)
+				return ret;
+		}
+
+		/* Else there is actually data on the socket */
+		if (PQconsumeInput(conn) == 0)
+		{
+			fprintf(stderr,
+					_("%s: could not receive data from WAL stream: %s"),
+					progname, PQerrorMessage(conn));
+			return -1;
+		}
+
+		/* Now that we've consumed some input, try again */
+		rawlen = PQgetCopyData(conn, &copybuf, 1);
+		if (rawlen == 0)
+			return 0;
+	}
+	if (rawlen == -1)			/* end-of-streaming or error */
+		return -2;
+	if (rawlen == -2)
+	{
+		fprintf(stderr, _("%s: could not read COPY data: %s"),
+				progname, PQerrorMessage(conn));
+		return -1;
+	}
+
+	/* Return received messages to caller */
+	*buffer = copybuf;
+	return rawlen;
+}
+
+/*
+ * Process the keepalive message.
+ */
+static bool
+ProcessKeepaliveMsg(PGconn *conn, char *copybuf, int len,
+					XLogRecPtr blockpos, int64 *last_status)
+{
 	int			pos;
 	bool		replyRequested;
+	int64		now;
 
 	/*
 	 * Parse the keepalive message, enclosed in the CopyData message.
@@ -870,11 +992,11 @@ HandleCopyStream(PGconn *conn, XLogRecPtr startpos, uint32 timeline,
 	pos += 8;			/* skip walEnd */
 	pos += 8;			/* skip sendTime */
 
-			if (r < pos + 1)
+	if (len < pos + 1)
 	{
 		fprintf(stderr, _("%s: streaming header too small: %d\n"),
-						progname, r);
-				goto error;
+				progname, len);
+		return false;
 	}
 	replyRequested = copybuf[pos];
 
@@ -883,18 +1005,33 @@ HandleCopyStream(PGconn *conn, XLogRecPtr startpos, uint32 timeline,
 	{
 		now = feGetCurrentTimestamp();
 		if (!sendFeedback(conn, blockpos, now, false))
-					goto error;
-				last_status = now;
-			}
+			return false;
+		*last_status = now;
 	}
-		else if (copybuf[0] == 'w')
-		{
+
+	return true;
+}
+
+/*
+ * Process XLogData message.
+ */
+static bool
+ProcessXLogDataMsg(PGconn *conn, char *copybuf, int len,
+				   XLogRecPtr *blockpos, uint32 timeline,
+				   char *basedir, stream_stop_callback stream_stop,
+				   char *partial_suffix)
+{
+	int			xlogoff;
+	int			bytes_left;
+	int			bytes_written;
+	int			hdr_len;
+
 	/*
 	 * Once we've decided we don't want to receive any more, just
 	 * ignore any subsequent XLogData messages.
 	 */
-			if (!still_sending)
-				continue;
+	if (!(still_sending))
+		return true;
 
 	/*
 	 * Read the header of the XLogData message, enclosed in the
@@ -905,16 +1042,16 @@ HandleCopyStream(PGconn *conn, XLogRecPtr startpos, uint32 timeline,
 	hdr_len += 8;		/* dataStart */
 	hdr_len += 8;		/* walEnd */
 	hdr_len += 8;		/* sendTime */
-			if (r < hdr_len)
+	if (len < hdr_len)
 	{
 		fprintf(stderr, _("%s: streaming header too small: %d\n"),
-						progname, r);
-				goto error;
+				progname, len);
+		return false;
 	}
-			blockpos = fe_recvint64(&copybuf[1]);
+	*blockpos = fe_recvint64(&copybuf[1]);
 
 	/* Extract WAL location for this block */
-			xlogoff = blockpos % XLOG_SEG_SIZE;
+	xlogoff = *blockpos % XLOG_SEG_SIZE;
 
 	/*
 	 * Verify that the initial location in the stream matches where we
@@ -928,7 +1065,7 @@ HandleCopyStream(PGconn *conn, XLogRecPtr startpos, uint32 timeline,
 			fprintf(stderr,
 					_("%s: received transaction log record for offset %u with no file open\n"),
 					progname, xlogoff);
-					goto error;
+			return false;
 		}
 	}
 	else
@@ -940,11 +1077,11 @@ HandleCopyStream(PGconn *conn, XLogRecPtr startpos, uint32 timeline,
 			fprintf(stderr,
 					_("%s: got WAL data offset %08x, expected %08x\n"),
 					progname, xlogoff, (int) lseek(walfile, 0, SEEK_CUR));
-					goto error;
+			return false;
 		}
 	}
 
-			bytes_left = r - hdr_len;
+	bytes_left = len - hdr_len;
 	bytes_written = 0;
 
 	while (bytes_left)
@@ -962,11 +1099,11 @@ HandleCopyStream(PGconn *conn, XLogRecPtr startpos, uint32 timeline,
 
 		if (walfile == -1)
 		{
-					if (!open_walfile(blockpos, timeline,
+			if (!open_walfile(*blockpos, timeline,
 							  basedir, partial_suffix))
 			{
 				/* Error logged by open_walfile */
-						goto error;
+				return false;
 			}
 		}
 
@@ -978,160 +1115,81 @@ HandleCopyStream(PGconn *conn, XLogRecPtr startpos, uint32 timeline,
 					_("%s: could not write %u bytes to WAL file \"%s\": %s\n"),
 					progname, bytes_to_write, current_walfile_name,
 					strerror(errno));
-					goto error;
+			return false;
 		}
 
 		/* Write was successful, advance our position */
 		bytes_written += bytes_to_write;
 		bytes_left -= bytes_to_write;
-				blockpos += bytes_to_write;
+		*blockpos += bytes_to_write;
 		xlogoff += bytes_to_write;
 
 		/* Did we reach the end of a WAL segment? */
-				if (blockpos % XLOG_SEG_SIZE == 0)
+		if (*blockpos % XLOG_SEG_SIZE == 0)
 		{
-					if (!close_walfile(basedir, partial_suffix, blockpos))
+			if (!close_walfile(basedir, partial_suffix, *blockpos))
 				/* Error message written in close_walfile() */
-						goto error;
+				return false;
 
 			xlogoff = 0;
 
-					if (still_sending && stream_stop(blockpos, timeline, true))
+			if (still_sending && stream_stop(*blockpos, timeline, true))
 			{
 				if (PQputCopyEnd(conn, NULL) <= 0 || PQflush(conn))
 				{
 					fprintf(stderr, _("%s: could not send copy-end packet: %s"),
 							progname, PQerrorMessage(conn));
-							goto error;
+					return false;
 				}
 				still_sending = false;
-						break;	/* ignore the rest of this XLogData packet */
+				return true;	/* ignore the rest of this XLogData packet */
 			}
 		}
 	}
 	/* No more data left to write, receive next copy packet */
-		}
-		else
-		{
-			fprintf(stderr, _("%s: unrecognized streaming header: \"%c\"\n"),
-					progname, copybuf[0]);
-			goto error;
-		}
-	}
-
-error:
-	if (copybuf != NULL)
-		PQfreemem(copybuf);
-	return NULL;
-}
-
-/*
- * Wait until we can read CopyData message, or timeout.
- *
- * Returns 1 if data has become available for reading, 0 if timed out
- * or interrupted by signal, and -1 on an error.
- */
-static int
-CopyStreamPoll(PGconn *conn, long timeout_ms)
-{
-	int			ret;
-	fd_set		input_mask;
-	struct timeval timeout;
-	struct timeval *timeoutptr;
 
-	if (PQsocket(conn) < 0)
-	{
-		fprintf(stderr, _("%s: socket not open"), progname);
-		return -1;
-	}
-
-	FD_ZERO(&input_mask);
-	FD_SET(PQsocket(conn), &input_mask);
-
-	if (timeout_ms < 0)
-		timeoutptr = NULL;
-	else
-	{
-		timeout.tv_sec = timeout_ms / 1000L;
-		timeout.tv_usec = (timeout_ms % 1000L) * 1000L;
-		timeoutptr = &timeout;
-	}
-
-	ret = select(PQsocket(conn) + 1, &input_mask, NULL, NULL, timeoutptr);
-	if (ret == 0 || (ret < 0 && errno == EINTR))
-		return 0;		/* Got a timeout or signal */
-	else if (ret < 0)
-	{
-		fprintf(stderr, _("%s: select() failed: %s\n"),
-				progname, strerror(errno));
-		return -1;
-	}
-
-	return 1;
+	return true;
 }
 
 /*
- * Receive CopyData message available from XLOG stream, blocking for
- * maximum of 'timeout' ms.
- *
- * If data was received, returns the length of the data. *buffer is set to
- * point to a buffer holding the received message. The buffer is only valid
- * until the next CopyStreamReceive call.
- *
- * 0 if no data was available within timeout, or wait was interrupted
- * by signal. -1 on error. -2 if the server ended the COPY.
+ * Handle end of the copy stream.
  */
-static int
-CopyStreamReceive(PGconn *conn, long timeout, char **buffer)
+static PGresult *
+HandleEndOfCopyStream(PGconn *conn, char *copybuf,
+					  XLogRecPtr blockpos, char *basedir, char *partial_suffix,
+					  XLogRecPtr *stoppos)
 {
-	char	   *copybuf = NULL;
-	int			rawlen;
-
-	if (*buffer != NULL)
-		PQfreemem(*buffer);
-	*buffer = NULL;
+	PGresult   *res = PQgetResult(conn);
 
-	/* Try to receive a CopyData message */
-	rawlen = PQgetCopyData(conn, &copybuf, 1);
-	if (rawlen == 0)
-	{
 	/*
-		 * No data available. Wait for some to appear, but not longer than
-		 * the specified timeout, so that we can ping the server.
+	 * The server closed its end of the copy stream.  If we haven't
+	 * closed ours already, we need to do so now, unless the server
+	 * threw an error, in which case we don't.
 	 */
-		if (timeout != 0)
+	if (still_sending)
 	{
-			int		ret;
-
-			ret = CopyStreamPoll(conn, timeout);
-			if (ret <= 0)
-				return ret;
+		if (!close_walfile(basedir, partial_suffix, blockpos))
+		{
+			/* Error message written in close_walfile() */
+			PQclear(res);
+			return NULL;
 		}
-
-		/* Else there is actually data on the socket */
-		if (PQconsumeInput(conn) == 0)
+		if (PQresultStatus(res) == PGRES_COPY_IN)
+		{
+			if (PQputCopyEnd(conn, NULL) <= 0 || PQflush(conn))
 			{
 				fprintf(stderr,
-					_("%s: could not receive data from WAL stream: %s"),
+						_("%s: could not send copy-end packet: %s"),
 						progname, PQerrorMessage(conn));
-			return -1;
+				PQclear(res);
+				return NULL;
 			}
-
-		/* Now that we've consumed some input, try again */
-		rawlen = PQgetCopyData(conn, &copybuf, 1);
-		if (rawlen == 0)
-			return 0;
+			res = PQgetResult(conn);
 		}
-	if (rawlen == -1)			/* end-of-streaming or error */
-		return -2;
-	if (rawlen == -2)
-	{
-		fprintf(stderr, _("%s: could not read COPY data: %s"),
-				progname, PQerrorMessage(conn));
-		return -1;
+		still_sending = false;
 	}
-
-	/* Return received messages to caller */
-	*buffer = copybuf;
-	return rawlen;
+	if (copybuf != NULL)
+		PQfreemem(copybuf);
+	*stoppos = blockpos;
+	return res;
 }