From ea7e504d9ac90ba0dfa8731f7b3465e07d19baf8 Mon Sep 17 00:00:00 2001
From: bovenshi <bovenshi@tencent.com>
Date: Fri, 25 Nov 2022 15:42:01 +0800
Subject: [PATCH] Fix snapshot builder's state after restoring snapshot from
 disk.

Set snapshot builder's state after restoring snapshot from disk, since consistent snapshots have no next phase.

Per reports from bowenshi.

Discussion: https://www.postgresql.org/message-id/17695-6be9277c9295985f%40postgresql.org
---
 src/backend/replication/logical/snapbuild.c      |   6 ++
 src/test/recovery/t/034_replication_snapbuild.pl | 100 +++++++++++++++++++++++
 2 files changed, 106 insertions(+)
 create mode 100644 src/test/recovery/t/034_replication_snapbuild.pl

diff --git a/src/backend/replication/logical/snapbuild.c b/src/backend/replication/logical/snapbuild.c
index b5b018e..03e7458 100644
--- a/src/backend/replication/logical/snapbuild.c
+++ b/src/backend/replication/logical/snapbuild.c
@@ -1964,6 +1964,12 @@ SnapBuildRestore(SnapBuild *builder, XLogRecPtr lsn)
 	/* consistent snapshots have no next phase */
 	Assert(ondisk.builder.next_phase_at == InvalidTransactionId);
 
+	/* 
+	 * We are now in consistent point and have no next phase, copy the state
+	 * from the snapshot in disk.
+	 */
+	builder->next_phase_at = ondisk.builder.next_phase_at;
+
 	/* ok, we think the snapshot is sensible, copy over everything important */
 	builder->xmin = ondisk.builder.xmin;
 	builder->xmax = ondisk.builder.xmax;
diff --git a/src/test/recovery/t/034_replication_snapbuild.pl b/src/test/recovery/t/034_replication_snapbuild.pl
new file mode 100644
index 0000000..9a40f25
--- /dev/null
+++ b/src/test/recovery/t/034_replication_snapbuild.pl
@@ -0,0 +1,100 @@
+# Copyright (c) 2021-2022, PostgreSQL Global Development Group
+
+# Test replication snapshot rebuild.
+use strict;
+use warnings;
+use TestLib;
+use PostgresNode;
+use Test::More;
+use DBI;
+use Test::More tests => 1;
+
+# set slot number
+my $slot_num = 100;
+
+# initialize master node
+my $master = get_new_node('master');
+$master->init;
+$master->append_conf(
+	'postgresql.conf', qq(
+wal_level=logical
+max_connections=1000
+max_worker_processes=500
+max_wal_senders=500
+max_replication_slots=500
+max_logical_replication_workers=500
+log_error_verbosity = verbose
+log_statement = all
+));
+$master->start;
+my $master_port = $master->port;
+
+# master init
+for (my $i = 0; $i<$slot_num; $i = $i+1)
+{
+	my $table_name=qq(t_$i);
+	my $pub_name=qq(pub_$i);
+	$master->safe_psql('postgres', "create table $table_name(a int);");
+	$master->safe_psql('postgres', "create publication $pub_name for table $table_name;");
+}
+
+# create subscription
+my $subscriber = get_new_node('subscriber');
+$subscriber->init;
+$subscriber->append_conf(
+	'postgresql.conf', qq(
+wal_level=logical
+max_connections=1000
+max_worker_processes=500
+max_replication_slots=500
+log_error_verbosity = verbose
+max_logical_replication_workers=500
+log_statement = all
+));
+$subscriber->start;
+
+# subscriber init
+for (my $i = 0; $i<$slot_num; $i = $i+1)
+{
+	my $table_name=qq(t_$i);
+	my $pub_name=qq(pub_$i);
+	my $sub_name=qq(sub_$i);
+
+	$subscriber->safe_psql('postgres',
+	"create table $table_name(a int);");
+}
+
+my $script = $master->basedir . '/pgbench_script';
+for (my $i = 0; $i<$slot_num; $i = $i+1)
+{
+	my $table_name=qq(t_$i);
+
+	append_to_file($script,
+	"INSERT INTO $table_name SELECT FROM generate_series(1,100);");
+}
+
+my $pid = fork();
+if ($pid == 0)
+{
+	sleep(1);
+	for (my $i = 0; $i<$slot_num; $i = $i+1)
+	{
+		my $pub_name=qq(pub_$i);
+		my $sub_name=qq(sub_$i);
+
+		$subscriber->safe_psql('postgres',
+		"create subscription $sub_name connection 'port=$master_port dbname=postgres user=postgres' publication $pub_name;");
+	}
+	exec "echo", "\nbye\n";
+}
+elsif ($pid < 0)
+{
+	die "failed to fork new process";
+}
+
+$master->run_log(
+	[qw(pgbench --no-vacuum --client=10 -U postgres
+	-T 25 --file), $script ]);
+
+my $result = $master->safe_psql('postgres', "select count(*) from pg_replication_slots;");
+is($result, qq($slot_num), "Create slots successfully.")
\ No newline at end of file
-- 
2.9.3