Store tx block when tx is aborted

DataDog · Jul 29, 2024 · e9cc2d0 · e9cc2d0
1 parent 0d1fc8f
commit e9cc2d0
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 28 deletions.
diff --git a/expected/extended.out b/expected/extended.out
@@ -22,14 +22,15 @@ BEGIN; select $1 \bind \g
 ERROR:  bind message supplies 0 parameters, but prepared statement "" requires 1
 ROLLBACK;
 SELECT span_type, span_operation, parameters, lvl FROM peek_ordered_spans;
-   span_type    | span_operation | parameters | lvl 
-----------------+----------------+------------+-----
- Utility query  | BEGIN;         |            |   1
- ProcessUtility | ProcessUtility |            |   2
- Select query   | select $1      |            |   1
- Utility query  | ROLLBACK;      |            |   1
- ProcessUtility | ProcessUtility |            |   2
-(5 rows)
+    span_type     |  span_operation  | parameters | lvl 
+------------------+------------------+------------+-----
+ TransactionBlock | TransactionBlock |            |   1
+ Utility query    | BEGIN;           |            |   2
+ ProcessUtility   | ProcessUtility   |            |   3
+ Select query     | select $1        |            |   2
+ Utility query    | ROLLBACK;        |            |   1
+ ProcessUtility   | ProcessUtility   |            |   2
+(6 rows)
 
 CALL clean_spans();
 -- Execute queries with extended protocol within an explicit transaction

diff --git a/expected/json.out b/expected/json.out
@@ -74,19 +74,20 @@ SELECT trace_id, name,
     lvl FROM peek_ordered_json_spans;
              trace_id             |                               name                               | subxact_count | has_shared_hit | has_wal_records | has_wal_bytes | has_startup | lvl 
 ----------------------------------+------------------------------------------------------------------+---------------+----------------+-----------------+---------------+-------------+-----
- 00000000000000000000000000000001 | BEGIN;                                                           |               |                |                 |               |             |   1
- 00000000000000000000000000000001 | ProcessUtility                                                   |               |                |                 |               |             |   2
- 00000000000000000000000000000001 | SAVEPOINT s1;                                                    |               |                |                 |               |             |   1
- 00000000000000000000000000000001 | ProcessUtility                                                   |               |                |                 |               |             |   2
- 00000000000000000000000000000001 | INSERT INTO pg_tracing_test VALUES(generate_series($1, $2), $3); |               | t              | t               | t             |             |   1
- 00000000000000000000000000000001 | Planner                                                          |               |                |                 |               |             |   2
- 00000000000000000000000000000001 | ExecutorRun                                                      |               |                |                 |               |             |   2
- 00000000000000000000000000000001 | Insert on pg_tracing_test                                        |             1 | t              | t               | t             |             |   3
- 00000000000000000000000000000001 | ProjectSet                                                       |             1 |                |                 |               |             |   4
- 00000000000000000000000000000001 | Result                                                           |             1 |                |                 |               |             |   5
- 00000000000000000000000000000001 | ROLLBACK;                                                        |             1 |                |                 |               |             |   1
- 00000000000000000000000000000001 | ProcessUtility                                                   |             1 |                |                 |               |             |   2
-(12 rows)
+ 00000000000000000000000000000001 | TransactionBlock                                                 |               |                |                 |               |             |   1
+ 00000000000000000000000000000001 | BEGIN;                                                           |               |                |                 |               |             |   2
+ 00000000000000000000000000000001 | ProcessUtility                                                   |               |                |                 |               |             |   3
+ 00000000000000000000000000000001 | SAVEPOINT s1;                                                    |               |                |                 |               |             |   2
+ 00000000000000000000000000000001 | ProcessUtility                                                   |               |                |                 |               |             |   3
+ 00000000000000000000000000000001 | INSERT INTO pg_tracing_test VALUES(generate_series($1, $2), $3); |               | t              | t               | t             |             |   2
+ 00000000000000000000000000000001 | Planner                                                          |               |                |                 |               |             |   3
+ 00000000000000000000000000000001 | ExecutorRun                                                      |               |                |                 |               |             |   3
+ 00000000000000000000000000000001 | Insert on pg_tracing_test                                        |             1 | t              | t               | t             |             |   4
+ 00000000000000000000000000000001 | ProjectSet                                                       |             1 |                |                 |               |             |   5
+ 00000000000000000000000000000001 | Result                                                           |             1 |                |                 |               |             |   6
+ 00000000000000000000000000000001 | ROLLBACK;                                                        |             1 |                |                 |               |             |   2
+ 00000000000000000000000000000001 | ProcessUtility                                                   |             1 |                |                 |               |             |   3
+(13 rows)
 
 CALL clean_spans();
 -- Test parameters and deparse_info

diff --git a/src/pg_tracing.c b/src/pg_tracing.c
@@ -2191,6 +2191,7 @@ pg_tracing_xact_callback(XactEvent event, void *arg)
 				}
 				break;
 			}
+		case XACT_EVENT_ABORT:
 		case XACT_EVENT_COMMIT:
 			current_ts = GetCurrentTimestamp();
 			end_nested_level(&current_ts);
@@ -2218,13 +2219,6 @@ pg_tracing_xact_callback(XactEvent event, void *arg)
 			end_nested_level(NULL);
 			end_tracing();
 			break;
-		case XACT_EVENT_ABORT:
-			/* TODO: Create an abort span */
-			end_nested_level(NULL);
-			end_tracing();
-			reset_span(&commit_span);
-			reset_span(&tx_block_span);
-			break;
 		default:
 			break;
 	}