#!/bin/bash
# Test pg_rewind stuff

DATA_BASE=$HOME/data
MASTER_PORT=5432
STANDBY1_PORT=5433
STANDBY2_PORT=5434
STANDBY3_PORT=5435
ARCHIVE_DATA=$HOME/archive/5432

create_standby()
{
	ROOT_PORT=$1
	STANDBY_PORT=$2

	pg_basebackup -D $DATA_BASE/$STANDBY_PORT -p $ROOT_PORT > /dev/null 2>&1
	cat >> $DATA_BASE/$STANDBY_PORT/postgresql.conf <<EOF
port = $STANDBY_PORT
EOF
	cat >> $DATA_BASE/$STANDBY_PORT/recovery.conf <<EOF
primary_conninfo = 'port=$ROOT_PORT application_name=node_$STANDBY_PORT'
restore_command = 'cp -i $ARCHIVE_DATA/%f %p'
standby_mode = on
EOF
	pg_ctl start -s -w -D $DATA_BASE/$STANDBY_PORT
}

operate_all_nodes()
{
	ACTION=$1
	pg_ctl $ACTION -w -s -D $DATA_BASE/$STANDBY1_PORT # Timeline 1->2
	pg_ctl $ACTION -w -s -D $DATA_BASE/$STANDBY2_PORT # Timeline 1->3
	pg_ctl $ACTION -w -s -D $DATA_BASE/$STANDBY3_PORT # Timeline 1->2->4
	pg_ctl $ACTION -w -s -D $DATA_BASE/$MASTER_PORT # Timeline 1
	sleep 1
}

replug_node()
{
	SOURCE_PORT=$1
	TARGET_PORT=$2
	echo "Source = $SOURCE_PORT, Target = $TARGET_PORT"
	# Rewind...
	set -e
	pg_rewind --source-pgdata=$DATA_BASE/$SOURCE_PORT --target-pgdata=$DATA_BASE/$TARGET_PORT
	set +e

	# ... And replug
	cat >> $DATA_BASE/$TARGET_PORT/postgresql.conf <<EOF
port = $TARGET_PORT
EOF
	cat > $DATA_BASE/$TARGET_PORT/recovery.conf <<EOF
primary_conninfo = 'port=$SOURCE_PORT application_name=node_$TARGET_PORT'
restore_command = 'cp -i $ARCHIVE_DATA/%f %p'
standby_mode = on
EOF
	pg_ctl start -s -w -D $DATA_BASE/$SOURCE_PORT
	pg_ctl start -s -w -D $DATA_BASE/$TARGET_PORT
}

# Create cluster
rm -rf $DATA_BASE/$MASTER_PORT
rm -rf $DATA_BASE/$STANDBY1_PORT
rm -rf $DATA_BASE/$STANDBY2_PORT
rm -rf $DATA_BASE/$STANDBY3_PORT
rm -rf $ARCHIVE_DATA
mkdir -p $ARCHIVE_DATA
initdb -D $DATA_BASE/$MASTER_PORT > /dev/null 2>&1
# archive_command does not care if partial segments overlap, hence
# forcing it. partial segments aren't used by pg_rewind as well...
cat >> $DATA_BASE/$MASTER_PORT/postgresql.conf <<EOF
wal_level = hot_standby
wal_log_hints = on
log_directory = pg_log
logging_collector = on
max_wal_senders = 10
hot_standby = on
wal_keep_segments = 25
archive_mode = on
archive_command = 'cp -f %p $ARCHIVE_DATA/%f'
EOF
cat >> $DATA_BASE/$MASTER_PORT/pg_hba.conf <<EOF
host replication $USER 127.0.0.1/32 trust
host replication $USER ::1/128 trust
local replication $USER trust
EOF
pg_ctl start -s -w -D $DATA_BASE/$MASTER_PORT
createdb -p $MASTER_PORT $USER > /dev/null 2>&1

# Create cluster as follows:
# master (5432)
#  /        \
# 1 (5433)   2 (5434)
# |
# 3 (5435)
create_standby $MASTER_PORT $STANDBY1_PORT
create_standby $MASTER_PORT $STANDBY2_PORT
create_standby $STANDBY1_PORT $STANDBY3_PORT

# Create data on master node
psql -Atq -c 'CREATE TABLE aa (a int)'
psql -Atq -c 'INSERT INTO aa VALUES (1)'
sleep 1 # Let room to replay

# Time to promote and fork each standby
# Standby 1 selects timeline 2
pg_ctl promote -s -D $DATA_BASE/$STANDBY1_PORT -w
sleep 1
psql -Atq -c 'INSERT INTO aa VALUES (2)' -p $MASTER_PORT
psql -Atq -c 'INSERT INTO aa VALUES (3)' -p $STANDBY1_PORT
sleep 1

# Standby 2 selects timeline 3
pg_ctl promote -s -D $DATA_BASE/$STANDBY2_PORT -w
sleep 2
psql -Atq -c 'INSERT INTO aa VALUES (4)' -p $MASTER_PORT
psql -Atq -c 'INSERT INTO aa VALUES (5)' -p $STANDBY2_PORT
sleep 1

# Standby 3 switches from timeline 2 to 4
pg_ctl promote -s -D $DATA_BASE/$STANDBY3_PORT -w
sleep 1
psql -Atq -c 'INSERT INTO aa VALUES (6)' -p $STANDBY1_PORT
psql -Atq -c 'INSERT INTO aa VALUES (7)' -p $STANDBY3_PORT
sleep 1

# Switch segments for all nodes
psql -Atq -c 'SELECT pg_switch_xlog()' -p $MASTER_PORT
psql -Atq -c 'SELECT pg_switch_xlog()' -p $STANDBY1_PORT
psql -Atq -c 'SELECT pg_switch_xlog()' -p $STANDBY2_PORT
psql -Atq -c 'SELECT pg_switch_xlog()' -p $STANDBY3_PORT

# Clean up
operate_all_nodes stop

# And now play with pg_rewind and the cluster...
#replug_node $STANDBY3_PORT $MASTER_PORT # OK
#replug_node $MASTER_PORT $STANDBY3_PORT # OK
#replug_node $STANDBY2_PORT $MASTER_PORT # OK
#replug_node $MASTER_PORT $STANDBY2_PORT # OK
#replug_node $STANDBY1_PORT $MASTER_PORT # OK
#replug_node $MASTER_PORT $STANDBY1_PORT # OK
#replug_node $STANDBY1_PORT $STANDBY2_PORT # OK
#replug_node $STANDBY2_PORT $STANDBY1_PORT # OK
#replug_node $STANDBY1_PORT $STANDBY3_PORT # OK
#replug_node $STANDBY3_PORT $STANDBY1_PORT # OK
#replug_node $STANDBY2_PORT $STANDBY3_PORT # OK

# Replug things such as 4->3->2->1
#replug_node $STANDBY3_PORT $STANDBY2_PORT # OK
#operate_all_nodes stop
#replug_node $STANDBY2_PORT $STANDBY1_PORT # OK
#operate_all_nodes start
#operate_all_nodes stop
#replug_node $STANDBY1_PORT $MASTER_PORT # OK
#operate_all_nodes start

# Replug things such as 3->1->4->2
replug_node $STANDBY2_PORT $MASTER_PORT # OK
operate_all_nodes stop
replug_node $MASTER_PORT $STANDBY3_PORT # OK
operate_all_nodes start
operate_all_nodes stop
replug_node $STANDBY3_PORT $STANDBY1_PORT # OK
operate_all_nodes start