Processes that are blocked on advisory locks show up in wait edges

Assign the distributed transaction id before trying to acquire the
executor advisory locks. This is useful to show this backend in citus
lock graphs (e.g., dump_global_wait_edges() and citus_lock_waits).

src/backend/distributed/executor/multi_router_executor.c

@@ -1351,6 +1351,13 @@ ExecuteModifyTasks(List *taskList, bool expectResults, ParamListInfo paramListIn
 		LockPartitionRelations(firstShardInterval->relationId, RowExclusiveLock);
 	}
 
+	/*
+	 * Assign the distributed transaction id before trying to acquire the
+	 * executor advisory locks. This is useful to show this backend in citus
+	 * lock graphs (e.g., dump_global_wait_edges() and citus_lock_waits).
+	 */
+	BeginOrContinueCoordinatedTransaction();
+
 	/*
 	 * Ensure that there are no concurrent modifications on the same
 	 * shards. In general, for DDL commands, we already obtained the
@@ -1362,8 +1369,6 @@ ExecuteModifyTasks(List *taskList, bool expectResults, ParamListInfo paramListIn
 	 */
 	AcquireExecutorMultiShardLocks(taskList);
 
-	BeginOrContinueCoordinatedTransaction();
-
 	if (MultiShardCommitProtocol == COMMIT_PROTOCOL_2PC ||
 		firstTask->replicationModel == REPLICATION_MODEL_2PC)
 	{

src/test/regress/expected/isolation_get_distributed_wait_queries.out

@@ -1,4 +1,4 @@
-Parsed test spec with 3 sessions
+Parsed test spec with 4 sessions
 
 starting permutation: s1-begin s1-update-ref-table-from-coordinator s2-start-session-level-connection s2-begin-on-worker s2-update-ref-table s3-select-distributed-waiting-queries s1-commit s2-commit-worker s2-stop-connection
 step s1-begin: 
@@ -747,3 +747,103 @@ stop_session_level_connection_to_node
 restore_isolation_tester_func
 
                
+
+starting permutation: s1-begin s1-update-on-the-coordinator s2-update-on-the-coordinator s3-select-distributed-waiting-queries s1-commit
+step s1-begin: 
+	BEGIN;
+
+step s1-update-on-the-coordinator: 
+	UPDATE tt1 SET value_1 = 4;
+
+step s2-update-on-the-coordinator: 
+	UPDATE tt1 SET value_1 = 4;
+ <waiting ...>
+step s3-select-distributed-waiting-queries: 
+	SELECT blocked_statement, current_statement_in_blocking_process, waiting_node_name, blocking_node_name, waiting_node_port, blocking_node_port FROM citus_lock_waits;
+
+blocked_statementcurrent_statement_in_blocking_processwaiting_node_nameblocking_node_namewaiting_node_portblocking_node_port
+
+
+	UPDATE tt1 SET value_1 = 4;
+
+	UPDATE tt1 SET value_1 = 4;
+coordinator_hostcoordinator_host57636          57636          
+step s1-commit: 
+	COMMIT;
+
+step s2-update-on-the-coordinator: <... completed>
+restore_isolation_tester_func
+
+               
+
+starting permutation: s1-start-session-level-connection s1-begin-on-worker s1-update-dist-table s4-start-session-level-connection s4-begin-on-worker s4-update-dist-table s3-select-distributed-waiting-queries s1-commit-worker s4-commit-worker s1-stop-connection s4-stop-connection
+step s1-start-session-level-connection: 
+	SELECT start_session_level_connection_to_node('localhost', 57637);
+
+start_session_level_connection_to_node
+
+               
+step s1-begin-on-worker: 
+	SELECT run_commands_on_session_level_connection_to_node('BEGIN');
+
+run_commands_on_session_level_connection_to_node
+
+               
+step s1-update-dist-table: 
+	SELECT run_commands_on_session_level_connection_to_node('UPDATE tt1 SET value_1 = 4');
+
+run_commands_on_session_level_connection_to_node
+
+               
+step s4-start-session-level-connection: 
+	SELECT start_session_level_connection_to_node('localhost', 57637);
+
+start_session_level_connection_to_node
+
+               
+step s4-begin-on-worker: 
+	SELECT run_commands_on_session_level_connection_to_node('BEGIN');
+
+run_commands_on_session_level_connection_to_node
+
+               
+step s4-update-dist-table: 
+	SELECT run_commands_on_session_level_connection_to_node('UPDATE tt1 SET value_1 = 5');
+ <waiting ...>
+step s3-select-distributed-waiting-queries: 
+	SELECT blocked_statement, current_statement_in_blocking_process, waiting_node_name, blocking_node_name, waiting_node_port, blocking_node_port FROM citus_lock_waits;
+
+blocked_statementcurrent_statement_in_blocking_processwaiting_node_nameblocking_node_namewaiting_node_portblocking_node_port
+
+UPDATE tt1 SET value_1 = 5UPDATE tt1 SET value_1 = 4localhost      localhost      57637          57637          
+step s1-commit-worker: 
+    SELECT run_commands_on_session_level_connection_to_node('COMMIT');
+
+run_commands_on_session_level_connection_to_node
+
+               
+step s4-update-dist-table: <... completed>
+run_commands_on_session_level_connection_to_node
+
+               
+step s4-commit-worker: 
+    SELECT run_commands_on_session_level_connection_to_node('COMMIT');
+
+run_commands_on_session_level_connection_to_node
+
+               
+step s1-stop-connection: 
+	SELECT stop_session_level_connection_to_node();
+
+stop_session_level_connection_to_node
+
+               
+step s4-stop-connection: 
+	SELECT stop_session_level_connection_to_node();
+
+stop_session_level_connection_to_node
+
+               
+restore_isolation_tester_func
+
+               

src/test/regress/specs/isolation_get_distributed_wait_queries.spec

@@ -120,6 +120,11 @@ step "s1-stop-connection"
 	SELECT stop_session_level_connection_to_node();
 }
 
+step "s1-update-on-the-coordinator"
+{
+	UPDATE tt1 SET value_1 = 4;
+}
+
 step "s1-commit"
 {
 	COMMIT;
@@ -127,6 +132,11 @@ step "s1-commit"
 
 session "s2"
 
+step "s2-begin"
+{
+	COMMIT;
+}
+
 step "s2-start-session-level-connection"
 {
 	SELECT start_session_level_connection_to_node('localhost', 57638);
@@ -172,6 +182,11 @@ step "s2-stop-connection"
 	SELECT stop_session_level_connection_to_node();
 }
 
+step "s2-update-on-the-coordinator"
+{
+	UPDATE tt1 SET value_1 = 4;
+}
+
 step "s2-commit-worker"
 {
     SELECT run_commands_on_session_level_connection_to_node('COMMIT');
@@ -184,6 +199,33 @@ step "s3-select-distributed-waiting-queries"
 	SELECT blocked_statement, current_statement_in_blocking_process, waiting_node_name, blocking_node_name, waiting_node_port, blocking_node_port FROM citus_lock_waits;
 }
 
+# session s1 and s4 executes the commands on the same worker node
+session "s4"
+
+step "s4-start-session-level-connection"
+{
+	SELECT start_session_level_connection_to_node('localhost', 57637);
+}
+
+step "s4-begin-on-worker"
+{
+	SELECT run_commands_on_session_level_connection_to_node('BEGIN');
+}
+
+step "s4-update-dist-table"
+{
+	SELECT run_commands_on_session_level_connection_to_node('UPDATE tt1 SET value_1 = 5');
+}
+
+step "s4-stop-connection"
+{
+	SELECT stop_session_level_connection_to_node();
+}
+step "s4-commit-worker"
+{
+    SELECT run_commands_on_session_level_connection_to_node('COMMIT');
+}
+
 permutation "s1-begin" "s1-update-ref-table-from-coordinator" "s2-start-session-level-connection" "s2-begin-on-worker" "s2-update-ref-table" "s3-select-distributed-waiting-queries" "s1-commit" "s2-commit-worker" "s2-stop-connection"
 permutation "s1-start-session-level-connection" "s1-begin-on-worker" "s1-update-ref-table" "s2-start-session-level-connection" "s2-begin-on-worker" "s2-update-ref-table" "s3-select-distributed-waiting-queries" "s1-commit-worker" "s2-commit-worker" "s1-stop-connection" "s2-stop-connection"
 permutation "s1-start-session-level-connection" "s1-begin-on-worker" "s1-update-dist-table" "s2-start-session-level-connection" "s2-begin-on-worker" "s2-update-dist-table" "s3-select-distributed-waiting-queries" "s1-commit-worker" "s2-commit-worker" "s1-stop-connection" "s2-stop-connection"
@@ -195,3 +237,9 @@ permutation "s1-start-session-level-connection" "s1-begin-on-worker" "s1-copy-to
 permutation "s1-start-session-level-connection" "s1-begin-on-worker" "s1-copy-to-ref-table" "s2-start-session-level-connection" "s2-begin-on-worker" "s2-copy-to-ref-table" "s3-select-distributed-waiting-queries" "s1-commit-worker" "s2-commit-worker" "s1-stop-connection" "s2-stop-connection"
 permutation "s1-start-session-level-connection" "s1-begin-on-worker" "s1-select-for-update" "s2-start-session-level-connection" "s2-begin-on-worker" "s2-update-ref-table" "s3-select-distributed-waiting-queries" "s1-commit-worker" "s2-commit-worker" "s1-stop-connection" "s2-stop-connection"
 permutation "s2-start-session-level-connection" "s2-begin-on-worker" "s2-insert-into-ref-table" "s1-begin" "s1-alter-table" "s3-select-distributed-waiting-queries" "s2-commit-worker" "s1-commit" "s2-stop-connection"
+
+# make sure that multi-shard modification queries
+# show up in the waiting processes even if they are
+# blocked on the same node
+permutation "s1-begin" "s1-update-on-the-coordinator" "s2-update-on-the-coordinator" "s3-select-distributed-waiting-queries" "s1-commit"
+permutation "s1-start-session-level-connection" "s1-begin-on-worker" "s1-update-dist-table" "s4-start-session-level-connection" "s4-begin-on-worker" "s4-update-dist-table" "s3-select-distributed-waiting-queries" "s1-commit-worker" "s4-commit-worker" "s1-stop-connection" "s4-stop-connection"