dimitri · dimitri · Apr 28, 2023 · Apr 27, 2023 · Apr 27, 2023
diff --git a/src/bin/pgcopydb/ld_apply.c b/src/bin/pgcopydb/ld_apply.c
@@ -303,6 +303,83 @@ stream_apply_sync_sentinel(StreamApplyContext *context)
 }
 
 
+/*
+ * stream_apply_send_sync_sentinel sends a query to sync with the pgcopydb
+ * sentinel table using the libpq async API. Use the associated function
+ * stream_apply_fetch_sync_sentinel to check for availability of the result and
+ * fetch it when it's available.
+ */
+bool
+stream_apply_send_sync_sentinel(StreamApplyContext *context)
+{
+	PGSQL *src = &(context->src);
+
+	if (src->connectionStatementType != PGSQL_CONNECTION_MULTI_STATEMENT)
+	{
+		log_error("BUG: stream_apply_send_sync_sentinel called "
+				  "in SINGLE statement mode");
+		return false;
+	}
+
+	if (context->sentinelQueryInProgress)
+	{
+		log_error("BUG: stream_apply_send_sync_sentinel already in progress");
+		return false;
+	}
+
+	if (!pgsql_init(src, context->source_pguri, PGSQL_CONN_SOURCE))
+	{
+		/* errors have already been logged */
+		return false;
+	}
+
+	/* limit the amount of logging of the apply process */
+	src->logSQL = true;
+
+	if (!pgsql_send_sync_sentinel_apply(src, context->previousLSN))
+	{
+		log_error("Failed to sync progress with the pgcopydb sentinel");
+		return false;
+	}
+
+	context->sentinelSyncTime = time(NULL);
+	context->sentinelQueryInProgress = true;
+
+	return true;
+}
+
+
+/*
+ * stream_apply_fetch_sync_sentinel checks to see if the result of the
+ * pgcopydb sentinel sync query is available and fetches it then.
+ */
+bool
+stream_apply_fetch_sync_sentinel(StreamApplyContext *context)
+{
+	PGSQL *src = &(context->src);
+
+	bool retry;
+	CopyDBSentinel sentinel = { 0 };
+
+	if (!pgsql_fetch_sync_sentinel_apply(src, &retry, &sentinel))
+	{
+		log_error("Failed to fetch sentinel update query results");
+		return false;
+	}
+
+	if (!retry)
+	{
+		context->sentinelQueryInProgress = false;
+
+		context->apply = sentinel.apply;
+		context->endpos = sentinel.endpos;
+		context->startpos = sentinel.startpos;
+	}
+
+	return true;
+}
+
+
 /*
  * stream_apply_file connects to the target database system and applies the
  * given SQL file as prepared by the stream_transform_file function.

diff --git a/src/bin/pgcopydb/ld_replay.c b/src/bin/pgcopydb/ld_replay.c
@@ -118,7 +118,17 @@ stream_apply_replay(StreamSpecs *specs)
 	/* we might still have to disconnect now */
 	(void) pgsql_finish(&(context->pgsql));
 
-	log_notice("Replayed %lld messages", (long long) readerContext.lineno);
+	/* make sure to send a last round of sentinel update before exit */
+	if (!stream_apply_sync_sentinel(context))
+	{
+		log_error("Failed to update pgcopydb.sentinel replay_lsn to %X/%X",
+				  LSN_FORMAT_ARGS(context->replay_lsn));
+		return false;
+	}
+
+	log_notice("Replayed %lld messages up to replay_lsn %X/%X",
+			   (long long) readerContext.lineno,
+			   LSN_FORMAT_ARGS(context->replay_lsn));
 
 	return true;
 }
@@ -156,7 +166,30 @@ stream_replay_line(void *ctx, const char *line, bool *stop)
 		case STREAM_ACTION_COMMIT:
 		case STREAM_ACTION_KEEPALIVE:
 		{
-			(void) stream_apply_sync_sentinel(context);
+			uint64_t now = time(NULL);
+
+			if (context->sentinelQueryInProgress)
+			{
+				if (!stream_apply_fetch_sync_sentinel(context))
+				{
+					/* errors have already been logged */
+					return false;
+				}
+			}
+
+			/* rate limit to 1 update per second */
+			else if (1 < (now - context->sentinelSyncTime))
+			{
+				/* we're going to keep the connection around */
+				context->src.connectionStatementType =
+					PGSQL_CONNECTION_MULTI_STATEMENT;
+
+				if (!stream_apply_send_sync_sentinel(context))
+				{
+					/* errors have already been logged */
+					return false;
+				}
+			}
 			break;
 		}
 

diff --git a/src/bin/pgcopydb/ld_stream.h b/src/bin/pgcopydb/ld_stream.h
@@ -121,7 +121,14 @@ typedef struct StreamApplyContext
 {
 	CDCPaths paths;
 
+	/* target connection */
 	PGSQL pgsql;
+
+	/* source connection to publish sentinel updates */
+	PGSQL src;
+	bool sentinelQueryInProgress;
+	uint64_t sentinelSyncTime;
+
 	char source_pguri[MAXCONNINFO];
 	char target_pguri[MAXCONNINFO];
 	char origin[BUFSIZE];
@@ -531,6 +538,8 @@ bool stream_apply_wait_for_sentinel(StreamSpecs *specs,
 									StreamApplyContext *context);
 
 bool stream_apply_sync_sentinel(StreamApplyContext *context);
+bool stream_apply_send_sync_sentinel(StreamApplyContext *context);
+bool stream_apply_fetch_sync_sentinel(StreamApplyContext *context);
 
 bool stream_apply_file(StreamApplyContext *context);