diff --git a/Makefile b/Makefile index ba8fa2c..6bf24f9 100644 --- a/Makefile +++ b/Makefile @@ -19,7 +19,7 @@ OBJS = \ src/pg_tracing_span.o \ src/version_compat.o -REGRESSCHECKS = setup utility select insert trigger cursor +REGRESSCHECKS = setup utility select insert trigger cursor transaction ifeq ($(PG_VERSION),15) REGRESSCHECKS += trigger_15 else diff --git a/expected/cursor.out b/expected/cursor.out index 516f481..16e8df9 100644 --- a/expected/cursor.out +++ b/expected/cursor.out @@ -1,10 +1,10 @@ -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-0000000000000001-01'*/ BEGIN; +DECLARE c CURSOR FOR SELECT * from pg_tracing_test; +FETCH FORWARD 20 from c \gset more than one row returned for \gset -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000003-01'*/ FETCH BACKWARD 10 from c \gset +FETCH BACKWARD 10 from c \gset more than one row returned for \gset -/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000004-01'*/ CLOSE c; +CLOSE c; COMMIT; -- First declare -- +----------------------------------------+ @@ -20,7 +20,8 @@ SELECT span_id AS span_a_id, get_epoch(span_start) as span_a_start, get_epoch(span_end) as span_a_end from pg_tracing_peek_spans - where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000001' \gset + where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000001' + AND span_operation='DECLARE c CURSOR FOR SELECT * from pg_tracing_test;' \gset SELECT span_id AS span_b_id, get_epoch(span_start) as span_b_start, get_epoch(span_end) as span_b_end @@ -63,7 +64,8 @@ SELECT span_id AS span_a_id, get_epoch(span_start) as span_a_start, get_epoch(span_end) as span_a_end from pg_tracing_peek_spans - where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000002' \gset + where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000001' + AND span_operation='FETCH FORWARD 20 from c' \gset SELECT span_id AS span_b_id, get_epoch(span_start) as span_b_start, get_epoch(span_end) as span_b_end @@ -98,7 +100,8 @@ SELECT span_id AS span_a_id, get_epoch(span_start) as span_a_start, get_epoch(span_end) as span_a_end from pg_tracing_peek_spans - where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000003' \gset + where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000001' + AND span_operation='FETCH BACKWARD 10 from c' \gset SELECT span_id AS span_b_id, get_epoch(span_start) as span_b_start, get_epoch(span_end) as span_b_end @@ -131,6 +134,8 @@ SELECT :span_c_start >= :span_b_start as nested_declare_starts_after_parent, SELECT span_type, span_operation, lvl from peek_ordered_spans where trace_id='00000000000000000000000000000001'; span_type | span_operation | lvl ----------------+-----------------------------------------------------+----- + Utility query | BEGIN; | 1 + ProcessUtility | ProcessUtility | 2 Utility query | DECLARE c CURSOR FOR SELECT * from pg_tracing_test; | 1 ProcessUtility | ProcessUtility | 2 Select query | DECLARE c CURSOR FOR SELECT * from pg_tracing_test; | 3 @@ -146,7 +151,9 @@ SELECT span_type, span_operation, lvl from peek_ordered_spans where trace_id='00 Utility query | CLOSE c; | 1 ProcessUtility | ProcessUtility | 2 Select query | DECLARE c CURSOR FOR SELECT * from pg_tracing_test; | 3 -(15 rows) + Utility query | COMMIT; | 1 + ProcessUtility | ProcessUtility | 2 +(19 rows) -- Clean created spans CALL clean_spans(); diff --git a/expected/select.out b/expected/select.out index a257780..86ff5ac 100644 --- a/expected/select.out +++ b/expected/select.out @@ -211,13 +211,17 @@ SELECT 1\; SELECT 1, 2; (1 row) SELECT span_type, span_operation, parameters, lvl from peek_ordered_spans; - span_type | span_operation | parameters | lvl ---------------+----------------+------------+----- - Select query | SELECT $1; | $1 = 1 | 1 - Planner | Planner | | 2 - Executor | ExecutorRun | | 2 - Result | Result | | 3 -(4 rows) + span_type | span_operation | parameters | lvl +--------------+----------------+----------------+----- + Select query | SELECT $1; | $1 = 1 | 1 + Planner | Planner | | 2 + Executor | ExecutorRun | | 2 + Result | Result | | 3 + Select query | SELECT $1, $2; | $1 = 1, $2 = 2 | 1 + Planner | Planner | | 2 + Executor | ExecutorRun | | 2 + Result | Result | | 3 +(8 rows) CALL clean_spans(); -- Check standalone trace diff --git a/expected/transaction.out b/expected/transaction.out new file mode 100644 index 0000000..62b721a --- /dev/null +++ b/expected/transaction.out @@ -0,0 +1,51 @@ +-- Only trace queries with sample flag +SET pg_tracing.sample_rate = 0.0; +SET pg_tracing.caller_sample_rate = 1.0; +-- Set tracecontext at the start of the transaction +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ BEGIN; SELECT 1; COMMIT; + ?column? +---------- + 1 +(1 row) + +SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000001'; + span_type | span_operation | lvl +----------------+----------------+----- + Utility query | BEGIN; | 1 + ProcessUtility | ProcessUtility | 2 + Select query | SELECT $1; | 1 + Planner | Planner | 2 + Executor | ExecutorRun | 2 + Result | Result | 3 + Utility query | COMMIT; | 1 + ProcessUtility | ProcessUtility | 2 +(8 rows) + +-- Test with override of the trace context in the middle of a transaction +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000002-0000000000000002-01'*/ BEGIN; +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000003-0000000000000003-01'*/ SELECT 1; COMMIT; + ?column? +---------- + 1 +(1 row) + +SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000002'; + span_type | span_operation | lvl +----------------+----------------+----- + Utility query | BEGIN; | 1 + ProcessUtility | ProcessUtility | 2 + Utility query | COMMIT; | 1 + ProcessUtility | ProcessUtility | 2 +(4 rows) + +SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000003'; + span_type | span_operation | lvl +--------------+----------------+----- + Select query | SELECT $1; | 1 + Planner | Planner | 2 + Executor | ExecutorRun | 2 + Result | Result | 3 +(4 rows) + +CALL clean_spans(); +CALL reset_settings(); diff --git a/sql/cursor.sql b/sql/cursor.sql index ad19880..4211ed4 100644 --- a/sql/cursor.sql +++ b/sql/cursor.sql @@ -1,8 +1,8 @@ -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; +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ BEGIN; +DECLARE c CURSOR FOR SELECT * from pg_tracing_test; +FETCH FORWARD 20 from c \gset +FETCH BACKWARD 10 from c \gset +CLOSE c; COMMIT; -- First declare @@ -20,7 +20,8 @@ SELECT span_id AS span_a_id, get_epoch(span_start) as span_a_start, get_epoch(span_end) as span_a_end from pg_tracing_peek_spans - where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000001' \gset + where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000001' + AND span_operation='DECLARE c CURSOR FOR SELECT * from pg_tracing_test;' \gset SELECT span_id AS span_b_id, get_epoch(span_start) as span_b_start, get_epoch(span_end) as span_b_end @@ -57,7 +58,8 @@ SELECT span_id AS span_a_id, get_epoch(span_start) as span_a_start, get_epoch(span_end) as span_a_end from pg_tracing_peek_spans - where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000002' \gset + where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000001' + AND span_operation='FETCH FORWARD 20 from c' \gset SELECT span_id AS span_b_id, get_epoch(span_start) as span_b_start, get_epoch(span_end) as span_b_end @@ -85,7 +87,8 @@ SELECT span_id AS span_a_id, get_epoch(span_start) as span_a_start, get_epoch(span_end) as span_a_end from pg_tracing_peek_spans - where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000003' \gset + where trace_id='00000000000000000000000000000001' AND parent_id='0000000000000001' + AND span_operation='FETCH BACKWARD 10 from c' \gset SELECT span_id AS span_b_id, get_epoch(span_start) as span_b_start, get_epoch(span_end) as span_b_end diff --git a/sql/transaction.sql b/sql/transaction.sql new file mode 100644 index 0000000..6f6dc42 --- /dev/null +++ b/sql/transaction.sql @@ -0,0 +1,18 @@ +-- Only trace queries with sample flag +SET pg_tracing.sample_rate = 0.0; +SET pg_tracing.caller_sample_rate = 1.0; + +-- Set tracecontext at the start of the transaction +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ BEGIN; SELECT 1; COMMIT; + +SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000001'; + +-- Test with override of the trace context in the middle of a transaction +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000002-0000000000000002-01'*/ BEGIN; +/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000003-0000000000000003-01'*/ SELECT 1; COMMIT; + +SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000002'; +SELECT span_type, span_operation, lvl FROM peek_ordered_spans where trace_id='00000000000000000000000000000003'; + +CALL clean_spans(); +CALL reset_settings(); diff --git a/src/pg_tracing.c b/src/pg_tracing.c index 40d60d0..e9690ff 100644 --- a/src/pg_tracing.c +++ b/src/pg_tracing.c @@ -180,8 +180,8 @@ static struct pgTracingTraceContext parsed_trace_context; /* trace context used in nested levels or within executor hooks */ static struct pgTracingTraceContext executor_trace_context; -/* Latest trace id observed */ -static TraceId latest_trace_id; +/* traceparent extracted at the start of a transaction */ +static pgTracingTraceparent tx_start_traceparent; /* Latest local transaction id traced */ static LocalTransactionId latest_lxid = InvalidLocalTransactionId; @@ -955,8 +955,15 @@ set_trace_id(struct pgTracingTraceContext *trace_context) { bool new_lxid = is_new_lxid(); + /* trace id was already set, use it */ if (!traceid_zero(trace_context->traceparent.trace_id)) { + /* + * Keep track of the trace context in case there are multiple + * statement in the transaction + */ + if (new_lxid) + tx_start_traceparent = trace_context->traceparent; return; } @@ -966,7 +973,9 @@ set_trace_id(struct pgTracingTraceContext *trace_context) */ if (!new_lxid) { - trace_context->traceparent.trace_id = latest_trace_id; + /* We're in the same transaction, use the transaction trace context */ + Assert(!traceid_zero(tx_start_traceparent.trace_id)); + trace_context->traceparent = tx_start_traceparent; return; } @@ -978,7 +987,8 @@ set_trace_id(struct pgTracingTraceContext *trace_context) trace_context->traceparent.trace_id.traceid_left = pg_prng_int64(&pg_global_prng_state); trace_context->traceparent.trace_id.traceid_right = pg_prng_int64(&pg_global_prng_state); - latest_trace_id = trace_context->traceparent.trace_id; + /* We're at the begining of a new local transaction, save trace context */ + tx_start_traceparent = trace_context->traceparent; } /* @@ -1520,6 +1530,9 @@ pg_tracing_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jst bool new_lxid = is_new_lxid(); bool is_root_level = exec_nested_level + plan_nested_level == 0; + if (new_lxid) + /* We have a new local transaction, reset the begin tx traceparent */ + reset_traceparent(&tx_start_traceparent); if (!pg_tracing_mem_ctx->isReset && new_lxid && is_root_level) @@ -1553,8 +1566,22 @@ pg_tracing_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jst extract_trace_context(trace_context, pstate, query->queryId); if (!trace_context->traceparent.sampled) - /* Query is not sampled, nothing to do */ - return; + { + if (is_root_level && !new_lxid && tx_start_traceparent.sampled) + { + /* + * We're in the same local transaction, use the trace context + * extracted at the begining of the transaction + */ + Assert(!traceid_zero(tx_start_traceparent.trace_id)); + trace_context->traceparent = tx_start_traceparent; + } + else + { + /* Query is not sampled, nothing to do. */ + return; + } + } /* * We want to avoid calling GetCurrentTimestamp at the start of post parse