CDC implementation for Citus using Logical Replication

.circleci/config.yml

@@ -841,6 +841,11 @@ workflows:
           name: 'test-15_check-failure'
           image: citus/failtester
           make: check-failure
+      - test-citus:
+          <<:  *test-citus-15
+          name: 'test-15_check-cdc'
+          make: check-cdc
+
 
       - tap-test-citus: &tap-test-citus-13
           name: 'test-13_tap-recovery'
@@ -978,6 +983,7 @@ workflows:
             - test-13-14_check-pg-upgrade
             - test-14-15_check-pg-upgrade
             - test-13_check-citus-upgrade
+            - test-15_check-cdc
 
       - ch_benchmark:
           requires: [build-13]

src/backend/distributed/commands/alter_table.c

@@ -55,6 +55,7 @@
 #include "distributed/multi_partitioning_utils.h"
 #include "distributed/reference_table_utils.h"
 #include "distributed/relation_access_tracking.h"
+#include "distributed/replication_origin_session_utils.h"
 #include "distributed/shared_library_init.h"
 #include "distributed/shard_utils.h"
 #include "distributed/worker_protocol.h"
@@ -402,7 +403,11 @@ UndistributeTable(TableConversionParameters *params)
 	params->conversionType = UNDISTRIBUTE_TABLE;
 	params->shardCountIsNull = true;
 	TableConversionState *con = CreateTableConversion(params);
-	return ConvertTable(con);
+
+	SetupReplicationOriginLocalSession();
+	TableConversionReturn *conv = ConvertTable(con);
+	ResetReplicationOriginLocalSession();
+	return conv;
 }
 
 

src/backend/distributed/commands/create_distributed_table.c

@@ -2219,12 +2219,12 @@ CopyLocalDataIntoShards(Oid distributedRelationId)
 	EState *estate = CreateExecutorState();
 	ExprContext *econtext = GetPerTupleExprContext(estate);
 	econtext->ecxt_scantuple = slot;
-
+	const bool nonPublishableData = false;
 	DestReceiver *copyDest =
 		(DestReceiver *) CreateCitusCopyDestReceiver(distributedRelationId,
 													 columnNameList,
 													 partitionColumnIndex,
-													 estate, NULL);
+													 estate, NULL, nonPublishableData);
 
 	/* initialise state for writing to shards, we'll open connections on demand */
 	copyDest->rStartup(copyDest, 0, tupleDescriptor);

src/backend/distributed/commands/local_multi_copy.c

@@ -36,6 +36,7 @@
 #include "distributed/local_multi_copy.h"
 #include "distributed/shard_utils.h"
 #include "distributed/version_compat.h"
+#include "distributed/replication_origin_session_utils.h"
 
 /* managed via GUC, default is 512 kB */
 int LocalCopyFlushThresholdByte = 512 * 1024;
@@ -46,7 +47,7 @@ static void AddSlotToBuffer(TupleTableSlot *slot, CitusCopyDestReceiver *copyDes
 static bool ShouldAddBinaryHeaders(StringInfo buffer, bool isBinary);
 static bool ShouldSendCopyNow(StringInfo buffer);
 static void DoLocalCopy(StringInfo buffer, Oid relationId, int64 shardId,
-						CopyStmt *copyStatement, bool isEndOfCopy);
+						CopyStmt *copyStatement, bool isEndOfCopy, bool isPublishable);
 static int ReadFromLocalBufferCallback(void *outBuf, int minRead, int maxRead);
 
 
@@ -94,7 +95,7 @@ WriteTupleToLocalShard(TupleTableSlot *slot, CitusCopyDestReceiver *copyDest, in
 		bool isEndOfCopy = false;
 		DoLocalCopy(localCopyOutState->fe_msgbuf, copyDest->distributedRelationId,
 					shardId,
-					copyDest->copyStatement, isEndOfCopy);
+					copyDest->copyStatement, isEndOfCopy, copyDest->isPublishable);
 		resetStringInfo(localCopyOutState->fe_msgbuf);
 	}
 }
@@ -133,7 +134,7 @@ FinishLocalCopyToShard(CitusCopyDestReceiver *copyDest, int64 shardId,
 	}
 	bool isEndOfCopy = true;
 	DoLocalCopy(localCopyOutState->fe_msgbuf, copyDest->distributedRelationId, shardId,
-				copyDest->copyStatement, isEndOfCopy);
+				copyDest->copyStatement, isEndOfCopy, copyDest->isPublishable);
 }
 
 
@@ -197,14 +198,18 @@ ShouldSendCopyNow(StringInfo buffer)
  */
 static void
 DoLocalCopy(StringInfo buffer, Oid relationId, int64 shardId, CopyStmt *copyStatement,
-			bool isEndOfCopy)
+			bool isEndOfCopy, bool isPublishable)
 {
 	/*
 	 * Set the buffer as a global variable to allow ReadFromLocalBufferCallback
 	 * to read from it. We cannot pass additional arguments to
 	 * ReadFromLocalBufferCallback.
 	 */
 	LocalCopyBuffer = buffer;
+	if (!isPublishable)
+	{
+		SetupReplicationOriginLocalSession();
+	}
 
 	Oid shardOid = GetTableLocalShardOid(relationId, shardId);
 	Relation shard = table_open(shardOid, RowExclusiveLock);
@@ -219,6 +224,10 @@ DoLocalCopy(StringInfo buffer, Oid relationId, int64 shardId, CopyStmt *copyStat
 	EndCopyFrom(cstate);
 
 	table_close(shard, NoLock);
+	if (!isPublishable)
+	{
+		ResetReplicationOriginLocalSession();
+	}
 	free_parsestate(pState);
 }
 

src/backend/distributed/commands/multi_copy.c

@@ -85,6 +85,7 @@
 #include "distributed/relation_access_tracking.h"
 #include "distributed/remote_commands.h"
 #include "distributed/remote_transaction.h"
+#include "distributed/replication_origin_session_utils.h"
 #include "distributed/resource_lock.h"
 #include "distributed/shard_pruning.h"
 #include "distributed/shared_connection_stats.h"
@@ -270,7 +271,8 @@ static CopyConnectionState * GetConnectionState(HTAB *connectionStateHash,
 static CopyShardState * GetShardState(uint64 shardId, HTAB *shardStateHash,
 									  HTAB *connectionStateHash,
 									  bool *found, bool shouldUseLocalCopy, CopyOutState
-									  copyOutState, bool isColocatedIntermediateResult);
+									  copyOutState, bool isColocatedIntermediateResult,
+									  bool isPublishable);
 static MultiConnection * CopyGetPlacementConnection(HTAB *connectionStateHash,
 													ShardPlacement *placement,
 													bool colocatedIntermediateResult);
@@ -285,7 +287,8 @@ static void InitializeCopyShardState(CopyShardState *shardState,
 									 uint64 shardId,
 									 bool canUseLocalCopy,
 									 CopyOutState copyOutState,
-									 bool colocatedIntermediateResult);
+									 bool colocatedIntermediateResult, bool
+									 isPublishable);
 static void StartPlacementStateCopyCommand(CopyPlacementState *placementState,
 										   CopyStmt *copyStatement,
 										   CopyOutState copyOutState);
@@ -492,9 +495,11 @@ CopyToExistingShards(CopyStmt *copyStatement, QueryCompletion *completionTag)
 	ExprContext *executorExpressionContext = GetPerTupleExprContext(executorState);
 
 	/* set up the destination for the COPY */
+	const bool publishableData = true;
 	CitusCopyDestReceiver *copyDest = CreateCitusCopyDestReceiver(tableId, columnNameList,
 																  partitionColumnIndex,
-																  executorState, NULL);
+																  executorState, NULL,
+																  publishableData);
 
 	/* if the user specified an explicit append-to_shard option, write to it */
 	uint64 appendShardId = ProcessAppendToShardOption(tableId, copyStatement);
@@ -1934,7 +1939,7 @@ CopyFlushOutput(CopyOutState cstate, char *start, char *pointer)
 CitusCopyDestReceiver *
 CreateCitusCopyDestReceiver(Oid tableId, List *columnNameList, int partitionColumnIndex,
 							EState *executorState,
-							char *intermediateResultIdPrefix)
+							char *intermediateResultIdPrefix, bool isPublishable)
 {
 	CitusCopyDestReceiver *copyDest = (CitusCopyDestReceiver *) palloc0(
 		sizeof(CitusCopyDestReceiver));
@@ -1953,6 +1958,7 @@ CreateCitusCopyDestReceiver(Oid tableId, List *columnNameList, int partitionColu
 	copyDest->executorState = executorState;
 	copyDest->colocatedIntermediateResultIdPrefix = intermediateResultIdPrefix;
 	copyDest->memoryContext = CurrentMemoryContext;
+	copyDest->isPublishable = isPublishable;
 
 	return copyDest;
 }
@@ -2318,7 +2324,9 @@ CitusSendTupleToPlacements(TupleTableSlot *slot, CitusCopyDestReceiver *copyDest
 											   &cachedShardStateFound,
 											   copyDest->shouldUseLocalCopy,
 											   copyDest->copyOutState,
-											   isColocatedIntermediateResult);
+											   isColocatedIntermediateResult,
+											   copyDest->isPublishable);
+
 	if (!cachedShardStateFound)
 	{
 		firstTupleInShard = true;
@@ -2751,6 +2759,11 @@ ShutdownCopyConnectionState(CopyConnectionState *connectionState,
 	if (activePlacementState != NULL)
 	{
 		EndPlacementStateCopyCommand(activePlacementState, copyOutState);
+		if (!copyDest->isPublishable)
+		{
+			ResetReplicationOriginRemoteSession(
+				activePlacementState->connectionState->connection);
+		}
 	}
 
 	dlist_foreach(iter, &connectionState->bufferedPlacementList)
@@ -2764,6 +2777,10 @@ ShutdownCopyConnectionState(CopyConnectionState *connectionState,
 		SendCopyDataToPlacement(placementState->data, shardId,
 								connectionState->connection);
 		EndPlacementStateCopyCommand(placementState, copyOutState);
+		if (!copyDest->isPublishable)
+		{
+			ResetReplicationOriginRemoteSession(connectionState->connection);
+		}
 	}
 }
 
@@ -3436,15 +3453,16 @@ static CopyShardState *
 GetShardState(uint64 shardId, HTAB *shardStateHash,
 			  HTAB *connectionStateHash, bool *found, bool
 			  shouldUseLocalCopy, CopyOutState copyOutState,
-			  bool isColocatedIntermediateResult)
+			  bool isColocatedIntermediateResult, bool isPublishable)
 {
 	CopyShardState *shardState = (CopyShardState *) hash_search(shardStateHash, &shardId,
 																HASH_ENTER, found);
 	if (!*found)
 	{
 		InitializeCopyShardState(shardState, connectionStateHash,
 								 shardId, shouldUseLocalCopy,
-								 copyOutState, isColocatedIntermediateResult);
+								 copyOutState, isColocatedIntermediateResult,
+								 isPublishable);
 	}
 
 	return shardState;
@@ -3461,7 +3479,8 @@ InitializeCopyShardState(CopyShardState *shardState,
 						 HTAB *connectionStateHash, uint64 shardId,
 						 bool shouldUseLocalCopy,
 						 CopyOutState copyOutState,
-						 bool colocatedIntermediateResult)
+						 bool colocatedIntermediateResult,
+						 bool isPublishable)
 {
 	ListCell *placementCell = NULL;
 	int failedPlacementCount = 0;
@@ -3532,6 +3551,12 @@ InitializeCopyShardState(CopyShardState *shardState,
 			RemoteTransactionBeginIfNecessary(connection);
 		}
 
+		if (!isPublishable)
+		{
+			/*elog(LOG,"InitializeCopyShardState: calling SetupReplicationOriginRemoteSession conn id: %lu", connection->connectionId); */
+			SetupReplicationOriginRemoteSession(connection);
+		}
+
 		CopyPlacementState *placementState = palloc0(sizeof(CopyPlacementState));
 		placementState->shardState = shardState;
 		placementState->data = makeStringInfo();

src/backend/distributed/connection/connection_management.c

@@ -1481,6 +1481,7 @@ AfterXactHostConnectionHandling(ConnectionHashEntry *entry, bool isCommit)
  * - Current cached connections is already at MaxCachedConnectionsPerWorker
  * - Connection is forced to close at the end of transaction
  * - Connection is not in OK state
+ * - Connection has a replication origin setup
  * - A transaction is still in progress (usually because we are cancelling a distributed transaction)
  * - A connection reached its maximum lifetime
  */
@@ -1500,6 +1501,7 @@ ShouldShutdownConnection(MultiConnection *connection, const int cachedConnection
 		   PQstatus(connection->pgConn) != CONNECTION_OK ||
 		   !RemoteTransactionIdle(connection) ||
 		   connection->requiresReplication ||
+		   connection->isReplicationOriginSessionSetup ||
 		   (MaxCachedConnectionLifetime >= 0 &&
 			MillisecondsToTimeout(connection->connectionEstablishmentStart,
 								  MaxCachedConnectionLifetime) <= 0);

src/backend/distributed/executor/insert_select_executor.c

@@ -409,11 +409,13 @@ ExecutePlanIntoColocatedIntermediateResults(Oid targetRelationId,
 																  columnNameList);
 
 	/* set up a DestReceiver that copies into the intermediate table */
+	const bool publishableData = true;
 	CitusCopyDestReceiver *copyDest = CreateCitusCopyDestReceiver(targetRelationId,
 																  columnNameList,
 																  partitionColumnIndex,
 																  executorState,
-																  intermediateResultIdPrefix);
+																  intermediateResultIdPrefix,
+																  publishableData);
 
 	ExecutePlanIntoDestReceiver(selectPlan, paramListInfo, (DestReceiver *) copyDest);
 
@@ -443,10 +445,12 @@ ExecutePlanIntoRelation(Oid targetRelationId, List *insertTargetList,
 																  columnNameList);
 
 	/* set up a DestReceiver that copies into the distributed table */
+	const bool publishableData = true;
 	CitusCopyDestReceiver *copyDest = CreateCitusCopyDestReceiver(targetRelationId,
 																  columnNameList,
 																  partitionColumnIndex,
-																  executorState, NULL);
+																  executorState, NULL,
+																  publishableData);
 
 	ExecutePlanIntoDestReceiver(selectPlan, paramListInfo, (DestReceiver *) copyDest);
 

src/backend/distributed/operations/worker_shard_copy.c

@@ -24,6 +24,7 @@
 #include "distributed/relation_utils.h"
 #include "distributed/version_compat.h"
 #include "distributed/local_executor.h"
+#include "distributed/replication_origin_session_utils.h"
 
 /*
  * LocalCopyBuffer is used in copy callback to return the copied rows.
@@ -80,6 +81,7 @@ static void LocalCopyToShard(ShardCopyDestReceiver *copyDest, CopyOutState
 							 localCopyOutState);
 static void ConnectToRemoteAndStartCopy(ShardCopyDestReceiver *copyDest);
 
+
 static bool
 CanUseLocalCopy(uint32_t destinationNodeId)
 {
@@ -103,6 +105,12 @@ ConnectToRemoteAndStartCopy(ShardCopyDestReceiver *copyDest)
 														 NULL /* database (current) */);
 	ClaimConnectionExclusively(copyDest->connection);
 
+
+	RemoteTransactionBeginIfNecessary(copyDest->connection);
+
+	SetupReplicationOriginRemoteSession(copyDest->connection);
+
+
 	StringInfo copyStatement = ConstructShardCopyStatement(
 		copyDest->destinationShardFullyQualifiedName,
 		copyDest->copyOutState->binary);
@@ -184,6 +192,8 @@ ShardCopyDestReceiverReceive(TupleTableSlot *slot, DestReceiver *dest)
 	CopyOutState copyOutState = copyDest->copyOutState;
 	if (copyDest->useLocalCopy)
 	{
+		/* Setup replication origin session for local copy*/
+
 		WriteLocalTuple(slot, copyDest);
 		if (copyOutState->fe_msgbuf->len > LocalCopyFlushThresholdByte)
 		{
@@ -259,6 +269,11 @@ ShardCopyDestReceiverStartup(DestReceiver *dest, int operation, TupleDesc
 	copyDest->columnOutputFunctions = ColumnOutputFunctions(inputTupleDescriptor,
 															copyOutState->binary);
 	copyDest->copyOutState = copyOutState;
+	if (copyDest->useLocalCopy)
+	{
+		/* Setup replication origin session for local copy*/
+		SetupReplicationOriginLocalSession();
+	}
 }
 
 
@@ -317,6 +332,9 @@ ShardCopyDestReceiverShutdown(DestReceiver *dest)
 
 		PQclear(result);
 		ForgetResults(copyDest->connection);
+
+		ResetReplicationOriginRemoteSession(copyDest->connection);
+
 		CloseConnection(copyDest->connection);
 	}
 }
@@ -329,6 +347,10 @@ static void
 ShardCopyDestReceiverDestroy(DestReceiver *dest)
 {
 	ShardCopyDestReceiver *copyDest = (ShardCopyDestReceiver *) dest;
+	if (copyDest->useLocalCopy)
+	{
+		ResetReplicationOriginLocalSession();
+	}
 
 	if (copyDest->copyOutState)
 	{