From bdce1058dbeee94ef441aa02bc0eb2ab390e199a Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Fri, 5 Apr 2024 10:18:30 +0200
Subject: [PATCH] Filter parent state candidate using exec nested level

---
 src/pg_tracing.c           | 2 +-
 src/pg_tracing.h           | 2 +-
 src/pg_tracing_planstate.c | 6 +++---
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/pg_tracing.c b/src/pg_tracing.c
index 7b42e91..d42430b 100644
--- a/src/pg_tracing.c
+++ b/src/pg_tracing.c
@@ -1300,7 +1300,7 @@ 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();
+		TracedPlanstate *parent_traced_planstate = get_parent_traced_planstate(exec_nested_level);
 		Span	   *latest_top_span = get_latest_top_span(exec_nested_level - 1);
 
 		/*
diff --git a/src/pg_tracing.h b/src/pg_tracing.h
index 88c5a4b..eeae260 100644
--- a/src/pg_tracing.h
+++ b/src/pg_tracing.h
@@ -246,7 +246,7 @@ extern void
 extern void
 			cleanup_planstarts(void);
 extern TracedPlanstate *
-get_parent_traced_planstate(void);
+get_parent_traced_planstate(int nested_level);
 extern void
 			drop_traced_planstate(int exec_nested_level);
 
diff --git a/src/pg_tracing_planstate.c b/src/pg_tracing_planstate.c
index 920b8ab..b6d33c9 100644
--- a/src/pg_tracing_planstate.c
+++ b/src/pg_tracing_planstate.c
@@ -325,20 +325,20 @@ get_traced_planstate(PlanState *planstate)
  * Get possible parent traced_planstate
  */
 TracedPlanstate *
-get_parent_traced_planstate(void)
+get_parent_traced_planstate(int nested_level)
 {
 	TracedPlanstate *traced_planstate;
 
 	if (index_planstart >= 2)
 	{
 		traced_planstate = &traced_planstates[index_planstart - 2];
-		if (nodeTag(traced_planstate->planstate->plan) == T_ProjectSet)
+		if (traced_planstate->nested_level == nested_level && nodeTag(traced_planstate->planstate->plan) == T_ProjectSet)
 			return traced_planstate;
 	}
 	if (index_planstart >= 1)
 	{
 		traced_planstate = &traced_planstates[index_planstart - 1];
-		if (nodeTag(traced_planstate->planstate->plan) == T_Result)
+		if (traced_planstate->nested_level == nested_level && nodeTag(traced_planstate->planstate->plan) == T_Result)
 			return traced_planstate;
 	}
 	return NULL;