Simplify metadata lock acquisition for DML

src/backend/distributed/executor/multi_router_executor.c

@@ -73,7 +73,7 @@ bool AllModificationsCommutative = false;
 bool EnableDeadlockPrevention = true;
 
 /* functions needed during run phase */
-static void ReacquireMetadataLocks(List *taskList);
+static void AcquireMetadataLocks(List *taskList);
 static ShardPlacementAccess * CreatePlacementAccess(ShardPlacement *placement,
 													ShardPlacementAccessType accessType);
 static void ExecuteSingleModifyTask(CitusScanState *scanState, Task *task,
@@ -100,22 +100,12 @@ static bool ConsumeQueryResult(MultiConnection *connection, bool failOnError,
 
 
 /*
- * ReacquireMetadataLocks re-acquires the metadata locks that are normally
+ * AcquireMetadataLocks acquires metadata locks on each of the anchor
- * acquired during planning.
+ * shards in the task list to prevent a shard being modified while it
- *
+ * is being copied.
- * If we are executing a prepared statement, then planning might have
- * happened in a separate transaction and advisory locks are no longer
- * held. If a shard is currently being repaired/copied/moved, then
- * obtaining the locks will fail and this function throws an error to
- * prevent executing a stale plan.
- *
- * If we are executing a non-prepared statement or planning happened in
- * the same transaction, then we already have the locks and obtain them
- * again here. Since we always release these locks at the end of the
- * transaction, this is effectively a no-op.
  */
 static void
-ReacquireMetadataLocks(List *taskList)
+AcquireMetadataLocks(List *taskList)
 {
 	ListCell *taskCell = NULL;
 
@@ -130,28 +120,7 @@ ReacquireMetadataLocks(List *taskList)
 	{
 		Task *task = (Task *) lfirst(taskCell);
 
-		/*
+		LockShardDistributionMetadata(task->anchorShardId, ShareLock);
-		 * Only obtain metadata locks for modifications to allow reads to
-		 * proceed during shard copy.
-		 */
-		if (task->taskType == MODIFY_TASK &&
-			!TryLockShardDistributionMetadata(task->anchorShardId, ShareLock))
-		{
-			/*
-			 * We could error out immediately to give quick feedback to the
-			 * client, but this might complicate flow control and our default
-			 * behaviour during shard copy is to block.
-			 *
-			 * Block until the lock becomes available such that the next command
-			 * will likely succeed and use the serialization failure error code
-			 * to signal to the client that it should retry the current command.
-			 */
-			LockShardDistributionMetadata(task->anchorShardId, ShareLock);
-
-			ereport(ERROR, (errcode(ERRCODE_T_R_SERIALIZATION_FAILURE),
-							errmsg("prepared modifications cannot be executed on "
-								   "a shard while it is being copied")));
-		}
 	}
 }
 
@@ -457,7 +426,7 @@ CitusModifyBeginScan(CustomScanState *node, EState *estate, int eflags)
 	}
 
 	/* prevent concurrent placement changes */
-	ReacquireMetadataLocks(taskList);
+	AcquireMetadataLocks(taskList);
 
 	/* assign task placements */
 	workerJob->taskList = FirstReplicaAssignTaskList(taskList);

src/test/regress/expected/isolation_dml_vs_repair.out

@@ -78,7 +78,7 @@ step s2-invalidate-57638:
     UPDATE pg_dist_shard_placement SET shardstate = '3' WHERE shardid = (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_dml_vs_repair'::regclass) AND nodeport = 57638;
 
 step s1-display: 
-    SELECT * FROM test_dml_vs_repair WHERE test_id = 1;
+    SELECT * FROM test_dml_vs_repair WHERE test_id = 1 ORDER BY test_id;
 
 test_id        data           
 
@@ -90,7 +90,7 @@ step s2-revalidate-57638:
     UPDATE pg_dist_shard_placement SET shardstate = '1' WHERE shardid = (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_dml_vs_repair'::regclass) AND nodeport = 57638;
 
 step s1-display: 
-    SELECT * FROM test_dml_vs_repair WHERE test_id = 1;
+    SELECT * FROM test_dml_vs_repair WHERE test_id = 1 ORDER BY test_id;
 
 test_id        data           
 
@@ -122,15 +122,15 @@ step s2-commit:
     COMMIT;
 
 step s1-prepared-insertone: <... completed>
-error in steps s2-commit s1-prepared-insertone: ERROR:  prepared modifications cannot be executed on a shard while it is being copied
 step s2-invalidate-57638: 
     UPDATE pg_dist_shard_placement SET shardstate = '3' WHERE shardid = (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_dml_vs_repair'::regclass) AND nodeport = 57638;
 
 step s1-display: 
-    SELECT * FROM test_dml_vs_repair WHERE test_id = 1;
+    SELECT * FROM test_dml_vs_repair WHERE test_id = 1 ORDER BY test_id;
 
 test_id        data           
 
+1              1              
 1              1              
 step s2-invalidate-57637: 
     UPDATE pg_dist_shard_placement SET shardstate = '3' WHERE shardid = (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_dml_vs_repair'::regclass) AND nodeport = 57637;
@@ -139,10 +139,11 @@ step s2-revalidate-57638:
     UPDATE pg_dist_shard_placement SET shardstate = '1' WHERE shardid = (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_dml_vs_repair'::regclass) AND nodeport = 57638;
 
 step s1-display: 
-    SELECT * FROM test_dml_vs_repair WHERE test_id = 1;
+    SELECT * FROM test_dml_vs_repair WHERE test_id = 1 ORDER BY test_id;
 
 test_id        data           
 
+1              1              
 1              1              
 
 starting permutation: s2-invalidate-57637 s1-insertone s1-prepared-insertall s2-begin s2-repair s1-prepared-insertall s2-commit s2-invalidate-57638 s1-display s2-invalidate-57637 s2-revalidate-57638 s1-display
@@ -174,17 +175,18 @@ step s2-commit:
     COMMIT;
 
 step s1-prepared-insertall: <... completed>
-error in steps s2-commit s1-prepared-insertall: ERROR:  prepared modifications cannot be executed on a shard while it is being copied
 step s2-invalidate-57638: 
     UPDATE pg_dist_shard_placement SET shardstate = '3' WHERE shardid = (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_dml_vs_repair'::regclass) AND nodeport = 57638;
 
 step s1-display: 
-    SELECT * FROM test_dml_vs_repair WHERE test_id = 1;
+    SELECT * FROM test_dml_vs_repair WHERE test_id = 1 ORDER BY test_id;
 
 test_id        data           
 
 1              1              
 1              2              
+1              2              
+1              3              
 step s2-invalidate-57637: 
     UPDATE pg_dist_shard_placement SET shardstate = '3' WHERE shardid = (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_dml_vs_repair'::regclass) AND nodeport = 57637;
 
@@ -192,9 +194,11 @@ step s2-revalidate-57638:
     UPDATE pg_dist_shard_placement SET shardstate = '1' WHERE shardid = (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_dml_vs_repair'::regclass) AND nodeport = 57638;
 
 step s1-display: 
-    SELECT * FROM test_dml_vs_repair WHERE test_id = 1;
+    SELECT * FROM test_dml_vs_repair WHERE test_id = 1 ORDER BY test_id;
 
 test_id        data           
 
 1              1              
 1              2              
+1              2              
+1              3              

src/test/regress/specs/isolation_dml_vs_repair.spec

@@ -47,7 +47,7 @@ step "s1-prepared-insertall"
 
 step "s1-display"
 {
-    SELECT * FROM test_dml_vs_repair WHERE test_id = 1;
+    SELECT * FROM test_dml_vs_repair WHERE test_id = 1 ORDER BY test_id;
 }
 
 step "s1-commit"
@@ -102,8 +102,8 @@ permutation "s1-insertone" "s2-invalidate-57637" "s1-begin" "s1-insertall" "s2-r
 # verify that modifications wait for shard repair
 permutation "s2-invalidate-57637" "s2-begin" "s2-repair" "s1-insertone" "s2-commit" "s2-invalidate-57638" "s1-display" "s2-invalidate-57637" "s2-revalidate-57638" "s1-display"
 
-# verify that prepared plain modifications wait for shard repair (and then fail to avoid race)
+# verify that prepared plain modifications wait for shard repair
 permutation "s2-invalidate-57637" "s1-prepared-insertone" "s2-begin" "s2-repair" "s1-prepared-insertone" "s2-commit" "s2-invalidate-57638" "s1-display" "s2-invalidate-57637" "s2-revalidate-57638" "s1-display"
 
-# verify that prepared INSERT ... SELECT waits for shard repair  (and then fail to avoid race)
+# verify that prepared INSERT ... SELECT waits for shard repair
 permutation "s2-invalidate-57637" "s1-insertone" "s1-prepared-insertall" "s2-begin" "s2-repair" "s1-prepared-insertall" "s2-commit" "s2-invalidate-57638" "s1-display" "s2-invalidate-57637" "s2-revalidate-57638" "s1-display"