From e3c7a50dd744caa9b29f6882d77cc3ba43955cd4 Mon Sep 17 00:00:00 2001
From: "houzj.fnst" <houzj.fnst@cn.fujitsu.com>
Date: Mon, 10 Oct 2022 11:55:58 +0800
Subject: [PATCH v38 4/6] Retrict parallel for partitioned table

Disallow replicating from or to a partitioned table in parallel streaming
mode. This is to avoid the deadlock cases when the partitioned table's
structure is different between publisher and subscriber.

---
 doc/src/sgml/protocol.sgml                    |  11 ++
 doc/src/sgml/ref/create_subscription.sgml     |   8 +-
 .../replication/logical/applyparallelworker.c |   7 --
 src/backend/replication/logical/proto.c       |  12 +-
 src/backend/replication/logical/relation.c    |  24 ++--
 src/backend/replication/logical/worker.c      |   2 +-
 src/backend/replication/pgoutput/pgoutput.c   |   4 +-
 src/include/replication/logicalproto.h        |   6 +-
 .../t/032_streaming_parallel_safety.pl        | 113 ++++++++++++++++--
 9 files changed, 158 insertions(+), 29 deletions(-)

diff --git a/doc/src/sgml/protocol.sgml b/doc/src/sgml/protocol.sgml
index c6a0e80b3f..0566153e62 100644
--- a/doc/src/sgml/protocol.sgml
+++ b/doc/src/sgml/protocol.sgml
@@ -6352,6 +6352,17 @@ psql "dbname=postgres replication=database" -c "IDENTIFY_SYSTEM;"
        </listitem>
       </varlistentry>
 
+      <varlistentry>
+       <term>Int8</term>
+       <listitem>
+        <para>
+         The relation kind (same as
+         <structfield>relkind</structfield> in <structname>pg_class</structname>).
+         This field is available since protocol version 4.
+        </para>
+       </listitem>
+      </varlistentry>
+
       <varlistentry>
        <term>Int16</term>
        <listitem>
diff --git a/doc/src/sgml/ref/create_subscription.sgml b/doc/src/sgml/ref/create_subscription.sgml
index 79250acaf5..6a8c839dca 100644
--- a/doc/src/sgml/ref/create_subscription.sgml
+++ b/doc/src/sgml/ref/create_subscription.sgml
@@ -245,11 +245,15 @@ CREATE SUBSCRIPTION <replaceable class="parameter">subscription_name</replaceabl
           transaction is committed. Note that if an error happens in a
           parallel apply worker, the finish LSN of the remote transaction
           might not be reported in the server log.
-          There are two prerequisites for using <literal>parallel</literal>
+          There are some prerequisites for using <literal>parallel</literal>
           mode: 1) the unique column in the table on the subscriber-side must
           also be the unique column on the publisher-side; 2) there cannot be
           any non-immutable functions used by the subscriber-side replicated
-          table.
+          table; 3) Changes on a partitioned table (or on its partitions)
+          contained in the publication cannot be published using the identity
+          and schema of the partitioned table. And the non-partitioned table on
+          the publisher-side cannot be a partitioned table on the
+          subscriber-side.
          </para>
         </listitem>
        </varlistentry>
diff --git a/src/backend/replication/logical/applyparallelworker.c b/src/backend/replication/logical/applyparallelworker.c
index 88473aa664..921b6c52ab 100644
--- a/src/backend/replication/logical/applyparallelworker.c
+++ b/src/backend/replication/logical/applyparallelworker.c
@@ -1105,13 +1105,6 @@ parallel_apply_relation_check(LogicalRepRelMapEntry *rel)
 	if (!MyLogicalRepWorker->parallel_apply)
 		return;
 
-	/*
-	 * Partition table checks are done later in function
-	 * apply_handle_tuple_routing.
-	 */
-	if (rel->localrel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
-		return;
-
 	if (!rel->parallel_apply_valid)
 		logicalrep_rel_mark_parallel_apply(rel);
 
diff --git a/src/backend/replication/logical/proto.c b/src/backend/replication/logical/proto.c
index bd406a075f..cdda5906f9 100644
--- a/src/backend/replication/logical/proto.c
+++ b/src/backend/replication/logical/proto.c
@@ -728,7 +728,7 @@ logicalrep_write_message(StringInfo out, TransactionId xid, XLogRecPtr lsn,
  */
 void
 logicalrep_write_rel(StringInfo out, TransactionId xid, Relation rel,
-					 Bitmapset *columns)
+					 Bitmapset *columns, bool write_relkind)
 {
 	char	   *relname;
 
@@ -749,6 +749,10 @@ logicalrep_write_rel(StringInfo out, TransactionId xid, Relation rel,
 	/* send replica identity */
 	pq_sendbyte(out, rel->rd_rel->relreplident);
 
+	/* Send relation kind if needed. */
+	if (write_relkind)
+		pq_sendbyte(out, rel->rd_rel->relkind);
+
 	/* send the attribute info */
 	logicalrep_write_attrs(out, rel, columns);
 }
@@ -757,7 +761,7 @@ logicalrep_write_rel(StringInfo out, TransactionId xid, Relation rel,
  * Read the relation info from stream and return as LogicalRepRelation.
  */
 LogicalRepRelation *
-logicalrep_read_rel(StringInfo in)
+logicalrep_read_rel(StringInfo in, bool read_relkind)
 {
 	LogicalRepRelation *rel = palloc(sizeof(LogicalRepRelation));
 
@@ -770,6 +774,10 @@ logicalrep_read_rel(StringInfo in)
 	/* Read the replica identity. */
 	rel->replident = pq_getmsgbyte(in);
 
+	/* Read the relation kind if needed. */
+	if (read_relkind)
+		rel->relkind = pq_getmsgbyte(in);
+
 	/* Get attribute description */
 	logicalrep_read_attrs(in, rel);
 
diff --git a/src/backend/replication/logical/relation.c b/src/backend/replication/logical/relation.c
index 695244110b..e5efdb3bd4 100644
--- a/src/backend/replication/logical/relation.c
+++ b/src/backend/replication/logical/relation.c
@@ -224,6 +224,7 @@ logicalrep_relmap_update(LogicalRepRelation *remoterel)
 		entry->remoterel.atttyps[i] = remoterel->atttyps[i];
 	}
 	entry->remoterel.replident = remoterel->replident;
+	entry->remoterel.relkind = remoterel->relkind;
 	entry->remoterel.attkeys = bms_copy(remoterel->attkeys);
 	entry->remoterel.attunique = bms_copy(remoterel->attunique);
 	MemoryContextSwitchTo(oldctx);
@@ -351,10 +352,11 @@ logicalrep_rel_mark_updatable(LogicalRepRelMapEntry *entry)
  * worker and assign the 'parallel_apply_valid' flag and
  * 'parallel_apply_details' field.
  *
- * There are two requirements for applying changes using a parallel apply
- * worker: 1) the unique column in the table on the subscriber-side should also
- * be the unique column on the publisher-side; 2) there cannot be any
- * non-immutable functions used by the subscriber-side replicated table.
+ * There are some requirements for applying changes using a parallel apply
+ * worker: 1) replicate from or into a partitioned table is not allowed; 2) the
+ * unique column in the table on the subscriber-side should also be the unique
+ * column on the publisher-side; 3) there cannot be any non-immutable functions
+ * used by the subscriber-side replicated table.
  *
  * Without these safety checks, the following scenario may occur: The parallel
  * apply worker locks a row when processing a streaming transaction, after
@@ -385,9 +387,17 @@ logicalrep_rel_mark_parallel_apply(LogicalRepRelMapEntry *entry)
 
 	entry->parallel_apply_valid = true;
 
+	/* 1) Check if replicate from or into a partitioned table. */
+	if (entry->remoterel.relkind == RELKIND_PARTITIONED_TABLE ||
+		entry->localrel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+	{
+		entry->parallel_apply_details = _("Replicate from or into a partitioned table is not allowed.");
+		return;
+	}
+
 	/*
-	 * First, check if the unique column in the relation on the
-	 * subscriber-side is also the unique column on the publisher-side.
+	 * 2) Check if the unique column in the relation on the subscriber-side
+	 * is also the unique column on the publisher-side.
 	 */
 	ukey = RelationGetIndexAttrBitmap(entry->localrel,
 									  INDEX_ATTR_BITMAP_KEY);
@@ -412,7 +422,7 @@ logicalrep_rel_mark_parallel_apply(LogicalRepRelMapEntry *entry)
 	}
 
 	/*
-	 * Then, check if there is any non-immutable function used by the
+	 * 3) Check if there is any non-immutable function used by the
 	 * subscriber-side relation. Look for functions in the following places:
 	 * a. trigger functions; b. Column default value expressions and domain
 	 * constraints; c. Constraint expressions;
diff --git a/src/backend/replication/logical/worker.c b/src/backend/replication/logical/worker.c
index 1e8163f26c..2c4c41aa97 100644
--- a/src/backend/replication/logical/worker.c
+++ b/src/backend/replication/logical/worker.c
@@ -2021,7 +2021,7 @@ apply_handle_relation(StringInfo s)
 	if (handle_streamed_transaction(LOGICAL_REP_MSG_RELATION, s))
 		return;
 
-	rel = logicalrep_read_rel(s);
+	rel = logicalrep_read_rel(s, MyLogicalRepWorker->parallel_apply);
 	logicalrep_relmap_update(rel);
 
 	/* Also reset all entries in the partition map that refer to remoterel. */
diff --git a/src/backend/replication/pgoutput/pgoutput.c b/src/backend/replication/pgoutput/pgoutput.c
index c3accf28e4..0d7abbeb1c 100644
--- a/src/backend/replication/pgoutput/pgoutput.c
+++ b/src/backend/replication/pgoutput/pgoutput.c
@@ -741,6 +741,8 @@ send_relation_and_attrs(Relation relation, TransactionId xid,
 						LogicalDecodingContext *ctx,
 						Bitmapset *columns)
 {
+	PGOutputData *data = (PGOutputData *) ctx->output_plugin_private;
+	bool		write_relkind = (data->streaming == SUBSTREAM_PARALLEL);
 	TupleDesc	desc = RelationGetDescr(relation);
 	int			i;
 
@@ -772,7 +774,7 @@ send_relation_and_attrs(Relation relation, TransactionId xid,
 	}
 
 	OutputPluginPrepareWrite(ctx, false);
-	logicalrep_write_rel(ctx->out, xid, relation, columns);
+	logicalrep_write_rel(ctx->out, xid, relation, columns, write_relkind);
 	OutputPluginWrite(ctx, false);
 }
 
diff --git a/src/include/replication/logicalproto.h b/src/include/replication/logicalproto.h
index d4bc92c5e2..eb12614335 100644
--- a/src/include/replication/logicalproto.h
+++ b/src/include/replication/logicalproto.h
@@ -248,8 +248,10 @@ extern List *logicalrep_read_truncate(StringInfo in,
 extern void logicalrep_write_message(StringInfo out, TransactionId xid, XLogRecPtr lsn,
 									 bool transactional, const char *prefix, Size sz, const char *message);
 extern void logicalrep_write_rel(StringInfo out, TransactionId xid,
-								 Relation rel, Bitmapset *columns);
-extern LogicalRepRelation *logicalrep_read_rel(StringInfo in);
+								 Relation rel, Bitmapset *columns,
+								 bool write_relkind);
+extern LogicalRepRelation *logicalrep_read_rel(StringInfo in,
+											   bool read_relkind);
 extern void logicalrep_write_typ(StringInfo out, TransactionId xid,
 								 Oid typoid);
 extern void logicalrep_read_typ(StringInfo in, LogicalRepTyp *ltyp);
diff --git a/src/test/subscription/t/032_streaming_parallel_safety.pl b/src/test/subscription/t/032_streaming_parallel_safety.pl
index 3aaa699c07..2195070434 100644
--- a/src/test/subscription/t/032_streaming_parallel_safety.pl
+++ b/src/test/subscription/t/032_streaming_parallel_safety.pl
@@ -34,7 +34,20 @@ $node_subscriber->start;
 $node_publisher->safe_psql('postgres', "CREATE TABLE test_tab1 (a int)");
 $node_publisher->safe_psql('postgres', "CREATE TABLE test_tab2 (a int)");
 $node_publisher->safe_psql('postgres',
-	"CREATE TABLE test_tab_partitioned (a int primary key, b varchar)");
+	"CREATE TABLE test_tab_partitioned (a int primary key, b varchar) PARTITION BY RANGE(a)"
+);
+$node_publisher->safe_psql('postgres',
+	"CREATE TABLE test_tab_partition (LIKE test_tab_partitioned)");
+$node_publisher->safe_psql('postgres',
+	"ALTER TABLE test_tab_partitioned ATTACH PARTITION test_tab_partition DEFAULT"
+);
+
+$node_publisher->safe_psql('postgres', "CREATE TABLE test_tab_schema_parent (a int)");
+$node_publisher->safe_psql('postgres', "CREATE TABLE test_tab_schema_child1 (a int)");
+
+$node_publisher->safe_psql('postgres', "CREATE TABLE test_tab_viaroot_parent (a int) PARTITION BY RANGE (a)");
+$node_publisher->safe_psql('postgres', "CREATE TABLE test_tab_viaroot_child1 PARTITION OF test_tab_viaroot_parent DEFAULT");
+$node_publisher->safe_psql('postgres', "CREATE TABLE test_tab_viaroot_child2 (a int)");
 
 # Setup structure on subscriber
 # We need to test normal table and partition table.
@@ -49,19 +62,33 @@ $node_subscriber->safe_psql('postgres',
 	"ALTER TABLE test_tab_partitioned ATTACH PARTITION test_tab_partition DEFAULT"
 );
 
+$node_subscriber->safe_psql('postgres', "CREATE TABLE test_tab_schema_parent (a int) PARTITION BY RANGE (a)");
+$node_subscriber->safe_psql('postgres', "CREATE TABLE test_tab_schema_child1 PARTITION OF test_tab_schema_parent DEFAULT");
+
+$node_subscriber->safe_psql('postgres', "CREATE TABLE test_tab_viaroot_parent (a int) PARTITION BY RANGE (a)");
+$node_subscriber->safe_psql('postgres', "CREATE TABLE test_tab_viaroot_child1 PARTITION OF test_tab_viaroot_parent FOR VALUES FROM (MINVALUE) TO (2);");
+$node_subscriber->safe_psql('postgres', "CREATE TABLE test_tab_viaroot_child2 PARTITION OF test_tab_viaroot_parent FOR VALUES FROM (2) TO (MAXVALUE)");
+
 # Setup logical replication
 my $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
 $node_publisher->safe_psql('postgres',
 	"CREATE PUBLICATION tap_pub_normal FOR TABLE test_tab1, test_tab2");
 $node_publisher->safe_psql('postgres',
 	"CREATE PUBLICATION tap_pub_partitioned FOR TABLE test_tab_partitioned");
+$node_publisher->safe_psql('postgres',
+	"CREATE PUBLICATION test_pub_partitioned_schema FOR TABLE test_tab_schema_parent, test_tab_schema_child1;"
+);
+$node_publisher->safe_psql('postgres',
+	"CREATE PUBLICATION test_pub_partitioned_viaroot FOR TABLE test_tab_viaroot_parent, test_tab_viaroot_child1, test_tab_viaroot_child2 WITH (PUBLISH_VIA_PARTITION_ROOT);"
+);
 
 my $appname = 'tap_sub';
 $node_subscriber->safe_psql(
 	'postgres', "
 	CREATE SUBSCRIPTION tap_sub
 	CONNECTION '$publisher_connstr application_name=$appname'
-	PUBLICATION tap_pub_normal, tap_pub_partitioned
+	PUBLICATION tap_pub_normal, tap_pub_partitioned,
+	test_pub_partitioned_schema, test_pub_partitioned_viaroot
 	WITH (streaming = parallel, copy_data = false)");
 
 $node_publisher->wait_for_catchup($appname);
@@ -129,7 +156,7 @@ $node_publisher->safe_psql('postgres',
 );
 
 $node_subscriber->wait_for_log(
-	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partitioned" using subscription parameter streaming = parallel/,
+	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partition" using subscription parameter streaming = parallel/,
 	$offset);
 
 # Drop the unique index on the subscriber, now it works.
@@ -253,7 +280,7 @@ $offset = -s $node_subscriber->logfile;
 $node_publisher->safe_psql('postgres', "DELETE FROM test_tab_partitioned");
 
 $node_subscriber->wait_for_log(
-	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partitioned" using subscription parameter streaming = parallel/,
+	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partition" using subscription parameter streaming = parallel/,
 	$offset);
 
 # Drop the trigger on the subscriber, now it works.
@@ -343,7 +370,7 @@ $node_publisher->safe_psql('postgres',
 );
 
 $node_subscriber->wait_for_log(
-	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partitioned" using subscription parameter streaming = parallel/,
+	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partition" using subscription parameter streaming = parallel/,
 	$offset);
 
 # Drop default value on the subscriber, now it works.
@@ -498,7 +525,7 @@ $offset = -s $node_subscriber->logfile;
 $node_publisher->safe_psql('postgres', "DELETE FROM test_tab_partitioned");
 
 $node_subscriber->wait_for_log(
-	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partitioned" using subscription parameter streaming = parallel/,
+	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partition" using subscription parameter streaming = parallel/,
 	$offset);
 
 # Drop constraint on the subscriber, now it works.
@@ -594,7 +621,7 @@ $node_publisher->safe_psql('postgres',
 );
 
 $node_subscriber->wait_for_log(
-	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partitioned" using subscription parameter streaming = parallel/,
+	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_partition" using subscription parameter streaming = parallel/,
 	$offset);
 
 # Drop the foreign key constraint on the subscriber, now it works.
@@ -610,6 +637,78 @@ $result =
 is($result, qq(5000),
 	'data replicated to subscriber after dropping the foreign key');
 
+# ============================================================================
+# It is not allowed that the non-partitioned table on the publisher-side is a
+# partitioned table on the subscriber-side. Check the error reported by
+# parallel worker in this case.
+# ============================================================================
+
+# Check the subscriber log from now on.
+$offset = -s $node_subscriber->logfile;
+
+$in .= q{
+BEGIN;
+INSERT INTO test_tab_schema_parent SELECT i FROM generate_series(1, 5000) s(i);
+};
+$h->pump_nb;
+
+$node_publisher->safe_psql('postgres', "TRUNCATE test_tab_schema_child1");
+
+$in .= q{
+COMMIT;
+\q
+};
+$h->finish;
+
+$node_subscriber->wait_for_log(
+	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_schema_parent" using subscription parameter streaming = parallel/,
+	$offset);
+
+# Change the streaming option from "parallel" to "on", now it works.
+$node_subscriber->safe_psql('postgres',
+	"ALTER SUBSCRIPTION tap_sub SET (streaming = true)"
+);
+
+# Wait for this streaming transaction to be applied in the apply worker.
+$node_publisher->wait_for_catchup($appname);
+
+$result =
+  $node_subscriber->safe_psql('postgres', "SELECT count(*) FROM test_tab_schema_parent");
+is($result, qq(5000),
+	'data replicated to subscriber after unifying the table schema');
+
+# Reset the streaming option.
+$node_subscriber->safe_psql('postgres',
+	"ALTER SUBSCRIPTION tap_sub SET (streaming = parallel)"
+);
+
+# ============================================================================
+# It is not allowed that the partitioned table on the publisher-side is
+# published with specifying publish_via_partition_root. Check the error
+# reported by parallel worker in this case.
+# ============================================================================
+
+# Check the subscriber log from now on.
+$offset = -s $node_subscriber->logfile;
+
+$in .= q{
+BEGIN;
+INSERT INTO test_tab_viaroot_parent SELECT i FROM generate_series(1, 5000) s(i);
+};
+$h->pump_nb;
+
+$node_publisher->safe_psql('postgres', "TRUNCATE test_tab_viaroot_child2");
+
+$in .= q{
+COMMIT;
+\q
+};
+$h->finish;
+
+$node_subscriber->wait_for_log(
+	qr/ERROR: ( [A-Z0-9]+:)? cannot replicate target relation "public.test_tab_viaroot_parent" using subscription parameter streaming = parallel/,
+	$offset);
+
 $node_subscriber->stop;
 $node_publisher->stop;
 
-- 
2.23.0.windows.1

