From ec771b4c08fcc0c9ec50e8418d911cb02bc929cb Mon Sep 17 00:00:00 2001 From: Anthonin Bonnefoy Date: Tue, 7 May 2024 16:39:51 +0200 Subject: [PATCH] Fix tracing of close cursor Closing a cursor will directly call ExecutorFinish. We need to create the top span for those cases. Also, as ExecutorEnd will be called, we will try to generate spans from planstate. However, traced planstates won't be available and representing planstate spans for cursor execution doesn't fit our current model. To fix this, we disable full query instrumentation for cursors. Fixes #10 --- expected/cursor.out | 6 ++++- sql/cursor.sql | 1 + src/pg_tracing.c | 58 ++++++++++++++++++++++++++++++++++----------- 3 files changed, 50 insertions(+), 15 deletions(-) diff --git a/expected/cursor.out b/expected/cursor.out index faa22a1..516f481 100644 --- a/expected/cursor.out +++ b/expected/cursor.out @@ -4,6 +4,7 @@ BEGIN; more than one row returned for \gset /*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000003-01'*/ FETCH BACKWARD 10 from c \gset more than one row returned for \gset +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000004-01'*/ CLOSE c; COMMIT; -- First declare -- +----------------------------------------+ @@ -142,7 +143,10 @@ SELECT span_type, span_operation, lvl from peek_ordered_spans where trace_id='00 ProcessUtility | ProcessUtility | 2 Select query | DECLARE c CURSOR FOR SELECT * from pg_tracing_test; | 3 Executor | ExecutorRun | 4 -(12 rows) + Utility query | CLOSE c; | 1 + ProcessUtility | ProcessUtility | 2 + Select query | DECLARE c CURSOR FOR SELECT * from pg_tracing_test; | 3 +(15 rows) -- Clean created spans CALL clean_spans(); diff --git a/sql/cursor.sql b/sql/cursor.sql index a5dec5b..ad19880 100644 --- a/sql/cursor.sql +++ b/sql/cursor.sql @@ -2,6 +2,7 @@ BEGIN; /*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ DECLARE c CURSOR FOR SELECT * from pg_tracing_test; /*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000002-01'*/ FETCH FORWARD 20 from c \gset /*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000003-01'*/ FETCH BACKWARD 10 from c \gset +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000004-01'*/ CLOSE c; COMMIT; -- First declare diff --git a/src/pg_tracing.c b/src/pg_tracing.c index d2329d2..0696086 100644 --- a/src/pg_tracing.c +++ b/src/pg_tracing.c @@ -220,6 +220,9 @@ static int plan_nested_level = 0; /* Current nesting depth of ExecutorRun+ProcessUtility calls */ int exec_nested_level = 0; +/* Whether we're in a cursor declaration */ +static bool within_declare_cursor = false; + /* Number of spans initially allocated at the start of a trace. */ static int pg_tracing_initial_allocated_spans = 25; @@ -1101,6 +1104,7 @@ cleanup_tracing(void) reset_trace_context(&root_trace_context); reset_trace_context(¤t_trace_context); max_nested_level = -1; + within_declare_cursor = false; current_trace_spans = NULL; per_level_buffers = NULL; cleanup_planstarts(); @@ -1712,9 +1716,17 @@ pg_tracing_ExecutorStart(QueryDesc *queryDesc, int eflags) /* * We only need full instrumentation if we generate spans from - * planstate + * planstate. + * + * If we're within a cursor declaration, we won't be able to generate + * spans from planstate due to cursor behaviour. Fetches will only + * call ExecutorRun and Closing the cursor will call ExecutorFinish + * and ExecutorEnd. We won't have the start of the planstate available + * and representing a node that has a fragmented execution doesn't fit + * our current model. So we disable full query instrumentation for + * cursors. */ - if (pg_tracing_planstate_spans) + if (pg_tracing_planstate_spans && !within_declare_cursor) queryDesc->instrument_options = INSTRUMENT_ALL; } @@ -1759,10 +1771,8 @@ pg_tracing_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 cou * ExecutorRun is the first hook called. Create the top span if it * doesn't already exist. */ - initialize_top_span(trace_context, queryDesc->operation, NULL, NULL, NULL, - queryDesc->sourceText, span_start_time, false); - - parent_id = get_latest_top_span(exec_nested_level)->span_id; + parent_id = initialize_top_span(trace_context, queryDesc->operation, NULL, NULL, NULL, + queryDesc->sourceText, span_start_time, false); /* Start ExecutorRun span as a new top span */ executor_run_span = allocate_new_top_span(); @@ -1780,12 +1790,12 @@ pg_tracing_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 cou add_parallel_context(trace_context, executor_run_span->span_id, per_level_buffers[exec_nested_level].query_id); - if (pg_tracing_planstate_spans) - - /* - * Setup ExecProcNode override to capture node start if planstate - * spans were requested - */ + /* + * Setup ExecProcNode override to capture node start if planstate + * spans were requested. If there's no query instrumentation, we can + * skip ExecProcNode override. + */ + if (pg_tracing_planstate_spans && queryDesc->planstate->instrument) setup_ExecProcNode_override(queryDesc); } @@ -1847,10 +1857,18 @@ pg_tracing_ExecutorFinish(QueryDesc *queryDesc) if (executor_sampled) { TimestampTz span_start_time = GetCurrentTimestamp(); - uint64 parent_id = get_latest_top_span(exec_nested_level)->span_id; + uint64 parent_id; + Span *executor_finish_span; + + /* + * When closing a cursor, only ExecutorFinish and ExecutorEnd will be + * called. Create the top span in this case. + */ + parent_id = initialize_top_span(trace_context, queryDesc->operation, NULL, NULL, NULL, + queryDesc->sourceText, span_start_time, false); /* Create ExecutorFinish as a new potential top span */ - Span *executor_finish_span = allocate_new_top_span(); + executor_finish_span = allocate_new_top_span(); begin_span(trace_context->traceparent.trace_id, executor_finish_span, SPAN_EXECUTOR_FINISH, @@ -1980,6 +1998,7 @@ pg_tracing_ProcessUtility(PlannedStmt *pstmt, const char *queryString, uint64 parent_id; TimestampTz span_end_time; TimestampTz span_start_time; + Node *parsetree; /* * Save track utility value since this value could be modified by a SET @@ -1994,6 +2013,17 @@ pg_tracing_ProcessUtility(PlannedStmt *pstmt, const char *queryString, *trace_context = root_trace_context; } + parsetree = pstmt->utilityStmt; + + /* + * Keep track if we're in a declare cursor as we want to disable query + * instrumentation in this case. + */ + if (nodeTag(parsetree) == T_DeclareCursorStmt) + within_declare_cursor = true; + else + within_declare_cursor = false; + if (!track_utility || !pg_tracing_enabled(trace_context, exec_nested_level)) { /*