While handling error, stop the node instrumentation if it was left ru…

…nning Calling InstrEndLoop will throw an error if the node is still running. Explicitely stop it to avoid this.
DataDog · Sep 10, 2024 · cd2e01a · cd2e01a
1 parent 3d6992a
commit cd2e01a
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 0 deletions.
diff --git a/expected/trigger.out b/expected/trigger.out
@@ -246,3 +246,30 @@ SELECT count(distinct query_id) from pg_tracing_consume_spans;
      4
 (1 row)
 
+-- Test foreign key failure
+CREATE TYPE rainbow AS ENUM ('red', 'orange', 'yellow', 'green', 'blue', 'purple');
+CREATE TABLE enumtest_parent (id rainbow PRIMARY KEY);
+CREATE TABLE enumtest_child (parent rainbow REFERENCES enumtest_parent);
+INSERT INTO enumtest_parent VALUES ('red');
+/*traceparent='00-00000000000000000000000000000004-0000000000000004-01'*/ INSERT INTO enumtest_child VALUES ('blue');
+ERROR:  insert or update on table "enumtest_child" violates foreign key constraint "enumtest_child_parent_fkey"
+DETAIL:  Key (parent)=(blue) is not present in table "enumtest_parent".
+SELECT trace_id, span_type, span_operation, lvl from peek_ordered_spans where trace_id='00000000000000000000000000000004';
+             trace_id             |   span_type    |                                                                  span_operation                                                                  | lvl 
+----------------------------------+----------------+--------------------------------------------------------------------------------------------------------------------------------------------------+-----
+ 00000000000000000000000000000004 | Insert query   | INSERT INTO enumtest_child VALUES ($1);                                                                                                          |   1
+ 00000000000000000000000000000004 | Planner        | Planner                                                                                                                                          |   2
+ 00000000000000000000000000000004 | ExecutorRun    | ExecutorRun                                                                                                                                      |   2
+ 00000000000000000000000000000004 | Insert         | Insert on enumtest_child                                                                                                                         |   3
+ 00000000000000000000000000000004 | Result         | Result                                                                                                                                           |   4
+ 00000000000000000000000000000004 | ExecutorFinish | ExecutorFinish                                                                                                                                   |   2
+ 00000000000000000000000000000004 | Select query   | SELECT $2 FROM ONLY "public"."enumtest_parent" x WHERE "id"::pg_catalog.anyenum OPERATOR(pg_catalog.=) $1::pg_catalog.anyenum FOR KEY SHARE OF x |   3
+ 00000000000000000000000000000004 | Planner        | Planner                                                                                                                                          |   4
+ 00000000000000000000000000000004 | ExecutorRun    | ExecutorRun                                                                                                                                      |   4
+ 00000000000000000000000000000004 | LockRows       | LockRows                                                                                                                                         |   5
+ 00000000000000000000000000000004 | IndexScan      | IndexScan using enumtest_parent_pkey on enumtest_parent x                                                                                        |   6
+(11 rows)
+
+-- Cleanup
+CALL reset_settings();
+CALL clean_spans();
diff --git a/sql/trigger.sql b/sql/trigger.sql
@@ -160,3 +160,15 @@ SELECT trace_id, span_type, span_operation, lvl from peek_ordered_spans where tr
 
 -- Check count of query_id
 SELECT count(distinct query_id) from pg_tracing_consume_spans;
+
+-- Test foreign key failure
+CREATE TYPE rainbow AS ENUM ('red', 'orange', 'yellow', 'green', 'blue', 'purple');
+CREATE TABLE enumtest_parent (id rainbow PRIMARY KEY);
+CREATE TABLE enumtest_child (parent rainbow REFERENCES enumtest_parent);
+INSERT INTO enumtest_parent VALUES ('red');
+/*traceparent='00-00000000000000000000000000000004-0000000000000004-01'*/ INSERT INTO enumtest_child VALUES ('blue');
+SELECT trace_id, span_type, span_operation, lvl from peek_ordered_spans where trace_id='00000000000000000000000000000004';
+
+-- Cleanup
+CALL reset_settings();
+CALL clean_spans();
diff --git a/src/pg_tracing.c b/src/pg_tracing.c
@@ -1045,6 +1045,9 @@ process_query_desc(const Traceparent * traceparent, const QueryDesc *queryDesc,
 {
 	NodeCounters *node_counters = &peek_active_span()->node_counters;
 
+    if (!queryDesc->totaltime->running)
+        return;
+
 	/* Process total counters */
 	if (queryDesc->totaltime)
 	{
@@ -1363,13 +1366,21 @@ handle_pg_error(const Traceparent * traceparent,
 	sql_error_code = geterrcode();
 
 	if (queryDesc != NULL)
+	{
+		/*
+		 * On error, instrumentation may have been left running, stop it to
+		 * avoid error thrown by InstrEndLoop
+		 */
+		if (!INSTR_TIME_IS_ZERO(queryDesc->totaltime->starttime))
+			InstrStopNode(queryDesc->totaltime, 0);
 
 		/*
 		 * 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)
 	{