Use parent planstate as span end for top level spans (#5)

Top spans end are currently ended at the end of the parent's ExecutorRun while nodes will rely on query instrumentation. This can lead to a situation where a node with nested queries like ProjectSet is shorter than its children. We already use traced_planstate to correctly assign parent for a nested top span. We will keep track of the traced_planstate in the span and use it to when ending the top span.
DataDog · Apr 15, 2024 · 5572fdf · 5572fdf
1 parent bdce105
commit 5572fdf
Show file tree

Hide file tree

Showing 6 changed files with 144 additions and 12 deletions.
diff --git a/Makefile b/Makefile
@@ -19,7 +19,7 @@ OBJS = \
 	src/pg_tracing_span.o
 
 REGRESSCHECKS = setup utility select extended insert trigger sample \
-				planstate planstate_bitmap planstate_hash \
+				planstate planstate_bitmap planstate_hash planstate_projectset \
 				parallel subxact full_buffer \
 				nested wal cleanup
 REGRESSCHECKS_OPTS = --no-locale --encoding=UTF8 --temp-config pg_tracing.conf

diff --git a/expected/planstate_projectset.out b/expected/planstate_projectset.out
@@ -0,0 +1,55 @@
+/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ select information_schema._pg_expandarray('{0,1,2}'::int[]);
+ _pg_expandarray 
+-----------------
+ (0,1)
+ (1,2)
+ (2,3)
+(3 rows)
+
+SELECT span_operation, deparse_info, parameters, lvl from peek_ordered_spans where trace_id='00000000000000000000000000000001';
+                                                                                              span_operation                                                                                              | deparse_info |               parameters               | lvl 
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+----------------------------------------+-----
+ select information_schema._pg_expandarray($1::int[]);                                                                                                                                                    |              | $1 = '{0,1,2}'                         |   1
+ Planner                                                                                                                                                                                                  |              |                                        |   2
+ ExecutorRun                                                                                                                                                                                              |              |                                        |   2
+ ProjectSet                                                                                                                                                                                               |              |                                        |   3
+ Result                                                                                                                                                                                                   |              |                                        |   4
+ select $1[s], s operator(pg_catalog.-) pg_catalog.array_lower($1,$2) operator(pg_catalog.+) $3 from pg_catalog.generate_series(pg_catalog.array_lower($1,$4), pg_catalog.array_upper($1,$5), $6) as g(s) |              | $1 = 1, $2 = 1, $3 = 1, $4 = 1, $5 = 1 |   4
+ Planner                                                                                                                                                                                                  |              |                                        |   5
+(7 rows)
+
+-- +---------------------------------------------+
+-- | A: ProjectSet                               |
+-- ++-----------------+--+----------------+------+
+--  | B: Result       |  | C: TopSpan     |
+--  +-----------------+  +-+------------+-+
+--                         | D: Planner |
+--                         +------------+
+SELECT span_id AS span_a_id,
+        get_epoch(span_start) as span_a_start,
+        get_epoch(span_end) as span_a_end
+		from pg_tracing_peek_spans
+        where trace_id='00000000000000000000000000000001' AND span_operation='ProjectSet' \gset
+SELECT span_id AS span_b_id,
+        get_epoch(span_start) as span_b_start,
+        get_epoch(span_end) as span_b_end
+		from pg_tracing_peek_spans
+        where parent_id =:'span_a_id' and span_operation='Result' \gset
+SELECT span_id AS span_c_id,
+        get_epoch(span_start) as span_c_start,
+        get_epoch(span_end) as span_c_end
+		from pg_tracing_peek_spans
+        where parent_id =:'span_a_id' and span_type='Select query' \gset
+SELECT span_id AS span_d_id,
+        get_epoch(span_start) as span_d_start,
+        get_epoch(span_end) as span_d_end
+		from pg_tracing_peek_spans
+        where parent_id =:'span_c_id' and span_type='Planner' \gset
+SELECT :span_a_end >= :span_c_end as project_set_ends_after_nested_top_span;
+ project_set_ends_after_nested_top_span 
+----------------------------------------
+ t
+(1 row)
+
+-- Clean created spans
+CALL clean_spans();
diff --git a/sql/planstate_projectset.sql b/sql/planstate_projectset.sql
@@ -0,0 +1,36 @@
+/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ select information_schema._pg_expandarray('{0,1,2}'::int[]);
+SELECT span_operation, deparse_info, parameters, lvl from peek_ordered_spans where trace_id='00000000000000000000000000000001';
+
+-- +---------------------------------------------+
+-- | A: ProjectSet                               |
+-- ++-----------------+--+----------------+------+
+--  | B: Result       |  | C: TopSpan     |
+--  +-----------------+  +-+------------+-+
+--                         | D: Planner |
+--                         +------------+
+
+SELECT span_id AS span_a_id,
+        get_epoch(span_start) as span_a_start,
+        get_epoch(span_end) as span_a_end
+		from pg_tracing_peek_spans
+        where trace_id='00000000000000000000000000000001' AND span_operation='ProjectSet' \gset
+SELECT span_id AS span_b_id,
+        get_epoch(span_start) as span_b_start,
+        get_epoch(span_end) as span_b_end
+		from pg_tracing_peek_spans
+        where parent_id =:'span_a_id' and span_operation='Result' \gset
+SELECT span_id AS span_c_id,
+        get_epoch(span_start) as span_c_start,
+        get_epoch(span_end) as span_c_end
+		from pg_tracing_peek_spans
+        where parent_id =:'span_a_id' and span_type='Select query' \gset
+SELECT span_id AS span_d_id,
+        get_epoch(span_start) as span_d_start,
+        get_epoch(span_end) as span_d_end
+		from pg_tracing_peek_spans
+        where parent_id =:'span_c_id' and span_type='Planner' \gset
+
+SELECT :span_a_end >= :span_c_end as project_set_ends_after_nested_top_span;
+
+-- Clean created spans
+CALL clean_spans();
diff --git a/src/pg_tracing.c b/src/pg_tracing.c
@@ -712,9 +712,21 @@ end_nested_level(void)
 	for (int i = 0; i < top_spans->end; i++)
 	{
 		Span	   *top_span = top_spans->spans + i;
+		TimestampTz top_span_end = span_end_time;
 
 		if (!top_span->ended)
-			end_span(top_span, &span_end_time);
+		{
+			if (top_span->parent_planstate_index > -1)
+			{
+				/* We have a parent planstate, use it for the span end */
+				TracedPlanstate *traced_planstate = get_traced_planstate_from_index(top_span->parent_planstate_index);
+
+				/* Make sure to end instrumentation */
+				InstrEndLoop(traced_planstate->planstate->instrument);
+				top_span_end = get_span_end_from_planstate(traced_planstate->planstate, traced_planstate->node_start, span_end_time);
+			}
+			end_span(top_span, &top_span_end);
+		}
 	}
 
 finish:
@@ -1288,6 +1300,7 @@ begin_top_span(pgTracingTraceContext * trace_context, Span * top_span,
 	int			query_len;
 	const char *normalised_query;
 	uint64		parent_id;
+	int8		parent_planstate_index = -1;
 
 	/* in case of a cached plan, query might be unavailable */
 	if (query != NULL)
@@ -1300,14 +1313,23 @@ begin_top_span(pgTracingTraceContext * trace_context, Span * top_span,
 		parent_id = trace_context->traceparent.parent_id;
 	else
 	{
-		TracedPlanstate *parent_traced_planstate = get_parent_traced_planstate(exec_nested_level);
-		Span	   *latest_top_span = get_latest_top_span(exec_nested_level - 1);
+		TracedPlanstate *parent_traced_planstate = NULL;
+		Span	   *latest_top_span = NULL;
+
+		/*
+		 * We're in a nested level, check if we have a parent planstate and
+		 * use it as a parent span
+		 */
+		parent_planstate_index = get_parent_traced_planstate_index(exec_nested_level);
+		if (parent_planstate_index > -1)
+			parent_traced_planstate = get_traced_planstate_from_index(parent_planstate_index);
+		latest_top_span = get_latest_top_span(exec_nested_level - 1);
 
 		/*
 		 * Both planstate and previous top span can be the parent for the new
 		 * top span, we use the most recent as a parent
 		 */
-		if (parent_traced_planstate != NULL && parent_traced_planstate->node_start >= latest_top_span->start)
+		if (parent_traced_planstate != NULL && parent_traced_planstate->node_start > latest_top_span->start)
 			parent_id = parent_traced_planstate->span_id;
 		else
 			parent_id = latest_top_span->span_id;
@@ -1317,6 +1339,8 @@ begin_top_span(pgTracingTraceContext * trace_context, Span * top_span,
 			   command_type_to_span_type(commandType),
 			   NULL, parent_id,
 			   per_level_buffers[exec_nested_level].query_id, &start_time);
+	/* Keep track of the parent planstate index */
+	top_span->parent_planstate_index = parent_planstate_index;
 
 	if (IsParallelWorker())
 	{

diff --git a/src/pg_tracing.h b/src/pg_tracing.h
@@ -108,6 +108,8 @@ typedef struct Span
 								 * span's name */
 	bool		ended;			/* Track if the span was already ended.
 								 * Internal usage only */
+	int8		parent_planstate_index; /* Index to the parent planstate of
+										 * this span. Internal usage only */
 	int			be_pid;			/* Pid of the backend process */
 	Oid			user_id;		/* User ID when the span was created */
 	Oid			database_id;	/* Database ID where the span was created */
@@ -246,9 +248,13 @@ extern void
 extern void
 			cleanup_planstarts(void);
 extern TracedPlanstate *
-get_parent_traced_planstate(int nested_level);
+get_traced_planstate_from_index(int index);
+extern int
+			get_parent_traced_planstate_index(int nested_level);
 extern void
 			drop_traced_planstate(int exec_nested_level);
+extern TimestampTz
+			get_span_end_from_planstate(PlanState *planstate, TimestampTz plan_start, TimestampTz root_end);
 
 /* pg_tracing_query_process.c */
 extern const char *normalise_query_parameters(const JumbleState *jstate, const char *query,

diff --git a/src/pg_tracing_planstate.c b/src/pg_tracing_planstate.c
@@ -277,7 +277,7 @@ generate_span_from_custom_scan(CustomScanState *css, planstateTraceContext * pla
 /*
  * Get end time for a span node from the provided planstate.
  */
-static TimestampTz
+TimestampTz
 get_span_end_from_planstate(PlanState *planstate, TimestampTz plan_start, TimestampTz root_end)
 {
 	TimestampTz span_end_time;
@@ -322,26 +322,37 @@ get_traced_planstate(PlanState *planstate)
 }
 
 /*
- * Get possible parent traced_planstate
+ * Get traced_planstate from index
  */
 TracedPlanstate *
-get_parent_traced_planstate(int nested_level)
+get_traced_planstate_from_index(int index)
+{
+	Assert(index > -1);
+	Assert(index < max_planstart);
+	return &traced_planstates[index];
+}
+
+/*
+ * Get index in traced_planstates array of a possible parent traced_planstate
+ */
+int
+get_parent_traced_planstate_index(int nested_level)
 {
 	TracedPlanstate *traced_planstate;
 
 	if (index_planstart >= 2)
 	{
 		traced_planstate = &traced_planstates[index_planstart - 2];
 		if (traced_planstate->nested_level == nested_level && nodeTag(traced_planstate->planstate->plan) == T_ProjectSet)
-			return traced_planstate;
+			return index_planstart - 2;
 	}
 	if (index_planstart >= 1)
 	{
 		traced_planstate = &traced_planstates[index_planstart - 1];
 		if (traced_planstate->nested_level == nested_level && nodeTag(traced_planstate->planstate->plan) == T_Result)
-			return traced_planstate;
+			return index_planstart - 1;
 	}
-	return NULL;
+	return -1;
 }
 
 /*