Preallocate spans before ExecutorRun

src/pg_tracing.c

@@ -918,6 +918,24 @@ is_new_lxid()
 #endif
 }
 
+/*
+ * Preallocate space in the shared current_trace_spans to have at least slots_needed slots
+ */
+static void
+preallocate_spans(int slots_needed)
+{
+	int			target_size = current_trace_spans->end + slots_needed;
+
+	Assert(slots_needed > 0);
+
+	if (target_size > current_trace_spans->max)
+	{
+		current_trace_spans->max = target_size;
+		current_trace_spans = repalloc(current_trace_spans, sizeof(pgTracingSpans) +
+									   current_trace_spans->max * sizeof(Span));
+	}
+}
+
 /*
  * Store a span in the current_trace_spans buffer
  */
@@ -1023,7 +1041,7 @@ add_span_to_shared_buffer_locked(const Span * span)
  */
 static void
 process_query_desc(const Traceparent * traceparent, const QueryDesc *queryDesc,
-				   int sql_error_code, TimestampTz parent_end)
+				   int sql_error_code, bool deparse_plan, TimestampTz parent_end)
 {
 	NodeCounters *node_counters = &peek_active_span()->node_counters;
 
@@ -1050,7 +1068,7 @@ process_query_desc(const Traceparent * traceparent, const QueryDesc *queryDesc,
 		uint64		query_id = queryDesc->plannedstmt->queryId;
 
 		process_planstate(traceparent, queryDesc, sql_error_code,
-						  pg_tracing_deparse_plan, parent_id, query_id,
+						  deparse_plan, parent_id, query_id,
 						  parent_start, parent_end,
 						  parameters_buffer, plan_name_buffer);
 	}
@@ -1340,11 +1358,18 @@ handle_pg_error(const Traceparent * traceparent,
 {
 	int			sql_error_code;
 	Span	   *span;
+	int			max_trace_spans = current_trace_spans->max;
 
 	sql_error_code = geterrcode();
 
 	if (queryDesc != NULL)
-		process_query_desc(traceparent, queryDesc, sql_error_code, span_end_time);
+
+		/*
+		 * Within error, we need to avoid any possible allocation as this
+		 * could be an out of memory error. deparse plan relies on allocation
+		 * through lcons so we explicitely disable it.
+		 */
+		process_query_desc(traceparent, queryDesc, sql_error_code, false, span_end_time);
 	span = peek_active_span();
 	while (span != NULL)
 	{
@@ -1355,6 +1380,7 @@ handle_pg_error(const Traceparent * traceparent,
 		/* Get the next span in the stack */
 		span = peek_active_span();
 	};
+	Assert(current_trace_spans->max == max_trace_spans);
 }
 
 /*
@@ -1811,6 +1837,7 @@ pg_tracing_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 cou
 	TimestampTz span_end_time;
 	Span	   *executor_run_span;
 	Traceparent *traceparent = &executor_traceparent;
+	int			num_nodes;
 
 	if (!pg_tracing_enabled(traceparent, nested_level) || queryDesc->totaltime == NULL)
 	{
@@ -1865,6 +1892,13 @@ pg_tracing_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 cou
 	if (pg_tracing_planstate_spans && queryDesc->planstate->instrument)
 		setup_ExecProcNode_override(pg_tracing_mem_ctx, queryDesc);
 
+	/*
+	 * Preallocate enough space in the current_trace_spans to process the
+	 * queryDesc in the error handler without doing new allocations
+	 */
+	num_nodes = number_nodes_from_planstate(queryDesc->planstate);
+	preallocate_spans(num_nodes);
+
 	nested_level++;
 	PG_TRY();
 	{
@@ -2045,7 +2079,7 @@ pg_tracing_ExecutorEnd(QueryDesc *queryDesc)
 	 * the planstate if needed
 	 */
 	parent_end = per_level_infos[nested_level].executor_end;
-	process_query_desc(traceparent, queryDesc, 0, parent_end);
+	process_query_desc(traceparent, queryDesc, 0, pg_tracing_deparse_plan, parent_end);
 
 	/* No need to increment nested level here */
 	if (prev_ExecutorEnd)

src/pg_tracing.h

@@ -384,6 +384,8 @@ extern int
 			get_parent_traced_planstate_index(int nested_level);
 extern TimestampTz
 			get_span_end_from_planstate(PlanState *planstate, TimestampTz plan_start, TimestampTz root_end);
+extern int
+			number_nodes_from_planstate(PlanState *planstate);
 
 /* pg_tracing_query_process.c */
 extern const char *normalise_query_parameters(const SpanContext * span_context, Span * span,

src/pg_tracing_planstate.c

@@ -11,6 +11,7 @@
 #include "postgres.h"
 
 #include "common/pg_prng.h"
+#include "nodes/nodeFuncs.h"
 #include "pg_tracing.h"
 #include "utils/ruleutils.h"
 #include "utils/timestamp.h"
@@ -692,3 +693,26 @@ process_planstate(const Traceparent * traceparent, const QueryDesc *queryDesc,
 	/* We can get rid of all the traced planstate for this level */
 	drop_traced_planstates();
 }
+
+static bool
+number_nodes_walker(PlanState *planstate, void *context)
+{
+	int		   *num_nodes = context;
+
+	if (planstate == NULL)
+		return false;
+	*num_nodes += 1;
+	return planstate_tree_walker(planstate, number_nodes_walker, num_nodes);
+}
+
+/*
+ * Give the number of nodes for the provided planstate
+ */
+int
+number_nodes_from_planstate(PlanState *planstate)
+{
+	int			num_nodes = 1;
+
+	planstate_tree_walker(planstate, number_nodes_walker, &num_nodes);
+	return num_nodes;
+}