diff --git a/Makefile b/Makefile index 558e85e..ddb4e02 100644 --- a/Makefile +++ b/Makefile @@ -45,7 +45,7 @@ endif # PG 18 contains additional psql metacommand to test extended protocol ifeq ($(shell test $(PG_VERSION) -ge 18; echo $$?),0) -REGRESSCHECKS += psql_extended +REGRESSCHECKS += psql_extended psql_extended_tx endif REGRESSCHECKS += sample planstate planstate_bitmap planstate_hash \ diff --git a/expected/nested.out b/expected/nested.out index dcb6be3..6846185 100644 --- a/expected/nested.out +++ b/expected/nested.out @@ -169,15 +169,17 @@ LANGUAGE sql IMMUTABLE; (1 row) select span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000056'; - span_operation | lvl ------------------------------------------+----- - select test_immutable_function($1); | 1 - Planner | 2 - SELECT oid from pg_class where oid = a; | 3 - Planner | 4 - ExecutorRun | 2 - Result | 3 -(6 rows) + span_operation | lvl +----------------------------------------------------+----- + select test_immutable_function($1); | 1 + Planner | 2 + SELECT oid from pg_class where oid = a; | 3 + Planner | 4 + ExecutorRun | 4 + IndexOnlyScan using pg_class_oid_index on pg_class | 5 + ExecutorRun | 2 + Result | 3 +(8 rows) -- Create function with generate series CREATE OR REPLACE FUNCTION test_generate_series(IN anyarray, OUT x anyelement) diff --git a/expected/psql_extended.out b/expected/psql_extended.out index bb417a5..877157b 100644 --- a/expected/psql_extended.out +++ b/expected/psql_extended.out @@ -1,4 +1,4 @@ --- Simple query with extended protocol +-- Select with extended protocol /*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ SELECT $1, $2 \parse stmt1 \bind_named stmt1 1 2 \g ?column? | ?column? @@ -16,97 +16,14 @@ SELECT trace_id, span_type, span_operation, parameters, lvl FROM peek_ordered_sp 00000000000000000000000000000001 | Result | Result | | 3 (5 rows) -BEGIN; -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000002-0000000000000001-01'*/ select $1 \parse '' -\bind_named '' 1 \g +WITH max_end AS (select max(span_end) from pg_tracing_peek_spans) +SELECT span_end = max_end.max from pg_tracing_peek_spans, max_end + where span_type='Select query' AND parameters IS NOT NULL; ?column? ---------- - 1 + t (1 row) -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000003-0000000000000001-01'*/ select $1, $2 \parse '' -\bind_named '' 1 2 \g - ?column? | ?column? -----------+---------- - 1 | 2 -(1 row) - -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000004-0000000000000001-01'*/ select $1, $2, $3 \parse '' -\bind_named '' 1 2 3 \g - ?column? | ?column? | ?column? -----------+----------+---------- - 1 | 2 | 3 -(1 row) - -COMMIT; -SELECT trace_id, span_type, span_operation, parameters, lvl FROM peek_ordered_spans; - trace_id | span_type | span_operation | parameters | lvl -----------------------------------+------------------+-------------------+------------+----- - 00000000000000000000000000000001 | Select query | SELECT $1, $2 | | 1 - 00000000000000000000000000000001 | Select query | SELECT $1, $2 | {1,2} | 1 - 00000000000000000000000000000001 | Planner | Planner | | 2 - 00000000000000000000000000000001 | ExecutorRun | ExecutorRun | | 2 - 00000000000000000000000000000001 | Result | Result | | 3 - 00000000000000000000000000000002 | TransactionBlock | TransactionBlock | | 1 - 00000000000000000000000000000002 | Select query | select $1 | {1} | 2 - 00000000000000000000000000000002 | Planner | Planner | | 3 - 00000000000000000000000000000002 | ExecutorRun | ExecutorRun | | 3 - 00000000000000000000000000000002 | Result | Result | | 4 - 00000000000000000000000000000003 | Select query | SELECT $1, $2 | {1,2} | 2 - 00000000000000000000000000000003 | Planner | Planner | | 3 - 00000000000000000000000000000003 | ExecutorRun | ExecutorRun | | 3 - 00000000000000000000000000000003 | Result | Result | | 4 - 00000000000000000000000000000004 | Select query | select $1, $2, $3 | {1,2,3} | 2 - 00000000000000000000000000000004 | Planner | Planner | | 3 - 00000000000000000000000000000004 | ExecutorRun | ExecutorRun | | 3 - 00000000000000000000000000000004 | Result | Result | | 4 -(18 rows) - -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000005-0000000000000001-01'*/ BEGIN; -select $1 \parse '' -\bind_named '' 1 \g - ?column? ----------- - 1 -(1 row) - -select $1, $2 \parse '' -\bind_named '' 1 2 \g - ?column? | ?column? -----------+---------- - 1 | 2 -(1 row) - -select $1, $2, $3 \parse '' -\bind_named '' 1 2 3 \g - ?column? | ?column? | ?column? -----------+----------+---------- - 1 | 2 | 3 -(1 row) - -COMMIT; -SELECT span_type, span_operation, parameters, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000005'; - span_type | span_operation | parameters | lvl -------------------+-------------------+------------+----- - TransactionBlock | TransactionBlock | | 1 - Utility query | BEGIN; | | 2 - ProcessUtility | ProcessUtility | | 3 - Select query | select $1 | {1} | 2 - Planner | Planner | | 3 - ExecutorRun | ExecutorRun | | 3 - Result | Result | | 4 - Select query | SELECT $1, $2 | {1,2} | 2 - Planner | Planner | | 3 - ExecutorRun | ExecutorRun | | 3 - Result | Result | | 4 - Select query | select $1, $2, $3 | {1,2,3} | 2 - Planner | Planner | | 3 - ExecutorRun | ExecutorRun | | 3 - Result | Result | | 4 - Utility query | COMMIT; | | 2 - ProcessUtility | ProcessUtility | | 3 -(17 rows) - -- Cleanup -CALL clean_spans(); CALL reset_settings(); +CALL clean_spans(); diff --git a/expected/psql_extended_tx.out b/expected/psql_extended_tx.out new file mode 100644 index 0000000..0911a19 --- /dev/null +++ b/expected/psql_extended_tx.out @@ -0,0 +1,248 @@ +-- Only trace individual statements using unnamed prepared statement +BEGIN; +/*traceparent='00-00000000000000000000000000000002-0000000000000001-01'*/ select $1 \parse '' +\bind_named '' 1 \g + ?column? +---------- + 1 +(1 row) + +/*traceparent='00-00000000000000000000000000000003-0000000000000001-01'*/ select $1, $2 \parse '' +\bind_named '' 1 2 \g + ?column? | ?column? +----------+---------- + 1 | 2 +(1 row) + +/*traceparent='00-00000000000000000000000000000004-0000000000000001-01'*/ select $1, $2, $3 \parse '' +\bind_named '' 1 2 3 \g + ?column? | ?column? | ?column? +----------+----------+---------- + 1 | 2 | 3 +(1 row) + +COMMIT; +SELECT trace_id, span_type, span_operation, parameters, lvl FROM peek_ordered_spans; + trace_id | span_type | span_operation | parameters | lvl +----------------------------------+------------------+-------------------+------------+----- + 00000000000000000000000000000002 | TransactionBlock | TransactionBlock | | 1 + 00000000000000000000000000000002 | Select query | select $1 | {1} | 2 + 00000000000000000000000000000002 | Planner | Planner | | 3 + 00000000000000000000000000000002 | ExecutorRun | ExecutorRun | | 3 + 00000000000000000000000000000002 | Result | Result | | 4 + 00000000000000000000000000000003 | Select query | select $1, $2 | {1,2} | 2 + 00000000000000000000000000000003 | Planner | Planner | | 3 + 00000000000000000000000000000003 | ExecutorRun | ExecutorRun | | 3 + 00000000000000000000000000000003 | Result | Result | | 4 + 00000000000000000000000000000004 | Select query | select $1, $2, $3 | {1,2,3} | 2 + 00000000000000000000000000000004 | Planner | Planner | | 3 + 00000000000000000000000000000004 | ExecutorRun | ExecutorRun | | 3 + 00000000000000000000000000000004 | Result | Result | | 4 +(13 rows) + +CALL clean_spans(); +-- Mix begin in simple protocol with extended protocol usage +/*traceparent='00-00000000000000000000000000000005-0000000000000001-01'*/ BEGIN; +select $1 \parse '' +\bind_named '' 1 \g + ?column? +---------- + 1 +(1 row) + +select $1, $2 \parse '' +\bind_named '' 1 2 \g + ?column? | ?column? +----------+---------- + 1 | 2 +(1 row) + +select $1, $2, $3 \parse '' +\bind_named '' 1 2 3 \g + ?column? | ?column? | ?column? +----------+----------+---------- + 1 | 2 | 3 +(1 row) + +COMMIT; +SELECT trace_id, span_type, span_operation, parameters, lvl FROM peek_ordered_spans; + trace_id | span_type | span_operation | parameters | lvl +----------------------------------+------------------+-------------------+------------+----- + 00000000000000000000000000000005 | TransactionBlock | TransactionBlock | | 1 + 00000000000000000000000000000005 | Utility query | BEGIN; | | 2 + 00000000000000000000000000000005 | ProcessUtility | ProcessUtility | | 3 + 00000000000000000000000000000005 | Select query | select $1 | {1} | 2 + 00000000000000000000000000000005 | Planner | Planner | | 3 + 00000000000000000000000000000005 | ExecutorRun | ExecutorRun | | 3 + 00000000000000000000000000000005 | Result | Result | | 4 + 00000000000000000000000000000005 | Select query | select $1, $2 | {1,2} | 2 + 00000000000000000000000000000005 | Planner | Planner | | 3 + 00000000000000000000000000000005 | ExecutorRun | ExecutorRun | | 3 + 00000000000000000000000000000005 | Result | Result | | 4 + 00000000000000000000000000000005 | Select query | select $1, $2, $3 | {1,2,3} | 2 + 00000000000000000000000000000005 | Planner | Planner | | 3 + 00000000000000000000000000000005 | ExecutorRun | ExecutorRun | | 3 + 00000000000000000000000000000005 | Result | Result | | 4 + 00000000000000000000000000000005 | Utility query | COMMIT; | | 2 + 00000000000000000000000000000005 | ProcessUtility | ProcessUtility | | 3 +(17 rows) + +CALL clean_spans(); +-- Test using extended protocol for utility statement +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000006-0000000000000001-01'*/ BEGIN \parse begin_stmt +\bind_named begin_stmt \g +SELECT 1 \parse stmt_mix_1 +\bind_named stmt_mix_1 \g + ?column? +---------- + 1 +(1 row) + +SELECT 1, 2 \parse stmt_mix_2 +\bind_named stmt_mix_2 \g + ?column? | ?column? +----------+---------- + 1 | 2 +(1 row) + +COMMIT \parse commit_stmt +\bind_named commit_stmt \g +SELECT span_type, span_operation, parameters, lvl FROM peek_ordered_spans + where trace_id='00000000000000000000000000000006' AND lvl <= 2; + span_type | span_operation | parameters | lvl +------------------+------------------+------------+----- + Utility query | BEGIN | | 1 + TransactionBlock | TransactionBlock | | 1 + Utility query | BEGIN | | 2 + Select query | SELECT $1 | {1} | 2 + Select query | SELECT $1, $2 | {1,2} | 2 + Utility query | COMMIT | | 2 +(6 rows) + +SELECT get_epoch(span_end) as span_select_1_end + from pg_tracing_peek_spans where span_operation='SELECT $1' and trace_id='00000000000000000000000000000006' \gset +SELECT get_epoch(span_start) as span_select_2_start, + get_epoch(span_end) as span_select_2_end + from pg_tracing_peek_spans where span_operation='SELECT $1, $2' and trace_id='00000000000000000000000000000006' \gset +SELECT get_epoch(span_start) as span_commit_start + from pg_tracing_peek_spans where span_operation='COMMIT' and trace_id='00000000000000000000000000000006' \gset +SELECT :span_select_1_end <= :span_select_2_start, + :span_select_2_end <= :span_commit_start; + ?column? | ?column? +----------+---------- + t | t +(1 row) + +CALL clean_spans(); +-- Test with extended protocol while reusing prepared stmts +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ BEGIN; +\bind_named stmt_mix_1 \g + ?column? +---------- + 1 +(1 row) + +\bind_named stmt_mix_2 \g + ?column? | ?column? +----------+---------- + 1 | 2 +(1 row) + +SELECT 1, 2, 3 \parse stmt_mix_3 +\bind_named stmt_mix_3 \g + ?column? | ?column? | ?column? +----------+----------+---------- + 1 | 2 | 3 +(1 row) + +\bind_named commit_stmt \g +SELECT span_type, span_operation, parameters, lvl FROM peek_ordered_spans; + span_type | span_operation | parameters | lvl +------------------+-------------------+------------+----- + TransactionBlock | TransactionBlock | | 1 + Utility query | BEGIN; | | 2 + ProcessUtility | ProcessUtility | | 3 + Select query | SELECT 1 | | 2 + ExecutorRun | ExecutorRun | | 3 + Result | Result | | 4 + Select query | SELECT 1, 2 | | 2 + ExecutorRun | ExecutorRun | | 3 + Result | Result | | 4 + Select query | SELECT $1, $2, $3 | {1,2,3} | 2 + Planner | Planner | | 3 + ExecutorRun | ExecutorRun | | 3 + Result | Result | | 4 + Utility query | COMMIT | | 2 + ProcessUtility | ProcessUtility | | 3 +(15 rows) + +SELECT get_epoch(span_end) as span_select_1_end + from pg_tracing_peek_spans where span_operation='SELECT 1' and trace_id='00000000000000000000000000000001' \gset +SELECT get_epoch(span_start) as span_select_2_start, + get_epoch(span_end) as span_select_2_end + from pg_tracing_peek_spans where span_operation='SELECT 1, 2' and trace_id='00000000000000000000000000000001' \gset +SELECT get_epoch(span_start) as span_select_3_start, + get_epoch(span_end) as span_select_3_end + from pg_tracing_peek_spans where span_operation='SELECT $1, $2, $3' and trace_id='00000000000000000000000000000001' \gset +SELECT get_epoch(span_start) as span_commit_start + from pg_tracing_peek_spans where span_operation='COMMIT' and trace_id='00000000000000000000000000000001' \gset +SELECT :span_select_1_end <= :span_select_2_start, + :span_select_2_end <= :span_select_3_start, + :span_select_3_end <= :span_commit_start; + ?column? | ?column? | ?column? +----------+----------+---------- + t | t | t +(1 row) + +DEALLOCATE ALL; +CALL clean_spans(); +-- Test extended protocol while tracing everything +SET pg_tracing.sample_rate = 1.0; +BEGIN \parse begin_stmt +\bind_named begin_stmt \g +SELECT 1 \parse stmt_mix_1 +\bind_named stmt_mix_1 \g + ?column? +---------- + 1 +(1 row) + +SELECT 1, 2 \parse stmt_mix_2 +\bind_named stmt_mix_2 \g + ?column? | ?column? +----------+---------- + 1 | 2 +(1 row) + +COMMIT \parse commit_stmt +\bind_named commit_stmt \g +SELECT span_type, span_operation, parameters, lvl FROM peek_ordered_spans + WHERE lvl <= 2; + span_type | span_operation | parameters | lvl +------------------+------------------+------------+----- + Utility query | BEGIN | | 1 + TransactionBlock | TransactionBlock | | 1 + Utility query | BEGIN | | 2 + Select query | SELECT $1 | {1} | 2 + Select query | SELECT $1, $2 | {1,2} | 2 + Utility query | COMMIT | | 2 +(6 rows) + +SELECT get_epoch(span_end) as span_select_1_end + from pg_tracing_peek_spans where span_operation='SELECT $1' \gset +SELECT get_epoch(span_start) as span_select_2_start, + get_epoch(span_end) as span_select_2_end + from pg_tracing_peek_spans where span_operation='SELECT $1, $2' \gset +SELECT get_epoch(span_start) as span_commit_start + from pg_tracing_peek_spans where span_operation='COMMIT' \gset +SELECT :span_select_1_end <= :span_select_2_start, + :span_select_2_end <= :span_commit_start; + ?column? | ?column? +----------+---------- + t | t +(1 row) + +CALL clean_spans(); +-- Cleanup +CALL reset_settings(); +CALL clean_spans(); diff --git a/expected/utility.out b/expected/utility.out index 6a0318b..1b60f42 100644 --- a/expected/utility.out +++ b/expected/utility.out @@ -69,7 +69,13 @@ CREATE VIEW peek_spans_with_level AS ) SELECT * FROM list_trace_spans; -- Create utility view to keep order stable CREATE VIEW peek_ordered_spans AS - SELECT * FROM peek_spans_with_level order by span_start, lvl, span_end, deparse_info; + WITH oldest_start AS ( + SELECT min(span_start) as min_start + FROM pg_tracing_peek_spans + ) select *, + extract(MICROSECONDS FROM age(span_start, oldest_start.min_start)) as us_start, + extract(MICROSECONDS FROM age(span_end, oldest_start.min_start)) as us_end + FROM peek_spans_with_level, oldest_start order by span_start, lvl, span_end, deparse_info; -- Column type to convert json to record create type span_json as ("traceId" text, "parentSpanId" text, "spanId" text, name text, "startTimeUnixNano" text, "endTimeUnixNano" text, attributes jsonb, kind int, status jsonb); CREATE VIEW peek_json_spans AS diff --git a/sql/psql_extended.sql b/sql/psql_extended.sql index 4470a5d..c9c0ba9 100644 --- a/sql/psql_extended.sql +++ b/sql/psql_extended.sql @@ -1,31 +1,13 @@ --- Simple query with extended protocol +-- Select with extended protocol /*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ SELECT $1, $2 \parse stmt1 \bind_named stmt1 1 2 \g SELECT trace_id, span_type, span_operation, parameters, lvl FROM peek_ordered_spans; -BEGIN; -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000002-0000000000000001-01'*/ select $1 \parse '' -\bind_named '' 1 \g -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000003-0000000000000001-01'*/ select $1, $2 \parse '' -\bind_named '' 1 2 \g -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000004-0000000000000001-01'*/ select $1, $2, $3 \parse '' -\bind_named '' 1 2 3 \g -COMMIT; - -SELECT trace_id, span_type, span_operation, parameters, lvl FROM peek_ordered_spans; - -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000005-0000000000000001-01'*/ BEGIN; -select $1 \parse '' -\bind_named '' 1 \g -select $1, $2 \parse '' -\bind_named '' 1 2 \g -select $1, $2, $3 \parse '' -\bind_named '' 1 2 3 \g -COMMIT; - -SELECT span_type, span_operation, parameters, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000005'; +WITH max_end AS (select max(span_end) from pg_tracing_peek_spans) +SELECT span_end = max_end.max from pg_tracing_peek_spans, max_end + where span_type='Select query' AND parameters IS NOT NULL; -- Cleanup -CALL clean_spans(); CALL reset_settings(); +CALL clean_spans(); diff --git a/sql/utility.sql b/sql/utility.sql index d41d159..1ddefc6 100644 --- a/sql/utility.sql +++ b/sql/utility.sql @@ -78,7 +78,13 @@ CREATE VIEW peek_spans_with_level AS -- Create utility view to keep order stable CREATE VIEW peek_ordered_spans AS - SELECT * FROM peek_spans_with_level order by span_start, lvl, span_end, deparse_info; + WITH oldest_start AS ( + SELECT min(span_start) as min_start + FROM pg_tracing_peek_spans + ) select *, + extract(MICROSECONDS FROM age(span_start, oldest_start.min_start)) as us_start, + extract(MICROSECONDS FROM age(span_end, oldest_start.min_start)) as us_end + FROM peek_spans_with_level, oldest_start order by span_start, lvl, span_end, deparse_info; -- Column type to convert json to record create type span_json as ("traceId" text, "parentSpanId" text, "spanId" text, name text, "startTimeUnixNano" text, "endTimeUnixNano" text, attributes jsonb, kind int, status jsonb); diff --git a/src/pg_tracing.c b/src/pg_tracing.c index 2a43348..c499b22 100644 --- a/src/pg_tracing.c +++ b/src/pg_tracing.c @@ -232,6 +232,9 @@ static Span commit_span; /* Tx block span used to represent explicit transaction block */ static Span tx_block_span; +/* Candidate span end time for top span */ +static TimestampTz top_span_end_candidate; + /* * query id at level 0 of the current tracing. * Used to assign queryId to the commit span @@ -1203,8 +1206,8 @@ extract_trace_context(Traceparent * traceparent, ParseState *pstate, * try to apply sample rate and exit. */ statement_start_ts = GetCurrentStatementStartTimestamp(); - if (traceparent->sampled == 0 - && last_statement_check_for_sampling == statement_start_ts) + if (!traceparent->sampled && + last_statement_check_for_sampling == statement_start_ts) goto cleanup; /* First time we see this statement, save the time */ last_statement_check_for_sampling = statement_start_ts; @@ -1226,6 +1229,14 @@ extract_trace_context(Traceparent * traceparent, ParseState *pstate, reset_traceparent(traceparent); cleanup: + /* Check if the current transaction is sampled use it if it's the case */ + if (!traceparent->sampled && tx_traceparent.sampled) + { + /* We have a sampled transaction, use the tx's traceparent */ + Assert(!traceid_zero(tx_traceparent.trace_id)); + *traceparent = tx_traceparent; + } + /* No matter what happens, we want to update the latest_lxid seen */ update_latest_lxid(); } @@ -1260,6 +1271,7 @@ cleanup_tracing(void) within_declare_cursor = false; current_trace_spans = NULL; per_level_infos = NULL; + top_span_end_candidate = 0; cleanup_planstarts(); cleanup_active_spans(); } @@ -1339,7 +1351,7 @@ handle_pg_error(const Traceparent * traceparent, /* Assign the error code to the latest top span */ span->sql_error_code = sql_error_code; /* End and store the span */ - pop_active_span(&span_end_time); + pop_and_store_active_span(span_end_time); /* Get the next span in the stack */ span = peek_active_span(); }; @@ -1418,7 +1430,8 @@ end_nested_level(const TimestampTz *input_span_end_time) InstrEndLoop(traced_planstate->planstate->instrument); span_end_time = get_span_end_from_planstate(traced_planstate->planstate, traced_planstate->node_start, span_end_time); } - pop_active_span(&span_end_time); + + pop_and_store_active_span(span_end_time); span = peek_active_span(); } } @@ -1477,14 +1490,14 @@ initialise_span_context(SpanContext * span_context, } static bool -should_start_tx_block(const Node *utilityStmt) +should_start_tx_block(const Node *utilityStmt, HookType hook_type) { TransactionStmt *stmt; - if (GetCurrentTransactionStartTimestamp() != GetCurrentStatementStartTimestamp()) + if (hook_type == HOOK_PARSE && GetCurrentTransactionStartTimestamp() != GetCurrentStatementStartTimestamp()) /* There's an ongoing tx block, we can create the matching span */ return true; - if (utilityStmt != NULL && nodeTag(utilityStmt) == T_TransactionStmt) + if (hook_type == HOOK_UTILITY && utilityStmt != NULL && nodeTag(utilityStmt) == T_TransactionStmt) { stmt = (TransactionStmt *) utilityStmt; /* If we have an explicit BEGIN statement, start a tx block */ @@ -1497,7 +1510,7 @@ should_start_tx_block(const Node *utilityStmt) * Start a tx_block span if we have an explicit BEGIN utility statement */ static void -start_tx_block_span(const Node *utilityStmt, SpanContext * span_context) +start_tx_block_span(const Node *utilityStmt, SpanContext * span_context, HookType hook_type) { if (nested_level > 0) return; @@ -1506,7 +1519,7 @@ start_tx_block_span(const Node *utilityStmt, SpanContext * span_context) /* We already have an ongoing tx_block span */ return; - if (!should_start_tx_block(utilityStmt)) + if (!should_start_tx_block(utilityStmt, hook_type)) /* Not a candidate to start tx block */ return; @@ -1540,7 +1553,6 @@ pg_tracing_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jst { bool new_lxid = is_new_lxid(); bool is_root_level = nested_level == 0; - bool new_local_transaction_statement; Traceparent *traceparent = &parse_traceparent; SpanContext span_context; @@ -1578,17 +1590,16 @@ pg_tracing_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jst /* Evaluate if query is sampled or not */ extract_trace_context(traceparent, pstate, query->queryId); - new_local_transaction_statement = is_root_level && !new_lxid; - if (!traceparent->sampled && (new_local_transaction_statement && tx_traceparent.sampled)) - { - /* - * We're in the same traced local transaction, use the trace context - * extracted at the start of this transaction - */ - Assert(!traceid_zero(tx_traceparent.trace_id)); - *traceparent = tx_traceparent; - } + /* + * Keep track of the start of the last parse at root level. When explicit + * tx is used with extended protocol, the end of the previous command + * happens when the next command is bound, creating overlapping spans. To + * avoid this, we want to end the previous command as soon as the new + * command is parsed or bound + */ + if (nested_level == 0) + top_span_end_candidate = GetCurrentStatementStartTimestamp(); if (!traceparent->sampled) /* Query is not sampled, nothing to do. */ @@ -1603,7 +1614,7 @@ pg_tracing_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jst * With queries using simple protocol, post parse is a good place to start * tx_block span */ - start_tx_block_span(query->utilityStmt, &span_context); + start_tx_block_span(query->utilityStmt, &span_context, HOOK_PARSE); push_active_span(pg_tracing_mem_ctx, &span_context, command_type_to_span_type(query->commandType), HOOK_PARSE); @@ -1681,7 +1692,7 @@ pg_tracing_planner_hook(Query *query, const char *query_string, int cursorOption nested_level--; /* End planner span */ - pop_active_span(&span_end_time); + pop_and_store_active_span(span_end_time); /* If we have a prepared statement, add bound parameters to the top span */ if (pg_tracing_max_parameter_size && @@ -1715,6 +1726,13 @@ pg_tracing_ExecutorStart(QueryDesc *queryDesc, int eflags) /* We're at the root level, copy parse traceparent */ *traceparent = parse_traceparent; reset_traceparent(&parse_traceparent); + + /* + * Reset top span end candidate, either a new parse will set a new + * value if extended protocol is used or we will just use the current + * ts in ExecutorEnd + */ + top_span_end_candidate = 0; } /* @@ -1894,7 +1912,7 @@ pg_tracing_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 cou end_nested_level(&span_end_time); nested_level--; /* End ExecutorRun span and store it */ - pop_active_span(&span_end_time); + pop_and_store_active_span(span_end_time); per_level_infos[nested_level].executor_end = span_end_time; } @@ -1993,10 +2011,10 @@ pg_tracing_ExecutorFinish(QueryDesc *queryDesc) { span_end_time = GetCurrentTimestamp(); end_nested_level(&span_end_time); - pop_active_span(&span_end_time); + pop_and_store_active_span(span_end_time); } else - pop_active_span(NULL); + pop_active_span(); nested_level--; } @@ -2035,8 +2053,11 @@ pg_tracing_ExecutorEnd(QueryDesc *queryDesc) else standard_ExecutorEnd(queryDesc); - span_end_time = GetCurrentTimestamp(); - pop_active_span(&span_end_time); + if (nested_level == 0 && top_span_end_candidate > 0) + span_end_time = top_span_end_candidate; + else + span_end_time = GetCurrentTimestamp(); + pop_and_store_active_span(span_end_time); } /* @@ -2066,10 +2087,14 @@ pg_tracing_ProcessUtility(PlannedStmt *pstmt, const char *queryString, /* We're at root level, copy the parse traceparent */ *traceparent = parse_traceparent; reset_traceparent(&parse_traceparent); - } - if (nested_level == 0) + /* Update tracked query_id */ current_query_id = pstmt->queryId; + } + + /* Use tx_traceparent if the whole tx is traced */ + if (!traceparent->sampled && tx_traceparent.sampled) + *traceparent = tx_traceparent; if (!pg_tracing_enabled(traceparent, nested_level)) { @@ -2117,7 +2142,7 @@ pg_tracing_ProcessUtility(PlannedStmt *pstmt, const char *queryString, initialize_trace_level(); initialise_span_context(&span_context, traceparent, pstmt, NULL, NULL, queryString); - start_tx_block_span(pstmt->utilityStmt, &span_context); + start_tx_block_span(pstmt->utilityStmt, &span_context, HOOK_UTILITY); if (track_utility) { @@ -2174,9 +2199,22 @@ pg_tracing_ProcessUtility(PlannedStmt *pstmt, const char *queryString, } /* End ProcessUtility span and store it */ - pop_active_span(&span_end_time); + pop_and_store_active_span(span_end_time); + + if (nested_level == 0 && tx_block_span.span_id > 0) + { + /* + * Parent span may have been created during post parse and tx block + * started during process utility. Link the parent span to the tx + * block here + */ + Span *parent_span = peek_active_span(); + + Assert(parent_span->nested_level == 0); + parent_span->parent_id = tx_block_span.span_id; + } /* Also end and store parent active span */ - pop_active_span(&span_end_time); + pop_and_store_active_span(span_end_time); /* * If we're in an aborted transaction, xact callback won't be called so we diff --git a/src/pg_tracing.h b/src/pg_tracing.h index 5b3cb76..b7db4d6 100644 --- a/src/pg_tracing.h +++ b/src/pg_tracing.h @@ -47,12 +47,13 @@ typedef struct PlanCounters /* * Hook currently called */ -typedef enum HookPhase +typedef enum HookType { HOOK_PARSE, + HOOK_UTILITY, HOOK_PLANNER, HOOK_EXECUTOR, -} HookPhase; +} HookType; /* * Error code when parsing traceparent field @@ -410,14 +411,16 @@ extern const char *span_type_to_str(SpanType span_type); extern pgTracingStats get_empty_pg_tracing_stats(void); /* pg_tracing_active_spans.c */ -extern Span * pop_active_span(const TimestampTz *end_time); +extern Span * pop_and_store_active_span(const TimestampTz end_time); +extern Span * pop_active_span(void); extern Span * peek_active_span(void); extern Span * push_active_span(MemoryContext context, const SpanContext * span_context, - SpanType span_type, HookPhase hook_phase); + SpanType span_type, HookType hook_type); extern Span * push_child_active_span(MemoryContext context, const SpanContext * span_context, SpanType span_type); extern void cleanup_active_spans(void); +extern int num_active_spans(void); /* pg_tracing.c */ extern pgTracingSharedState * pg_tracing_shared_state; diff --git a/src/pg_tracing_active_spans.c b/src/pg_tracing_active_spans.c index 54fe79b..9675e67 100644 --- a/src/pg_tracing_active_spans.c +++ b/src/pg_tracing_active_spans.c @@ -30,6 +30,15 @@ cleanup_active_spans(void) active_spans = NULL; } +/* + * Return the number of active spans currently in the stack + */ +int +num_active_spans(void) +{ + return active_spans->end; +} + /* * Push a new active span to the active_spans stack */ @@ -88,13 +97,27 @@ static Span * peek_active_span_for_current_level(void) return span; } +/* + * Store and return the latest active span + */ +Span * +pop_and_store_active_span(const TimestampTz end_time) +{ + Span *span = pop_active_span(); + + if (span == NULL) + return NULL; + + end_span(span, &end_time); + store_span(span); + return span; +} + /* * Pop the latest active span - * - * If end_time is provided, the span is ended and stored. */ Span * -pop_active_span(const TimestampTz *end_time) +pop_active_span(void) { Span *span; @@ -102,11 +125,6 @@ pop_active_span(const TimestampTz *end_time) return NULL; span = &active_spans->spans[--active_spans->end]; - if (end_time != NULL) - { - end_span(span, end_time); - store_span(span); - } return span; } @@ -231,7 +249,7 @@ push_child_active_span(MemoryContext context, const SpanContext * span_context, */ Span * push_active_span(MemoryContext context, const SpanContext * span_context, SpanType span_type, - HookPhase hook_phase) + HookType hook_type) { Span *span = peek_active_span_for_current_level(); Span *parent_span = peek_active_span(); @@ -253,7 +271,7 @@ push_active_span(MemoryContext context, const SpanContext * span_context, SpanTy } else { - if (hook_phase != HOOK_PARSE || nested_level > 0) + if (hook_type != HOOK_PARSE || nested_level > 0) return span; /* diff --git a/src/pg_tracing_span.c b/src/pg_tracing_span.c index d5c0020..74afbcc 100644 --- a/src/pg_tracing_span.c +++ b/src/pg_tracing_span.c @@ -88,6 +88,7 @@ end_span(Span * span, const TimestampTz *end_time_input) else span->end = *end_time_input; + Assert(span->end >= span->start); if (span->type == SPAN_PLANNER || span->type == SPAN_PROCESS_UTILITY) {