Don't walk into Gather subplans if there's no local scan

Parallel queries executed without the leader participation won't execute the sub nodes in the Gather/GatherMerge, which is represented by need_to_scan_locally=false. Since the nodes won't be executed, we won't have a traced plan and there's no value to generate spans from those nodes. Just skip the whole subtree in those cases. Fixes #54
DataDog · Sep 6, 2024 · 2b481ac · 2b481ac
1 parent d92b743
commit 2b481ac
Show file tree

Hide file tree

Showing 5 changed files with 83 additions and 2 deletions.
diff --git a/expected/parallel.out b/expected/parallel.out
@@ -62,6 +62,35 @@ SELECT count(*) from pg_tracing_consume_spans where span_operation='ExecutorRun'
      7
 (1 row)
 
+CALL clean_spans();
+-- Test leaderless parallel query
+set parallel_setup_cost=0;
+set parallel_tuple_cost=0;
+set min_parallel_table_scan_size=0;
+set max_parallel_workers_per_gather=2;
+set parallel_leader_participation=false;
+/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ select 1 from pg_class limit 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000001' ORDER BY lvl, span_operation;
+  span_type   |          span_operation           | lvl 
+--------------+-----------------------------------+-----
+ Select query | select $1 from pg_class limit $2; |   1
+ ExecutorRun  | ExecutorRun                       |   2
+ Planner      | Planner                           |   2
+ Limit        | Limit                             |   3
+ Select query | Worker 0                          |   3
+ Select query | Worker 1                          |   3
+ ExecutorRun  | ExecutorRun                       |   4
+ ExecutorRun  | ExecutorRun                       |   4
+ Gather       | Gather                            |   4
+ SeqScan      | SeqScan on pg_class               |   5
+ SeqScan      | SeqScan on pg_class               |   5
+(11 rows)
+
 -- Cleanup
 CALL clean_spans();
 CALL reset_settings();
diff --git a/expected/setup.out b/expected/setup.out
@@ -17,6 +17,10 @@ AS $$
     SET pg_tracing.caller_sample_rate TO DEFAULT;
     SET pg_tracing.track_utility TO DEFAULT;
     SET pg_tracing.max_parameter_size TO DEFAULT;
+    SET parallel_setup_cost TO DEFAULT;
+    SET parallel_tuple_cost TO DEFAULT;
+    SET min_parallel_table_scan_size TO DEFAULT;
+    SET max_parallel_workers_per_gather TO DEFAULT;
 $$;
 CREATE OR REPLACE PROCEDURE reset_pg_tracing_test_table() AS $$
 BEGIN

diff --git a/sql/parallel.sql b/sql/parallel.sql
@@ -30,6 +30,17 @@ SELECT trace_id from pg_tracing_peek_spans group by trace_id;
 
 -- Check number of executor spans
 SELECT count(*) from pg_tracing_consume_spans where span_operation='ExecutorRun';
+CALL clean_spans();
+
+-- Test leaderless parallel query
+set parallel_setup_cost=0;
+set parallel_tuple_cost=0;
+set min_parallel_table_scan_size=0;
+set max_parallel_workers_per_gather=2;
+set parallel_leader_participation=false;
+/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ select 1 from pg_class limit 1;
+
+SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000001' ORDER BY lvl, span_operation;
 
 -- Cleanup
 CALL clean_spans();

diff --git a/sql/setup.sql b/sql/setup.sql
@@ -19,6 +19,10 @@ AS $$
     SET pg_tracing.caller_sample_rate TO DEFAULT;
     SET pg_tracing.track_utility TO DEFAULT;
     SET pg_tracing.max_parameter_size TO DEFAULT;
+    SET parallel_setup_cost TO DEFAULT;
+    SET parallel_tuple_cost TO DEFAULT;
+    SET min_parallel_table_scan_size TO DEFAULT;
+    SET max_parallel_workers_per_gather TO DEFAULT;
 $$;
 
 CREATE OR REPLACE PROCEDURE reset_pg_tracing_test_table() AS $$

diff --git a/src/pg_tracing_planstate.c b/src/pg_tracing_planstate.c
@@ -279,6 +279,36 @@ create_spans_from_custom_scan(CustomScanState *css, planstateTraceContext * plan
 	return last_end;
 }
 
+/*
+ * Check if the subplan will run or not
+ *
+ * With parallel queries, if leader participation is disabled, the leader won't execute the subplans
+ */
+static bool
+is_subplan_executed(PlanState *planstate)
+{
+	if (planstate == NULL)
+		return false;
+
+	switch (nodeTag(planstate))
+	{
+		case T_GatherState:
+			{
+				GatherState *splanstate = (GatherState *) planstate;
+
+				return splanstate->need_to_scan_locally;
+			}
+		case T_GatherMerge:
+			{
+				GatherMergeState *splanstate = (GatherMergeState *) planstate;
+
+				return splanstate->need_to_scan_locally;
+			}
+		default:
+			return true;
+	}
+}
+
 /*
  * Walk through the planstate tree generating a node span for each node.
  */
@@ -295,6 +325,7 @@ create_spans_from_planstate(PlanState *planstate, planstateTraceContext * planst
 	TimestampTz span_end;
 	TimestampTz child_end = 0;
 	bool		haschildren = false;
+	bool		subplan_executed = true;
 
 	/* The node was never executed, skip it */
 	if (planstate->instrument == NULL)
@@ -379,11 +410,13 @@ create_spans_from_planstate(PlanState *planstate, planstateTraceContext * planst
 	if (haschildren && planstateTraceContext->deparse_ctx)
 		planstateTraceContext->ancestors = lcons(planstate->plan, planstateTraceContext->ancestors);
 
+    /* We only need to walk the subplans if they are executed */
+	subplan_executed = is_subplan_executed(planstate);
 	/* Walk the outerplan */
-	if (outerPlanState(planstate))
+	if (outerPlanState(planstate) && subplan_executed)
 		create_spans_from_planstate(outerPlanState(planstate), planstateTraceContext, span_id, query_id, span_start, root_end, latest_end);
 	/* Walk the innerplan */
-	if (innerPlanState(planstate))
+	if (innerPlanState(planstate) && subplan_executed)
 		create_spans_from_planstate(innerPlanState(planstate), planstateTraceContext, span_id, query_id, span_start, root_end, latest_end);
 
 	/* Handle init plans */