Increase shard move test coverage by improving advisory locks

To be able to test non-blocking shard moves we take an advisory lock, so
we can pause the shard move at an interesting moment. Originally this
was during the logical replication catch up phase. But when I added
tests for the rebalancer progress I moved this lock before the initial
data copy. This allowed testing of the rebalance progress, but
inadvertently made our non-blocking tests not actually test if we held
unintended locks during logical replication catch up.

This fixes that by creating two types of advisory locks, one before the
copy and one after. This causes the tests to actually test their
intended scenario again.

Furthermore it starts using one of these locks for blocking shard moves
too. Which allowed me to reduce the complexity of the rebalance progress
test suite quite a bit. It also allowed enabling some flaky tests again,
because this stopped them from being flaky. And finally it allowed
testing of rebalance progress for blocking shard copy operations as well.

src/backend/distributed/operations/shard_split.c

@@ -615,7 +615,7 @@ BlockingShardSplit(SplitOperation splitOperation,
 					snapshotName, distributionColumnOverrides);
 
 		/* Used for testing */
-		ConflictOnlyWithIsolationTesting();
+		ConflictOnlyWithIsolationTesting(false);
 
 		ereport(LOG, (errmsg(
 						  "creating auxillary structures (indexes, stats, replicaindentities, triggers) for %s",

src/backend/distributed/operations/shard_transfer.c

@@ -1258,6 +1258,7 @@ CopyShardTablesViaBlockWrites(List *shardIntervalList, char *sourceNodeName,
 	}
 
 	CopyShardsToNode(sourceNode, targetNode, shardIntervalList, NULL);
+	ConflictOnlyWithIsolationTesting(false);
 
 	foreach_ptr(shardInterval, shardIntervalList)
 	{

src/backend/distributed/replication/multi_logical_replication.c

@@ -236,7 +236,7 @@ LogicallyReplicateShards(List *shardList, char *sourceNodeName, int sourceNodePo
 			logicalRepTargetList);
 
 		/* only useful for isolation testing, see the function comment for the details */
-		ConflictOnlyWithIsolationTesting();
+		ConflictOnlyWithIsolationTesting(true);
 
 		/*
 		 * We have to create the primary key (or any other replica identity)
@@ -402,7 +402,7 @@ CompleteNonBlockingShardTransfer(List *shardList,
 
 
 	/* only useful for isolation testing, see the function comment for the details */
-	ConflictOnlyWithIsolationTesting();
+	ConflictOnlyWithIsolationTesting(false);
 
 	/*
 	 * We're almost done, we'll block the writes to the shards that we're
@@ -1241,9 +1241,14 @@ CreateUncheckedForeignKeyConstraints(List *logicalRepTargetList)
  *
  * Note that since the cost of calling this function is pretty low, we prefer
  * to use it in non-assert builds as well not to diverge in the behaviour.
+ *
+ * Depending on its argument it takes a different lock. It can take either a
+ * lock that is taken before the data copy is started, or it takes a lock
+ * after the copy (while logical replication is happening). Both of these locks
+ * are useful to run tests at different moments.
  */
 extern void
-ConflictOnlyWithIsolationTesting()
+ConflictOnlyWithIsolationTesting(bool beforeCopy)
 {
 	LOCKTAG tag;
 	const bool sessionLock = false;
@@ -1252,8 +1257,18 @@ ConflictOnlyWithIsolationTesting()
 	if (RunningUnderIsolationTest)
 	{
 		/* we've picked random keys */
-		SET_LOCKTAG_ADVISORY(tag, MyDatabaseId, SHARD_MOVE_ADVISORY_LOCK_FIRST_KEY,
-							 SHARD_MOVE_ADVISORY_LOCK_SECOND_KEY, 2);
+		if (beforeCopy)
+		{
+			SET_LOCKTAG_ADVISORY(tag, MyDatabaseId,
+								 SHARD_MOVE_ADVISORY_LOCK_SECOND_KEY,
+								 SHARD_MOVE_ADVISORY_LOCK_FIRST_KEY, 2);
+		}
+		else
+		{
+			SET_LOCKTAG_ADVISORY(tag, MyDatabaseId,
+								 SHARD_MOVE_ADVISORY_LOCK_FIRST_KEY,
+								 SHARD_MOVE_ADVISORY_LOCK_SECOND_KEY, 2);
+		}
 
 		(void) LockAcquire(&tag, ExclusiveLock, sessionLock, dontWait);
 	}

src/include/distributed/multi_logical_replication.h

@@ -131,7 +131,7 @@ extern void LogicallyReplicateShards(List *shardList, char *sourceNodeName,
 									 int sourceNodePort, char *targetNodeName,
 									 int targetNodePort);
 
-extern void ConflictOnlyWithIsolationTesting(void);
+extern void ConflictOnlyWithIsolationTesting(bool beforeCopy);
 extern void CreateReplicaIdentities(List *subscriptionInfoList);
 extern void CreateReplicaIdentitiesOnNode(List *shardList,
 										  char *nodeName,

src/test/regress/expected/isolation_shard_rebalancer_progress.out

@@ -1,6 +1,6 @@
 Parsed test spec with 7 sessions
 
-starting permutation: s2-lock-1-start s1-rebalance-c1-block-writes s7-get-progress s2-unlock-1-start s1-commit s7-get-progress enable-deferred-drop
+starting permutation: s2-lock-1-start s1-rebalance-c1-block-writes s7-get-progress s2-unlock-1-start s1-wait s7-get-progress
 master_set_node_property
 ---------------------------------------------------------------------
 
@@ -12,8 +12,6 @@ step s2-lock-1-start:
  DELETE FROM separate WHERE test_id = 1;
 
 step s1-rebalance-c1-block-writes:
- BEGIN;
- SELECT * FROM get_rebalance_table_shards_plan('colocated1');
  SELECT rebalance_table_shards('colocated1', shard_transfer_mode:='block_writes');
  <waiting ...>
 step s7-get-progress: 
@@ -48,22 +46,12 @@ step s2-unlock-1-start:
  ROLLBACK;
 
 step s1-rebalance-c1-block-writes: <... completed>
-table_name|shardid|shard_size|sourcename|sourceport|targetname|targetport
----------------------------------------------------------------------
-colocated1|1500001|         0|localhost |     57637|localhost |     57638
-colocated2|1500005|         0|localhost |     57637|localhost |     57638
-colocated1|1500002|         0|localhost |     57637|localhost |     57638
-colocated2|1500006|         0|localhost |     57637|localhost |     57638
-(4 rows)
-
 rebalance_table_shards
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s1-commit:
- COMMIT;
-
+step s1-wait:
 step s7-get-progress:
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
@@ -88,11 +76,8 @@ table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname
 ---------------------------------------------------------------------
 (0 rows)
 
-step enable-deferred-drop:
- ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
 
-
-starting permutation: s3-lock-2-start s1-rebalance-c1-block-writes s7-get-progress s3-unlock-2-start s1-commit s7-get-progress enable-deferred-drop
+starting permutation: s3-lock-2-start s1-rebalance-c1-block-writes s7-get-progress s3-unlock-2-start s1-wait s7-get-progress
 master_set_node_property
 ---------------------------------------------------------------------
 
@@ -103,8 +88,6 @@ step s3-lock-2-start:
  DELETE FROM colocated1 WHERE test_id = 3;
 
 step s1-rebalance-c1-block-writes:
- BEGIN;
- SELECT * FROM get_rebalance_table_shards_plan('colocated1');
  SELECT rebalance_table_shards('colocated1', shard_transfer_mode:='block_writes');
  <waiting ...>
 step s7-get-progress: 
@@ -129,8 +112,8 @@ step s7-get-progress:
 
 table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
 ---------------------------------------------------------------------
-colocated1|1500001|     50000|localhost |     57637|                0|localhost |     57638|            50000|       2|move          |t               |t                   |f
-colocated2|1500005|    400000|localhost |     57637|                0|localhost |     57638|           400000|       2|move          |t               |t                   |f
+colocated1|1500001|     50000|localhost |     57637|            50000|localhost |     57638|            50000|       2|move          |t               |t                   |f
+colocated2|1500005|    400000|localhost |     57637|           400000|localhost |     57638|           400000|       2|move          |t               |t                   |f
 colocated1|1500002|    200000|localhost |     57637|           200000|localhost |     57638|                0|       1|move          |t               |t                   |f
 colocated2|1500006|      8000|localhost |     57637|             8000|localhost |     57638|                0|       1|move          |t               |t                   |f
 (4 rows)
@@ -139,22 +122,12 @@ step s3-unlock-2-start:
  ROLLBACK;
 
 step s1-rebalance-c1-block-writes: <... completed>
-table_name|shardid|shard_size|sourcename|sourceport|targetname|targetport
----------------------------------------------------------------------
-colocated1|1500001|         0|localhost |     57637|localhost |     57638
-colocated2|1500005|         0|localhost |     57637|localhost |     57638
-colocated1|1500002|         0|localhost |     57637|localhost |     57638
-colocated2|1500006|         0|localhost |     57637|localhost |     57638
-(4 rows)
-
 rebalance_table_shards
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s1-commit:
- COMMIT;
-
+step s1-wait:
 step s7-get-progress:
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
@@ -179,35 +152,22 @@ table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname
 ---------------------------------------------------------------------
 (0 rows)
 
-step enable-deferred-drop:
- ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
 
-
-starting permutation: s7-grab-lock s1-rebalance-c1-block-writes s7-get-progress s7-release-lock s1-commit s7-get-progress enable-deferred-drop
+starting permutation: s6-acquire-advisory-lock s1-rebalance-c1-block-writes s7-get-progress s6-release-advisory-lock s1-wait s7-get-progress
 master_set_node_property
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s7-grab-lock:
- BEGIN;
- SET LOCAL citus.max_adaptive_executor_pool_size = 1;
- SELECT 1 FROM colocated1 LIMIT 1;
- SELECT 1 FROM separate LIMIT 1;
+step s6-acquire-advisory-lock:
+    SELECT pg_advisory_lock(44000, 55152);
 
-?column?
+pg_advisory_lock
 ---------------------------------------------------------------------
-       1
-(1 row)
 
-?column?
----------------------------------------------------------------------
-       1
 (1 row)
 
 step s1-rebalance-c1-block-writes:
- BEGIN;
- SELECT * FROM get_rebalance_table_shards_plan('colocated1');
  SELECT rebalance_table_shards('colocated1', shard_transfer_mode:='block_writes');
  <waiting ...>
 step s7-get-progress: 
@@ -238,26 +198,21 @@ colocated1|1500002|    200000|localhost |     57637|           200000|localhost
 colocated2|1500006|      8000|localhost |     57637|             8000|localhost |     57638|                0|       0|move          |t               |t                   |f
 (4 rows)
 
-step s7-release-lock:
- COMMIT;
+step s6-release-advisory-lock:
+    SELECT pg_advisory_unlock(44000, 55152);
+
+pg_advisory_unlock
+---------------------------------------------------------------------
+t
+(1 row)
 
 step s1-rebalance-c1-block-writes: <... completed>
-table_name|shardid|shard_size|sourcename|sourceport|targetname|targetport
----------------------------------------------------------------------
-colocated1|1500001|         0|localhost |     57637|localhost |     57638
-colocated2|1500005|         0|localhost |     57637|localhost |     57638
-colocated1|1500002|         0|localhost |     57637|localhost |     57638
-colocated2|1500006|         0|localhost |     57637|localhost |     57638
-(4 rows)
-
 rebalance_table_shards
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s1-commit:
- COMMIT;
-
+step s1-wait:
 step s7-get-progress:
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
@@ -282,18 +237,15 @@ table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname
 ---------------------------------------------------------------------
 (0 rows)
 
-step enable-deferred-drop:
- ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
 
-
-starting permutation: s6-acquire-advisory-lock s1-rebalance-c1-online s7-get-progress s6-release-advisory-lock s1-commit s7-get-progress enable-deferred-drop
+starting permutation: s5-acquire-advisory-lock s1-rebalance-c1-online s7-get-progress s5-release-advisory-lock s1-wait s7-get-progress
 master_set_node_property
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s6-acquire-advisory-lock:
-    SELECT pg_advisory_lock(44000, 55152);
+step s5-acquire-advisory-lock:
+    SELECT pg_advisory_lock(55152, 44000);
 
 pg_advisory_lock
 ---------------------------------------------------------------------
@@ -301,8 +253,6 @@ pg_advisory_lock
 (1 row)
 
 step s1-rebalance-c1-online:
- BEGIN;
- SELECT * FROM get_rebalance_table_shards_plan('colocated1');
  SELECT rebalance_table_shards('colocated1', shard_transfer_mode:='force_logical');
  <waiting ...>
 step s7-get-progress: 
@@ -333,8 +283,8 @@ colocated1|1500002|    200000|localhost |     57637|           200000|localhost
 colocated2|1500006|      8000|localhost |     57637|             8000|localhost |     57638|                0|       0|move          |t               |t                   |f
 (4 rows)
 
-step s6-release-advisory-lock:
-    SELECT pg_advisory_unlock(44000, 55152);
+step s5-release-advisory-lock:
+    SELECT pg_advisory_unlock(55152, 44000);
 
 pg_advisory_unlock
 ---------------------------------------------------------------------
@@ -342,22 +292,12 @@ t
 (1 row)
 
 step s1-rebalance-c1-online: <... completed>
-table_name|shardid|shard_size|sourcename|sourceport|targetname|targetport
----------------------------------------------------------------------
-colocated1|1500001|         0|localhost |     57637|localhost |     57638
-colocated2|1500005|         0|localhost |     57637|localhost |     57638
-colocated1|1500002|         0|localhost |     57637|localhost |     57638
-colocated2|1500006|         0|localhost |     57637|localhost |     57638
-(4 rows)
-
 rebalance_table_shards
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s1-commit:
- COMMIT;
-
+step s1-wait:
 step s7-get-progress:
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
@@ -382,11 +322,93 @@ table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname
 ---------------------------------------------------------------------
 (0 rows)
 
-step enable-deferred-drop:
- ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
+
+starting permutation: s6-acquire-advisory-lock s1-rebalance-c1-online s7-get-progress s6-release-advisory-lock s1-wait s7-get-progress
+master_set_node_property
+---------------------------------------------------------------------
+
+(1 row)
+
+step s6-acquire-advisory-lock:
+    SELECT pg_advisory_lock(44000, 55152);
+
+pg_advisory_lock
+---------------------------------------------------------------------
+
+(1 row)
+
+step s1-rebalance-c1-online:
+ SELECT rebalance_table_shards('colocated1', shard_transfer_mode:='force_logical');
+ <waiting ...>
+step s7-get-progress: 
+ set LOCAL client_min_messages=NOTICE;
+ WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
+ SELECT
+  table_name,
+  shardid,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - shard_size) = (SELECT MIN(ABS(size - shard_size)) FROM possible_sizes )) shard_size,
+  sourcename,
+  sourceport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - source_shard_size) = (SELECT MIN(ABS(size - source_shard_size)) FROM possible_sizes )) source_shard_size,
+  targetname,
+  targetport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - target_shard_size) = (SELECT MIN(ABS(size - target_shard_size)) FROM possible_sizes )) target_shard_size,
+  progress,
+  operation_type,
+  source_lsn >= target_lsn as lsn_sanity_check,
+  source_lsn > '0/0' as source_lsn_available,
+  target_lsn > '0/0' as target_lsn_available
+ FROM get_rebalance_progress();
+
+table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
+---------------------------------------------------------------------
+colocated1|1500001|     50000|localhost |     57637|            50000|localhost |     57638|            50000|       1|move          |t               |t                   |t
+colocated2|1500005|    400000|localhost |     57637|           400000|localhost |     57638|           400000|       1|move          |t               |t                   |t
+colocated1|1500002|    200000|localhost |     57637|           200000|localhost |     57638|                0|       0|move          |t               |t                   |f
+colocated2|1500006|      8000|localhost |     57637|             8000|localhost |     57638|                0|       0|move          |t               |t                   |f
+(4 rows)
+
+step s6-release-advisory-lock:
+    SELECT pg_advisory_unlock(44000, 55152);
+
+pg_advisory_unlock
+---------------------------------------------------------------------
+t
+(1 row)
+
+step s1-rebalance-c1-online: <... completed>
+rebalance_table_shards
+---------------------------------------------------------------------
+
+(1 row)
+
+step s1-wait:
+step s7-get-progress:
+ set LOCAL client_min_messages=NOTICE;
+ WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
+ SELECT
+  table_name,
+  shardid,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - shard_size) = (SELECT MIN(ABS(size - shard_size)) FROM possible_sizes )) shard_size,
+  sourcename,
+  sourceport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - source_shard_size) = (SELECT MIN(ABS(size - source_shard_size)) FROM possible_sizes )) source_shard_size,
+  targetname,
+  targetport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - target_shard_size) = (SELECT MIN(ABS(size - target_shard_size)) FROM possible_sizes )) target_shard_size,
+  progress,
+  operation_type,
+  source_lsn >= target_lsn as lsn_sanity_check,
+  source_lsn > '0/0' as source_lsn_available,
+  target_lsn > '0/0' as target_lsn_available
+ FROM get_rebalance_progress();
+
+table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
+---------------------------------------------------------------------
+(0 rows)
 
 
-starting permutation: s2-lock-1-start s1-shard-move-c1-block-writes s7-get-progress s2-unlock-1-start s1-commit s7-get-progress enable-deferred-drop
+starting permutation: s2-lock-1-start s1-shard-move-c1-block-writes s7-get-progress s2-unlock-1-start s1-wait s7-get-progress
 master_set_node_property
 ---------------------------------------------------------------------
 
@@ -398,7 +420,6 @@ step s2-lock-1-start:
  DELETE FROM separate WHERE test_id = 1;
 
 step s1-shard-move-c1-block-writes:
- BEGIN;
  SELECT citus_move_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
  <waiting ...>
 step s7-get-progress: 
@@ -436,9 +457,7 @@ citus_move_shard_placement
 
 (1 row)
 
-step s1-commit:
- COMMIT;
-
+step s1-wait:
 step s7-get-progress:
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
@@ -463,34 +482,22 @@ table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname
 ---------------------------------------------------------------------
 (0 rows)
 
-step enable-deferred-drop:
- ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
 
-
-starting permutation: s7-grab-lock s1-shard-move-c1-block-writes s7-get-progress s7-release-lock s1-commit s7-get-progress enable-deferred-drop
+starting permutation: s6-acquire-advisory-lock s1-shard-move-c1-block-writes s7-get-progress s6-release-advisory-lock s1-wait s7-get-progress
 master_set_node_property
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s7-grab-lock:
- BEGIN;
- SET LOCAL citus.max_adaptive_executor_pool_size = 1;
- SELECT 1 FROM colocated1 LIMIT 1;
- SELECT 1 FROM separate LIMIT 1;
+step s6-acquire-advisory-lock:
+    SELECT pg_advisory_lock(44000, 55152);
 
-?column?
+pg_advisory_lock
 ---------------------------------------------------------------------
-       1
-(1 row)
 
-?column?
----------------------------------------------------------------------
-       1
 (1 row)
 
 step s1-shard-move-c1-block-writes:
- BEGIN;
  SELECT citus_move_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
  <waiting ...>
 step s7-get-progress: 
@@ -519,8 +526,13 @@ colocated1|1500001|     50000|localhost |     57637|            50000|localhost
 colocated2|1500005|    400000|localhost |     57637|           400000|localhost |     57638|           400000|       1|move          |t               |t                   |f
 (2 rows)
 
-step s7-release-lock:
- COMMIT;
+step s6-release-advisory-lock:
+    SELECT pg_advisory_unlock(44000, 55152);
+
+pg_advisory_unlock
+---------------------------------------------------------------------
+t
+(1 row)
 
 step s1-shard-move-c1-block-writes: <... completed>
 citus_move_shard_placement
@@ -528,9 +540,7 @@ citus_move_shard_placement
 
 (1 row)
 
-step s1-commit:
- COMMIT;
-
+step s1-wait:
 step s7-get-progress:
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
@@ -555,11 +565,8 @@ table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname
 ---------------------------------------------------------------------
 (0 rows)
 
-step enable-deferred-drop:
- ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
 
-
-starting permutation: s2-lock-1-start s1-shard-copy-c1-block-writes s7-get-progress s2-unlock-1-start s1-commit
+starting permutation: s2-lock-1-start s1-shard-copy-c1-block-writes s7-get-progress s2-unlock-1-start s1-wait
 master_set_node_property
 ---------------------------------------------------------------------
 
@@ -571,7 +578,6 @@ step s2-lock-1-start:
  DELETE FROM separate WHERE test_id = 1;
 
 step s1-shard-copy-c1-block-writes:
- BEGIN;
  UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid IN ('colocated1', 'colocated2');
  SELECT citus_copy_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, transfer_mode:='block_writes');
  <waiting ...>
@@ -610,11 +616,9 @@ citus_copy_shard_placement
 
 (1 row)
 
-step s1-commit:
- COMMIT;
+step s1-wait:
 
-
-starting permutation: s6-acquire-advisory-lock s1-shard-move-c1-online s7-get-progress s6-release-advisory-lock s1-commit s7-get-progress enable-deferred-drop
+starting permutation: s6-acquire-advisory-lock s1-shard-copy-c1-block-writes s7-get-progress s6-release-advisory-lock s1-wait s7-get-progress
 master_set_node_property
 ---------------------------------------------------------------------
 
@@ -628,8 +632,91 @@ pg_advisory_lock
 
 (1 row)
 
+step s1-shard-copy-c1-block-writes:
+ UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid IN ('colocated1', 'colocated2');
+ SELECT citus_copy_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, transfer_mode:='block_writes');
+ <waiting ...>
+step s7-get-progress: 
+ set LOCAL client_min_messages=NOTICE;
+ WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
+ SELECT
+  table_name,
+  shardid,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - shard_size) = (SELECT MIN(ABS(size - shard_size)) FROM possible_sizes )) shard_size,
+  sourcename,
+  sourceport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - source_shard_size) = (SELECT MIN(ABS(size - source_shard_size)) FROM possible_sizes )) source_shard_size,
+  targetname,
+  targetport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - target_shard_size) = (SELECT MIN(ABS(size - target_shard_size)) FROM possible_sizes )) target_shard_size,
+  progress,
+  operation_type,
+  source_lsn >= target_lsn as lsn_sanity_check,
+  source_lsn > '0/0' as source_lsn_available,
+  target_lsn > '0/0' as target_lsn_available
+ FROM get_rebalance_progress();
+
+table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
+---------------------------------------------------------------------
+colocated1|1500001|     50000|localhost |     57637|            50000|localhost |     57638|            50000|       1|copy          |t               |t                   |f
+colocated2|1500005|    400000|localhost |     57637|           400000|localhost |     57638|           400000|       1|copy          |t               |t                   |f
+(2 rows)
+
+step s6-release-advisory-lock:
+    SELECT pg_advisory_unlock(44000, 55152);
+
+pg_advisory_unlock
+---------------------------------------------------------------------
+t
+(1 row)
+
+step s1-shard-copy-c1-block-writes: <... completed>
+citus_copy_shard_placement
+---------------------------------------------------------------------
+
+(1 row)
+
+step s1-wait:
+step s7-get-progress:
+ set LOCAL client_min_messages=NOTICE;
+ WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
+ SELECT
+  table_name,
+  shardid,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - shard_size) = (SELECT MIN(ABS(size - shard_size)) FROM possible_sizes )) shard_size,
+  sourcename,
+  sourceport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - source_shard_size) = (SELECT MIN(ABS(size - source_shard_size)) FROM possible_sizes )) source_shard_size,
+  targetname,
+  targetport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - target_shard_size) = (SELECT MIN(ABS(size - target_shard_size)) FROM possible_sizes )) target_shard_size,
+  progress,
+  operation_type,
+  source_lsn >= target_lsn as lsn_sanity_check,
+  source_lsn > '0/0' as source_lsn_available,
+  target_lsn > '0/0' as target_lsn_available
+ FROM get_rebalance_progress();
+
+table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
+---------------------------------------------------------------------
+(0 rows)
+
+
+starting permutation: s5-acquire-advisory-lock s1-shard-move-c1-online s7-get-progress s5-release-advisory-lock s1-wait s7-get-progress
+master_set_node_property
+---------------------------------------------------------------------
+
+(1 row)
+
+step s5-acquire-advisory-lock:
+    SELECT pg_advisory_lock(55152, 44000);
+
+pg_advisory_lock
+---------------------------------------------------------------------
+
+(1 row)
+
 step s1-shard-move-c1-online:
- BEGIN;
  SELECT citus_move_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='force_logical');
  <waiting ...>
 step s7-get-progress: 
@@ -658,8 +745,8 @@ colocated1|1500001|     50000|localhost |     57637|            50000|localhost
 colocated2|1500005|    400000|localhost |     57637|           400000|localhost |     57638|             8000|       1|move          |t               |t                   |f
 (2 rows)
 
-step s6-release-advisory-lock:
-    SELECT pg_advisory_unlock(44000, 55152);
+step s5-release-advisory-lock:
+    SELECT pg_advisory_unlock(55152, 44000);
 
 pg_advisory_unlock
 ---------------------------------------------------------------------
@@ -672,9 +759,7 @@ citus_move_shard_placement
 
 (1 row)
 
-step s1-commit:
- COMMIT;
-
+step s1-wait:
 step s7-get-progress:
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
@@ -699,11 +784,8 @@ table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname
 ---------------------------------------------------------------------
 (0 rows)
 
-step enable-deferred-drop:
- ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
 
-
-starting permutation: s6-acquire-advisory-lock s1-shard-copy-c1-online s7-get-progress s6-release-advisory-lock s1-commit
+starting permutation: s6-acquire-advisory-lock s1-shard-move-c1-online s7-get-progress s6-release-advisory-lock s1-wait s7-get-progress
 master_set_node_property
 ---------------------------------------------------------------------
 
@@ -717,8 +799,90 @@ pg_advisory_lock
 
 (1 row)
 
+step s1-shard-move-c1-online:
+ SELECT citus_move_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='force_logical');
+ <waiting ...>
+step s7-get-progress: 
+ set LOCAL client_min_messages=NOTICE;
+ WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
+ SELECT
+  table_name,
+  shardid,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - shard_size) = (SELECT MIN(ABS(size - shard_size)) FROM possible_sizes )) shard_size,
+  sourcename,
+  sourceport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - source_shard_size) = (SELECT MIN(ABS(size - source_shard_size)) FROM possible_sizes )) source_shard_size,
+  targetname,
+  targetport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - target_shard_size) = (SELECT MIN(ABS(size - target_shard_size)) FROM possible_sizes )) target_shard_size,
+  progress,
+  operation_type,
+  source_lsn >= target_lsn as lsn_sanity_check,
+  source_lsn > '0/0' as source_lsn_available,
+  target_lsn > '0/0' as target_lsn_available
+ FROM get_rebalance_progress();
+
+table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
+---------------------------------------------------------------------
+colocated1|1500001|     50000|localhost |     57637|            50000|localhost |     57638|            50000|       1|move          |t               |t                   |t
+colocated2|1500005|    400000|localhost |     57637|           400000|localhost |     57638|           400000|       1|move          |t               |t                   |t
+(2 rows)
+
+step s6-release-advisory-lock:
+    SELECT pg_advisory_unlock(44000, 55152);
+
+pg_advisory_unlock
+---------------------------------------------------------------------
+t
+(1 row)
+
+step s1-shard-move-c1-online: <... completed>
+citus_move_shard_placement
+---------------------------------------------------------------------
+
+(1 row)
+
+step s1-wait:
+step s7-get-progress:
+ set LOCAL client_min_messages=NOTICE;
+ WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
+ SELECT
+  table_name,
+  shardid,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - shard_size) = (SELECT MIN(ABS(size - shard_size)) FROM possible_sizes )) shard_size,
+  sourcename,
+  sourceport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - source_shard_size) = (SELECT MIN(ABS(size - source_shard_size)) FROM possible_sizes )) source_shard_size,
+  targetname,
+  targetport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - target_shard_size) = (SELECT MIN(ABS(size - target_shard_size)) FROM possible_sizes )) target_shard_size,
+  progress,
+  operation_type,
+  source_lsn >= target_lsn as lsn_sanity_check,
+  source_lsn > '0/0' as source_lsn_available,
+  target_lsn > '0/0' as target_lsn_available
+ FROM get_rebalance_progress();
+
+table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
+---------------------------------------------------------------------
+(0 rows)
+
+
+starting permutation: s5-acquire-advisory-lock s1-shard-copy-c1-online s7-get-progress s5-release-advisory-lock s1-wait
+master_set_node_property
+---------------------------------------------------------------------
+
+(1 row)
+
+step s5-acquire-advisory-lock:
+    SELECT pg_advisory_lock(55152, 44000);
+
+pg_advisory_lock
+---------------------------------------------------------------------
+
+(1 row)
+
 step s1-shard-copy-c1-online:
- BEGIN;
  UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid IN ('colocated1', 'colocated2');
  SELECT citus_copy_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, transfer_mode:='force_logical');
  <waiting ...>
@@ -748,8 +912,8 @@ colocated1|1500001|     50000|localhost |     57637|            50000|localhost
 colocated2|1500005|    400000|localhost |     57637|           400000|localhost |     57638|             8000|       1|copy          |t               |t                   |f
 (2 rows)
 
-step s6-release-advisory-lock:
-    SELECT pg_advisory_unlock(44000, 55152);
+step s5-release-advisory-lock:
+    SELECT pg_advisory_unlock(55152, 44000);
 
 pg_advisory_unlock
 ---------------------------------------------------------------------
@@ -762,28 +926,25 @@ citus_copy_shard_placement
 
 (1 row)
 
-step s1-commit:
- COMMIT;
+step s1-wait:
 
-
-starting permutation: s2-lock-1-start s1-shard-move-c1-block-writes s4-shard-move-sep-block-writes s7-get-progress s2-unlock-1-start s1-commit s4-commit s7-get-progress enable-deferred-drop
+starting permutation: s6-acquire-advisory-lock s1-shard-copy-c1-online s7-get-progress s6-release-advisory-lock s1-wait s7-get-progress
 master_set_node_property
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s2-lock-1-start:
- BEGIN;
- DELETE FROM colocated1 WHERE test_id = 1;
- DELETE FROM separate WHERE test_id = 1;
+step s6-acquire-advisory-lock:
+    SELECT pg_advisory_lock(44000, 55152);
 
-step s1-shard-move-c1-block-writes:
- BEGIN;
- SELECT citus_move_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
- <waiting ...>
-step s4-shard-move-sep-block-writes: 
- BEGIN;
- SELECT citus_move_shard_placement(1500009, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
+pg_advisory_lock
+---------------------------------------------------------------------
+
+(1 row)
+
+step s1-shard-copy-c1-online:
+ UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid IN ('colocated1', 'colocated2');
+ SELECT citus_copy_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, transfer_mode:='force_logical');
  <waiting ...>
 step s7-get-progress: 
  set LOCAL client_min_messages=NOTICE;
@@ -805,6 +966,92 @@ step s7-get-progress:
   target_lsn > '0/0' as target_lsn_available
  FROM get_rebalance_progress();
 
+table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
+---------------------------------------------------------------------
+colocated1|1500001|     50000|localhost |     57637|            50000|localhost |     57638|            50000|       1|copy          |t               |t                   |t
+colocated2|1500005|    400000|localhost |     57637|           400000|localhost |     57638|           400000|       1|copy          |t               |t                   |t
+(2 rows)
+
+step s6-release-advisory-lock:
+    SELECT pg_advisory_unlock(44000, 55152);
+
+pg_advisory_unlock
+---------------------------------------------------------------------
+t
+(1 row)
+
+step s1-shard-copy-c1-online: <... completed>
+citus_copy_shard_placement
+---------------------------------------------------------------------
+
+(1 row)
+
+step s1-wait:
+step s7-get-progress:
+ set LOCAL client_min_messages=NOTICE;
+ WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
+ SELECT
+  table_name,
+  shardid,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - shard_size) = (SELECT MIN(ABS(size - shard_size)) FROM possible_sizes )) shard_size,
+  sourcename,
+  sourceport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - source_shard_size) = (SELECT MIN(ABS(size - source_shard_size)) FROM possible_sizes )) source_shard_size,
+  targetname,
+  targetport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - target_shard_size) = (SELECT MIN(ABS(size - target_shard_size)) FROM possible_sizes )) target_shard_size,
+  progress,
+  operation_type,
+  source_lsn >= target_lsn as lsn_sanity_check,
+  source_lsn > '0/0' as source_lsn_available,
+  target_lsn > '0/0' as target_lsn_available
+ FROM get_rebalance_progress();
+
+table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
+---------------------------------------------------------------------
+(0 rows)
+
+
+starting permutation: s2-lock-1-start s1-shard-move-c1-block-writes s4-shard-move-sep-block-writes-without-advisory-locks s7-get-progress-ordered s2-unlock-1-start s1-wait s4-commit s7-get-progress-ordered
+master_set_node_property
+---------------------------------------------------------------------
+
+(1 row)
+
+step s2-lock-1-start:
+ BEGIN;
+ DELETE FROM colocated1 WHERE test_id = 1;
+ DELETE FROM separate WHERE test_id = 1;
+
+step s1-shard-move-c1-block-writes:
+ SELECT citus_move_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
+ <waiting ...>
+step s4-shard-move-sep-block-writes-without-advisory-locks: 
+ BEGIN;
+ SET LOCAL citus.running_under_isolation_test = false;
+ SELECT citus_move_shard_placement(1500009, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
+ <waiting ...>
+step s7-get-progress-ordered: 
+ set LOCAL client_min_messages=NOTICE;
+ WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
+ SELECT
+  table_name,
+  shardid,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - shard_size) = (SELECT MIN(ABS(size - shard_size)) FROM possible_sizes )) shard_size,
+  sourcename,
+  sourceport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - source_shard_size) = (SELECT MIN(ABS(size - source_shard_size)) FROM possible_sizes )) source_shard_size,
+  targetname,
+  targetport,
+  ( SELECT size FROM possible_sizes WHERE ABS(size - target_shard_size) = (SELECT MIN(ABS(size - target_shard_size)) FROM possible_sizes )) target_shard_size,
+  progress,
+  operation_type,
+  source_lsn >= target_lsn as lsn_sanity_check,
+  source_lsn > '0/0' as source_lsn_available,
+  target_lsn > '0/0' as target_lsn_available
+ FROM get_rebalance_progress()
+ ORDER BY 1, 2, 3, 4, 5;
+
 table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
 ---------------------------------------------------------------------
 colocated1|1500001|     50000|localhost |     57637|            50000|localhost |     57638|                0|       1|move          |t               |t                   |f
@@ -821,19 +1068,17 @@ citus_move_shard_placement
 
 (1 row)
 
-step s4-shard-move-sep-block-writes: <... completed>
+step s4-shard-move-sep-block-writes-without-advisory-locks: <... completed>
 citus_move_shard_placement
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s1-commit:
- COMMIT;
-
+step s1-wait:
 step s4-commit:
  COMMIT;
 
-step s7-get-progress:
+step s7-get-progress-ordered:
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
  SELECT
@@ -851,47 +1096,36 @@ step s7-get-progress:
   source_lsn >= target_lsn as lsn_sanity_check,
   source_lsn > '0/0' as source_lsn_available,
   target_lsn > '0/0' as target_lsn_available
- FROM get_rebalance_progress();
+ FROM get_rebalance_progress()
+ ORDER BY 1, 2, 3, 4, 5;
 
 table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
 ---------------------------------------------------------------------
 (0 rows)
 
-step enable-deferred-drop:
- ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
 
-
-starting permutation: s7-grab-lock s1-shard-move-c1-block-writes s4-shard-move-sep-block-writes s7-get-progress s7-release-lock s1-commit s4-commit s7-get-progress enable-deferred-drop
+starting permutation: s6-acquire-advisory-lock s1-shard-move-c1-block-writes s4-shard-move-sep-block-writes s7-get-progress-ordered s6-release-advisory-lock s1-wait s4-commit s7-get-progress-ordered
 master_set_node_property
 ---------------------------------------------------------------------
 
 (1 row)
 
-step s7-grab-lock:
- BEGIN;
- SET LOCAL citus.max_adaptive_executor_pool_size = 1;
- SELECT 1 FROM colocated1 LIMIT 1;
- SELECT 1 FROM separate LIMIT 1;
+step s6-acquire-advisory-lock:
+    SELECT pg_advisory_lock(44000, 55152);
 
-?column?
+pg_advisory_lock
 ---------------------------------------------------------------------
-       1
-(1 row)
 
-?column?
----------------------------------------------------------------------
-       1
 (1 row)
 
 step s1-shard-move-c1-block-writes:
- BEGIN;
  SELECT citus_move_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
  <waiting ...>
 step s4-shard-move-sep-block-writes: 
  BEGIN;
  SELECT citus_move_shard_placement(1500009, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
  <waiting ...>
-step s7-get-progress: 
+step s7-get-progress-ordered: 
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
  SELECT
@@ -909,7 +1143,8 @@ step s7-get-progress:
   source_lsn >= target_lsn as lsn_sanity_check,
   source_lsn > '0/0' as source_lsn_available,
   target_lsn > '0/0' as target_lsn_available
- FROM get_rebalance_progress();
+ FROM get_rebalance_progress()
+ ORDER BY 1, 2, 3, 4, 5;
 
 table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
 ---------------------------------------------------------------------
@@ -918,8 +1153,13 @@ colocated2|1500005|    400000|localhost |     57637|           400000|localhost
 separate  |1500009|     50000|localhost |     57637|            50000|localhost |     57638|           200000|       1|move          |t               |t                   |f
 (3 rows)
 
-step s7-release-lock:
- COMMIT;
+step s6-release-advisory-lock:
+    SELECT pg_advisory_unlock(44000, 55152);
+
+pg_advisory_unlock
+---------------------------------------------------------------------
+t
+(1 row)
 
 step s1-shard-move-c1-block-writes: <... completed>
 citus_move_shard_placement
@@ -933,13 +1173,11 @@ citus_move_shard_placement
 
 (1 row)
 
-step s1-commit:
- COMMIT;
-
+step s1-wait:
 step s4-commit:
  COMMIT;
 
-step s7-get-progress:
+step s7-get-progress-ordered:
  set LOCAL client_min_messages=NOTICE;
  WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
  SELECT
@@ -957,12 +1195,10 @@ step s7-get-progress:
   source_lsn >= target_lsn as lsn_sanity_check,
   source_lsn > '0/0' as source_lsn_available,
   target_lsn > '0/0' as target_lsn_available
- FROM get_rebalance_progress();
+ FROM get_rebalance_progress()
+ ORDER BY 1, 2, 3, 4, 5;
 
 table_name|shardid|shard_size|sourcename|sourceport|source_shard_size|targetname|targetport|target_shard_size|progress|operation_type|lsn_sanity_check|source_lsn_available|target_lsn_available
 ---------------------------------------------------------------------
 (0 rows)
 
-step enable-deferred-drop:
- ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
-

src/test/regress/spec/isolation_shard_rebalancer_progress.spec

@@ -1,13 +1,9 @@
 setup
 {
-	-- We disable deffered drop, so we can easily trigger blocking the shard
-	-- move at the end of the move. This is done in a separate setup step,
-	-- because this cannot run in a transaction.
-	ALTER SYSTEM SET citus.defer_drop_after_shard_move TO OFF;
+	CALL citus_cleanup_orphaned_shards();
 }
 setup
 {
-	SELECT pg_reload_conf();
 	ALTER SEQUENCE pg_catalog.pg_dist_shardid_seq RESTART 1500001;
 	SET citus.shard_count TO 4;
 	SET citus.shard_replication_factor TO 1;
@@ -30,7 +26,6 @@ setup
 
 teardown
 {
-	SELECT pg_reload_conf();
 	DROP TABLE colocated2;
 	DROP TABLE colocated1;
 	DROP TABLE separate;
@@ -40,48 +35,37 @@ session "s1"
 
 step "s1-rebalance-c1-block-writes"
 {
-	BEGIN;
-	SELECT * FROM get_rebalance_table_shards_plan('colocated1');
 	SELECT rebalance_table_shards('colocated1', shard_transfer_mode:='block_writes');
 }
 
 step "s1-rebalance-c1-online"
 {
-	BEGIN;
-	SELECT * FROM get_rebalance_table_shards_plan('colocated1');
 	SELECT rebalance_table_shards('colocated1', shard_transfer_mode:='force_logical');
 }
 
 step "s1-shard-move-c1-block-writes"
 {
-	BEGIN;
 	SELECT citus_move_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
 }
 
 step "s1-shard-copy-c1-block-writes"
 {
-	BEGIN;
 	UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid IN ('colocated1', 'colocated2');
 	SELECT citus_copy_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, transfer_mode:='block_writes');
 }
 
 step "s1-shard-move-c1-online"
 {
-	BEGIN;
 	SELECT citus_move_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='force_logical');
 }
 
 step "s1-shard-copy-c1-online"
 {
-	BEGIN;
 	UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid IN ('colocated1', 'colocated2');
 	SELECT citus_copy_shard_placement(1500001, 'localhost', 57637, 'localhost', 57638, transfer_mode:='force_logical');
 }
 
-step "s1-commit"
-{
-	COMMIT;
-}
+step "s1-wait" {}
 
 session "s2"
 
@@ -118,11 +102,38 @@ step "s4-shard-move-sep-block-writes"
 	SELECT citus_move_shard_placement(1500009, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
 }
 
+// Running two shard moves at the same time can cause racy behaviour over who
+// gets the lock. For the test where this move is used in we don't rely on the
+// advisory locks. So we disable taking the advisory lock there to avoid the
+// racy lock acquisition with the other concurretn move.
+step "s4-shard-move-sep-block-writes-without-advisory-locks"
+{
+	BEGIN;
+	SET LOCAL citus.running_under_isolation_test = false;
+	SELECT citus_move_shard_placement(1500009, 'localhost', 57637, 'localhost', 57638, shard_transfer_mode:='block_writes');
+}
+
 step "s4-commit"
 {
 	COMMIT;
 }
 
+session "s5"
+
+// this advisory lock with (almost) random values are only used
+// for testing purposes. For details, check Citus' logical replication
+// source code
+step "s5-acquire-advisory-lock"
+{
+    SELECT pg_advisory_lock(55152, 44000);
+}
+
+step "s5-release-advisory-lock"
+{
+    SELECT pg_advisory_unlock(55152, 44000);
+}
+
+
 session "s6"
 
 // this advisory lock with (almost) random values are only used
@@ -141,22 +152,6 @@ step "s6-release-advisory-lock"
 
 session "s7"
 
-// get_rebalance_progress internally calls pg_total_relation_size on all the
-// shards. This means that it takes AccessShareLock on those shards. Because we
-// run with deferred drop that means that get_rebalance_progress actually waits
-// for the shard move to complete the drop. But we want to get the progress
-// before the shards are dropped. So we grab the locks first with a simple
-// query that reads from all shards. We force using a single connection because
-// get_rebalance_progress isn't smart enough to reuse the right connection for
-// the right shards and will simply use a single one for all of them.
-step "s7-grab-lock"
-{
-	BEGIN;
-	SET LOCAL citus.max_adaptive_executor_pool_size = 1;
-	SELECT 1 FROM colocated1 LIMIT 1;
-	SELECT 1 FROM separate LIMIT 1;
-}
-
 step "s7-get-progress"
 {
 	set LOCAL client_min_messages=NOTICE;
@@ -179,46 +174,57 @@ step "s7-get-progress"
 	FROM get_rebalance_progress();
 }
 
-step "s7-release-lock"
+// When getting progress from multiple monitors at the same time it can result
+// in random order of the tuples, because there's no defined order of the
+// monitors. So in those cases we need to order the output for consistent results.
+step "s7-get-progress-ordered"
 {
-	COMMIT;
+	set LOCAL client_min_messages=NOTICE;
+	WITH possible_sizes(size) as (VALUES (0), (8000), (50000), (200000), (400000))
+	SELECT
+		table_name,
+		shardid,
+		( SELECT size FROM possible_sizes WHERE ABS(size - shard_size) = (SELECT MIN(ABS(size - shard_size)) FROM possible_sizes )) shard_size,
+		sourcename,
+		sourceport,
+		( SELECT size FROM possible_sizes WHERE ABS(size - source_shard_size) = (SELECT MIN(ABS(size - source_shard_size)) FROM possible_sizes )) source_shard_size,
+		targetname,
+		targetport,
+		( SELECT size FROM possible_sizes WHERE ABS(size - target_shard_size) = (SELECT MIN(ABS(size - target_shard_size)) FROM possible_sizes )) target_shard_size,
+		progress,
+		operation_type,
+		source_lsn >= target_lsn as lsn_sanity_check,
+		source_lsn > '0/0' as source_lsn_available,
+		target_lsn > '0/0' as target_lsn_available
+	FROM get_rebalance_progress()
+	ORDER BY 1, 2, 3, 4, 5;
 }
 
-session "s8"
-
-// After running these tests we want to enable deferred-drop again. Sadly
-// the isolation tester framework does not support multiple teardown steps
-// and this cannot be run in a transaction. So we need to do it manually at
-// the end of the last test.
-step "enable-deferred-drop"
-{
-	ALTER SYSTEM RESET citus.defer_drop_after_shard_move;
-}
 // blocking rebalancer does what it should
-permutation "s2-lock-1-start" "s1-rebalance-c1-block-writes" "s7-get-progress" "s2-unlock-1-start" "s1-commit" "s7-get-progress" "enable-deferred-drop"
-permutation "s3-lock-2-start" "s1-rebalance-c1-block-writes" "s7-get-progress" "s3-unlock-2-start" "s1-commit" "s7-get-progress" "enable-deferred-drop"
-permutation "s7-grab-lock" "s1-rebalance-c1-block-writes" "s7-get-progress" "s7-release-lock" "s1-commit" "s7-get-progress" "enable-deferred-drop"
+permutation "s2-lock-1-start" "s1-rebalance-c1-block-writes" "s7-get-progress" "s2-unlock-1-start" "s1-wait" "s7-get-progress"
+permutation "s3-lock-2-start" "s1-rebalance-c1-block-writes" "s7-get-progress" "s3-unlock-2-start" "s1-wait" "s7-get-progress"
+permutation "s6-acquire-advisory-lock" "s1-rebalance-c1-block-writes" "s7-get-progress" "s6-release-advisory-lock" "s1-wait" "s7-get-progress"
 
 // online rebalancer
-permutation "s6-acquire-advisory-lock" "s1-rebalance-c1-online" "s7-get-progress" "s6-release-advisory-lock" "s1-commit" "s7-get-progress" "enable-deferred-drop"
-// Commented out due to flakyness
-// permutation "s7-grab-lock" "s1-rebalance-c1-online" "s7-get-progress" "s7-release-lock" "s1-commit" "s7-get-progress" "enable-deferred-drop"
+permutation "s5-acquire-advisory-lock" "s1-rebalance-c1-online" "s7-get-progress" "s5-release-advisory-lock" "s1-wait" "s7-get-progress"
+permutation "s6-acquire-advisory-lock" "s1-rebalance-c1-online" "s7-get-progress" "s6-release-advisory-lock" "s1-wait" "s7-get-progress"
 
 // blocking shard move
-permutation "s2-lock-1-start" "s1-shard-move-c1-block-writes" "s7-get-progress" "s2-unlock-1-start" "s1-commit" "s7-get-progress" "enable-deferred-drop"
-permutation "s7-grab-lock" "s1-shard-move-c1-block-writes" "s7-get-progress" "s7-release-lock" "s1-commit" "s7-get-progress" "enable-deferred-drop"
+permutation "s2-lock-1-start" "s1-shard-move-c1-block-writes" "s7-get-progress" "s2-unlock-1-start" "s1-wait" "s7-get-progress"
+permutation "s6-acquire-advisory-lock" "s1-shard-move-c1-block-writes" "s7-get-progress" "s6-release-advisory-lock" "s1-wait" "s7-get-progress"
 
 // blocking shard copy
-permutation "s2-lock-1-start" "s1-shard-copy-c1-block-writes" "s7-get-progress" "s2-unlock-1-start" "s1-commit"
+permutation "s2-lock-1-start" "s1-shard-copy-c1-block-writes" "s7-get-progress" "s2-unlock-1-start" "s1-wait"
+permutation "s6-acquire-advisory-lock" "s1-shard-copy-c1-block-writes" "s7-get-progress" "s6-release-advisory-lock" "s1-wait" "s7-get-progress"
 
 // online shard move
-permutation "s6-acquire-advisory-lock" "s1-shard-move-c1-online" "s7-get-progress" "s6-release-advisory-lock" "s1-commit" "s7-get-progress" "enable-deferred-drop"
-// Commented out due to flakyness
-// permutation "s7-grab-lock" "s1-shard-move-c1-online" "s7-get-progress" "s7-release-lock" "s1-commit" "s7-get-progress" "enable-deferred-drop"
+permutation "s5-acquire-advisory-lock" "s1-shard-move-c1-online" "s7-get-progress" "s5-release-advisory-lock" "s1-wait" "s7-get-progress"
+permutation "s6-acquire-advisory-lock" "s1-shard-move-c1-online" "s7-get-progress" "s6-release-advisory-lock" "s1-wait" "s7-get-progress"
 
 // online shard copy
-permutation "s6-acquire-advisory-lock" "s1-shard-copy-c1-online" "s7-get-progress" "s6-release-advisory-lock" "s1-commit"
+permutation "s5-acquire-advisory-lock" "s1-shard-copy-c1-online" "s7-get-progress" "s5-release-advisory-lock" "s1-wait"
+permutation "s6-acquire-advisory-lock" "s1-shard-copy-c1-online" "s7-get-progress" "s6-release-advisory-lock" "s1-wait" "s7-get-progress"
 
 // parallel blocking shard move
-permutation "s2-lock-1-start" "s1-shard-move-c1-block-writes" "s4-shard-move-sep-block-writes"("s1-shard-move-c1-block-writes") "s7-get-progress" "s2-unlock-1-start" "s1-commit" "s4-commit" "s7-get-progress" "enable-deferred-drop"
-permutation "s7-grab-lock" "s1-shard-move-c1-block-writes" "s4-shard-move-sep-block-writes"("s1-shard-move-c1-block-writes") "s7-get-progress" "s7-release-lock" "s1-commit" "s4-commit" "s7-get-progress" "enable-deferred-drop"
+permutation "s2-lock-1-start" "s1-shard-move-c1-block-writes" "s4-shard-move-sep-block-writes-without-advisory-locks"("s1-shard-move-c1-block-writes") "s7-get-progress-ordered" "s2-unlock-1-start" "s1-wait" "s4-commit" "s7-get-progress-ordered"
+permutation "s6-acquire-advisory-lock" "s1-shard-move-c1-block-writes" "s4-shard-move-sep-block-writes"("s1-shard-move-c1-block-writes") "s7-get-progress-ordered" "s6-release-advisory-lock"  "s1-wait" "s4-commit" "s7-get-progress-ordered"