diff --git a/src/test/isolation2/expected/segwalrep/restartpoint_remove_xlog.out b/src/test/isolation2/expected/segwalrep/restartpoint_remove_xlog.out
index 6341c0cdf3c84351e03748ea0d699046d5301594..f9342c18b23795dd8419c707db701080e2cc1fd9 100644
--- a/src/test/isolation2/expected/segwalrep/restartpoint_remove_xlog.out
+++ b/src/test/isolation2/expected/segwalrep/restartpoint_remove_xlog.out
@@ -21,12 +21,6 @@ CREATE
 
 create extension if not exists gp_inject_fault;
 CREATE
-
-create or replace function wait_for_replication_replay (retries int) returns bool as $$ declare i int; /* in func */ result bool; /* in func */ begin /* in func */ i := 0; /* in func */ -- Wait until the mirror (content 0) has replayed up to flush location loop /* in func */ SELECT flush_location = replay_location INTO result from gp_stat_replication where gp_segment_id = 0; /* in func */ if result then /* in func */ return true; /* in func */ end if; /* in func */ 
-if i >= retries then /* in func */ return false; /* in func */ end if; /* in func */ perform pg_sleep(0.1); /* in func */ i := i + 1; /* in func */ end loop; /* in func */ end; /* in func */ $$ language plpgsql;
-CREATE
-
-
 create table t_restart (a int);
 CREATE
 
@@ -70,7 +64,7 @@ CHECKPOINT
 -- and then if the mirror is promoted it will panic like this:
 -- FATAL","58P01","requested WAL segment pg_xlog/000000010000000000000003 has already been removed
 -- The call stack is: StartupXLOG()->PrescanPreparedTransactions()...
-select * from wait_for_replication_replay(5000);
+select * from wait_for_replication_replay(0, 5000);
  wait_for_replication_replay 
 -----------------------------
  t                           
diff --git a/src/test/isolation2/expected/uao_crash_compaction_column.out b/src/test/isolation2/expected/uao_crash_compaction_column.out
index c6aa6b5f3f4a69a1d3d4ee4b95f72c7149b2d7d1..047c732b7b2c8ee23c5b918f98bc85c3a089da5c 100644
--- a/src/test/isolation2/expected/uao_crash_compaction_column.out
+++ b/src/test/isolation2/expected/uao_crash_compaction_column.out
@@ -8,10 +8,6 @@
 include: helpers/server_helpers.sql;
 CREATE
 
-create or replace function wait_for_replication_replay (retries int) returns bool as $$ declare i int; /* in func */ result bool; /* in func */ begin /* in func */ i := 0; /* in func */ -- Wait until the mirror (content 0) has replayed up to flush location loop /* in func */ SELECT flush_location = replay_location INTO result from gp_stat_replication where gp_segment_id = 0; /* in func */ if result then /* in func */ return true; /* in func */ end if; /* in func */ 
-if i >= retries then /* in func */ return false; /* in func */ end if; /* in func */ perform pg_sleep(0.1); /* in func */ i := i + 1; /* in func */ end loop; /* in func */ end; /* in func */ $$ language plpgsql;
-CREATE
-
 3:CREATE extension if NOT EXISTS gp_inject_fault;
 CREATE
 3:SELECT role, preferred_role, content, mode, status FROM gp_segment_configuration;
@@ -643,7 +639,7 @@ VACUUM
 -- Make sure mirror is able to successfully replay all the truncate
 -- records generated and doesn't encounter the "WAL contains
 -- references to invalid pages" PANIC.
-6:SELECT * from wait_for_replication_replay(5000);
+6:SELECT * from wait_for_replication_replay(0, 5000);
  wait_for_replication_replay 
 -----------------------------
  t                           
diff --git a/src/test/isolation2/helpers/server_helpers.sql b/src/test/isolation2/helpers/server_helpers.sql
index 33ed7020422344236538b43f9674915249e13331..ffb59d4cca6f814708d99deace7b4657e460cc74 100644
--- a/src/test/isolation2/helpers/server_helpers.sql
+++ b/src/test/isolation2/helpers/server_helpers.sql
@@ -165,3 +165,26 @@ begin
 	return 'Fail';
 end;
 $$ language plpgsql;
+
+create or replace function wait_for_replication_replay (segid int, retries int) returns bool as
+$$
+declare
+	i int;
+	result bool;
+begin
+	i := 0;
+	-- Wait until the mirror/standby has replayed up to flush location
+	loop
+		SELECT flush_location = replay_location INTO result from gp_stat_replication where gp_segment_id = segid;
+		if result then
+			return true;
+		end if;
+
+		if i >= retries then
+		   return false;
+		end if;
+		perform pg_sleep(0.1);
+		i := i + 1;
+	end loop;
+end;
+$$ language plpgsql;
diff --git a/src/test/isolation2/sql/segwalrep/restartpoint_remove_xlog.sql b/src/test/isolation2/sql/segwalrep/restartpoint_remove_xlog.sql
index f58c223d8fb56040c613580a1e9525247aaeffe1..17ccd58d9ea59fdd9757398eb122689513ed5875 100644
--- a/src/test/isolation2/sql/segwalrep/restartpoint_remove_xlog.sql
+++ b/src/test/isolation2/sql/segwalrep/restartpoint_remove_xlog.sql
@@ -14,31 +14,6 @@ include: helpers/server_helpers.sql;
 !\retcode gpstop -u;
 
 create extension if not exists gp_inject_fault;
-
-create or replace function wait_for_replication_replay (retries int) returns bool as
-$$
-declare
-	i int; /* in func */
-	result bool; /* in func */
-begin /* in func */
-	i := 0; /* in func */
-	-- Wait until the mirror (content 0) has replayed up to flush location
-	loop /* in func */
-		SELECT flush_location = replay_location INTO result from gp_stat_replication where gp_segment_id = 0; /* in func */
-		if result then /* in func */
-			return true; /* in func */
-		end if; /* in func */
-
-		if i >= retries then /* in func */
-		   return false; /* in func */
-		end if; /* in func */
-		perform pg_sleep(0.1); /* in func */
-		i := i + 1; /* in func */
-	end loop; /* in func */
-end; /* in func */
-$$ language plpgsql;
-
-
 create table t_restart (a int);
 
 -- generate an orphaned prepare transaction.
@@ -65,7 +40,7 @@ checkpoint;
 -- and then if the mirror is promoted it will panic like this:
 -- FATAL","58P01","requested WAL segment pg_xlog/000000010000000000000003 has already been removed
 -- The call stack is: StartupXLOG()->PrescanPreparedTransactions()...
-select * from wait_for_replication_replay(5000);
+select * from wait_for_replication_replay(0, 5000);
 
 -- shutdown primary and make sure the segment is down
 -1U: select pg_ctl((SELECT datadir from gp_segment_configuration c
diff --git a/src/test/isolation2/sql/uao_crash_compaction_column.sql b/src/test/isolation2/sql/uao_crash_compaction_column.sql
index 7cfdf422cc7e6fbcd206e5e8a12fe7696cdf9f1b..204a88087b7ebb4ed8b879d339c876e4b2f880bc 100644
--- a/src/test/isolation2/sql/uao_crash_compaction_column.sql
+++ b/src/test/isolation2/sql/uao_crash_compaction_column.sql
@@ -7,29 +7,6 @@
 -- end_matchsubs
 include: helpers/server_helpers.sql;
 
-create or replace function wait_for_replication_replay (retries int) returns bool as
-$$
-declare
-	i int; /* in func */
-	result bool; /* in func */
-begin /* in func */
-	i := 0; /* in func */
-	-- Wait until the mirror (content 0) has replayed up to flush location
-	loop /* in func */
-		SELECT flush_location = replay_location INTO result from gp_stat_replication where gp_segment_id = 0; /* in func */
-		if result then /* in func */
-			return true; /* in func */
-		end if; /* in func */
-
-		if i >= retries then /* in func */
-		   return false; /* in func */
-		end if; /* in func */
-		perform pg_sleep(0.1); /* in func */
-		i := i + 1; /* in func */
-	end loop; /* in func */
-end; /* in func */
-$$ language plpgsql;
-
 3:CREATE extension if NOT EXISTS gp_inject_fault;
 3:SELECT role, preferred_role, content, mode, status FROM gp_segment_configuration;
 --
@@ -216,5 +193,5 @@ where c.role='p' and c.content=0), 'restart');
 -- Make sure mirror is able to successfully replay all the truncate
 -- records generated and doesn't encounter the "WAL contains
 -- references to invalid pages" PANIC.
-6:SELECT * from wait_for_replication_replay(5000);
+6:SELECT * from wait_for_replication_replay(0, 5000);
 6:SELECT gp_inject_fault('fts_probe', 'reset', 1);