diff --git a/src/pg_tracing.c b/src/pg_tracing.c index c499b22..f4e24dd 100644 --- a/src/pg_tracing.c +++ b/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) diff --git a/src/pg_tracing.h b/src/pg_tracing.h index b7db4d6..8500de3 100644 --- a/src/pg_tracing.h +++ b/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, diff --git a/src/pg_tracing_planstate.c b/src/pg_tracing_planstate.c index 01e7f98..5c04cac 100644 --- a/src/pg_tracing_planstate.c +++ b/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; +}