Add trace context propagation through GUC

DataDog · Jul 1, 2024 · f668150 · f668150
1 parent 64d5898
commit f668150
Show file tree

Hide file tree

Showing 4 changed files with 239 additions and 4 deletions.
diff --git a/Makefile b/Makefile
@@ -34,7 +34,7 @@ endif
 REGRESSCHECKS += sample planstate planstate_bitmap planstate_hash \
 				 planstate_projectset planstate_subplans planstate_union \
 				 parallel subxact full_buffer \
-				 nested wal cleanup
+				 guc nested wal cleanup
 
 REGRESSCHECKS_OPTS = --no-locale --encoding=UTF8 --temp-config pg_tracing.conf
 

diff --git a/expected/guc.out b/expected/guc.out
@@ -0,0 +1,117 @@
+-- Test trace context propagation through GUCs
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00-00000000000000000000000000000004-0000000000000004-01';
+SELECT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Test multiple statements
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00-fffffffffffffffffffffffffffffff5-0000000000000005-01';
+SELECT 2;
+ ?column? 
+----------
+        2
+(1 row)
+
+SELECT 3;
+ ?column? 
+----------
+        3
+(1 row)
+
+-- Check results for GUC propagation with simple and multiple statements
+select trace_id, span_operation, parameters, lvl from peek_ordered_spans;
+             trace_id             | span_operation | parameters | lvl 
+----------------------------------+----------------+------------+-----
+ 00000000000000000000000000000004 | SELECT $1;     | $1 = 1     |   1
+ 00000000000000000000000000000004 | Planner        |            |   2
+ 00000000000000000000000000000004 | ExecutorRun    |            |   2
+ 00000000000000000000000000000004 | Result         |            |   3
+ fffffffffffffffffffffffffffffff5 | SELECT $1;     | $1 = 2     |   1
+ fffffffffffffffffffffffffffffff5 | Planner        |            |   2
+ fffffffffffffffffffffffffffffff5 | ExecutorRun    |            |   2
+ fffffffffffffffffffffffffffffff5 | Result         |            |   3
+ fffffffffffffffffffffffffffffff5 | SELECT $1;     | $1 = 3     |   1
+ fffffffffffffffffffffffffffffff5 | Planner        |            |   2
+ fffffffffffffffffffffffffffffff5 | ExecutorRun    |            |   2
+ fffffffffffffffffffffffffffffff5 | Result         |            |   3
+(12 rows)
+
+CALL clean_spans();
+-- Mix SQLCommenter and GUC propagation
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00-fffffffffffffffffffffffffffffff6-0000000000000006-01';
+SELECT 2;
+ ?column? 
+----------
+        2
+(1 row)
+
+/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ SELECT 1, 1;
+ ?column? | ?column? 
+----------+----------
+        1 |        1
+(1 row)
+
+/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000009-0000000000000009-00'*/ SELECT 1, 2, 3;
+ ?column? | ?column? | ?column? 
+----------+----------+----------
+        1 |        2 |        3
+(1 row)
+
+SELECT 3;
+ ?column? 
+----------
+        3
+(1 row)
+
+-- Check mix SQLCommenter and GUC propagation
+select trace_id, span_operation, parameters, lvl from peek_ordered_spans;
+             trace_id             |   span_operation   |       parameters       | lvl 
+----------------------------------+--------------------+------------------------+-----
+ fffffffffffffffffffffffffffffff6 | SELECT $1;         | $1 = 2                 |   1
+ fffffffffffffffffffffffffffffff6 | Planner            |                        |   2
+ fffffffffffffffffffffffffffffff6 | ExecutorRun        |                        |   2
+ fffffffffffffffffffffffffffffff6 | Result             |                        |   3
+ fffffffffffffffffffffffffffffff6 | SELECT $1, $2;     | $1 = 1, $2 = 1         |   1
+ fffffffffffffffffffffffffffffff6 | Planner            |                        |   2
+ fffffffffffffffffffffffffffffff6 | ExecutorRun        |                        |   2
+ fffffffffffffffffffffffffffffff6 | Result             |                        |   3
+ fffffffffffffffffffffffffffffff6 | SELECT $1, $2, $3; | $1 = 1, $2 = 2, $3 = 3 |   1
+ fffffffffffffffffffffffffffffff6 | Planner            |                        |   2
+ fffffffffffffffffffffffffffffff6 | ExecutorRun        |                        |   2
+ fffffffffffffffffffffffffffffff6 | Result             |                        |   3
+ fffffffffffffffffffffffffffffff6 | SELECT $1;         | $1 = 3                 |   1
+ fffffffffffffffffffffffffffffff6 | Planner            |                        |   2
+ fffffffffffffffffffffffffffffff6 | ExecutorRun        |                        |   2
+ fffffffffffffffffffffffffffffff6 | Result             |                        |   3
+(16 rows)
+
+CALL clean_spans();
+-- Test statement after reset
+SET pg_tracing.trace_context TO default;
+SELECT 4;
+ ?column? 
+----------
+        4
+(1 row)
+
+-- Test no traceparent field
+SET pg_tracing.trace_context='dddbs=''postgres.db'',taceparent=''00-fffffffffffffffffffffffffffffff5-0000000000000005-01';
+ERROR:  invalid value for parameter "pg_tracing.trace_context": "dddbs='postgres.db',taceparent='00-fffffffffffffffffffffffffffffff5-0000000000000005-01"
+DETAIL:  Error parsing tracecontext: No traceparent field found
+-- Test incorrect trace id
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00-ffffffffffffffffffffffffffffff5-0000000000000005-01';
+ERROR:  invalid value for parameter "pg_tracing.trace_context": "dddbs='postgres.db',traceparent='00-ffffffffffffffffffffffffffffff5-0000000000000005-01"
+DETAIL:  Error parsing tracecontext: Traceparent field doesn't have the correct size
+-- Test wrong format
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00f-ffffffffffffffffffffffffffffff5-0000000000000005-01';
+ERROR:  invalid value for parameter "pg_tracing.trace_context": "dddbs='postgres.db',traceparent='00f-ffffffffffffffffffffffffffffff5-0000000000000005-01"
+DETAIL:  Error parsing tracecontext: Incorrect traceparent format
+-- GUC errors and no GUC tracecontext should not generate spans
+select count(*) = 0 from peek_ordered_spans;
+ ?column? 
+----------
+ t
+(1 row)
+
diff --git a/sql/guc.sql b/sql/guc.sql
@@ -0,0 +1,37 @@
+-- Test trace context propagation through GUCs
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00-00000000000000000000000000000004-0000000000000004-01';
+SELECT 1;
+
+-- Test multiple statements
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00-fffffffffffffffffffffffffffffff5-0000000000000005-01';
+SELECT 2;
+SELECT 3;
+
+-- Check results for GUC propagation with simple and multiple statements
+select trace_id, span_operation, parameters, lvl from peek_ordered_spans;
+CALL clean_spans();
+
+-- Mix SQLCommenter and GUC propagation
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00-fffffffffffffffffffffffffffffff6-0000000000000006-01';
+SELECT 2;
+/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000001-0000000000000001-01'*/ SELECT 1, 1;
+/*dddbs='postgres.db',traceparent='00-00000000000000000000000000000009-0000000000000009-00'*/ SELECT 1, 2, 3;
+SELECT 3;
+
+-- Check mix SQLCommenter and GUC propagation
+select trace_id, span_operation, parameters, lvl from peek_ordered_spans;
+CALL clean_spans();
+
+-- Test statement after reset
+SET pg_tracing.trace_context TO default;
+SELECT 4;
+
+-- Test no traceparent field
+SET pg_tracing.trace_context='dddbs=''postgres.db'',taceparent=''00-fffffffffffffffffffffffffffffff5-0000000000000005-01';
+-- Test incorrect trace id
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00-ffffffffffffffffffffffffffffff5-0000000000000005-01';
+-- Test wrong format
+SET pg_tracing.trace_context='dddbs=''postgres.db'',traceparent=''00f-ffffffffffffffffffffffffffffff5-0000000000000005-01';
+
+-- GUC errors and no GUC tracecontext should not generate spans
+select count(*) = 0 from peek_ordered_spans;
diff --git a/src/pg_tracing.c b/src/pg_tracing.c
@@ -130,6 +130,8 @@ static int	pg_tracing_otel_naptime;	/* Delay between upload of spans to
 										 * otel collector */
 static int	pg_tracing_otel_connect_timeout_ms; /* Connect timeout to the otel
 												 * collector */
+static char *guc_tracecontext_str = NULL;	/* Trace context string propagated
+											 * through GUC variable */
 
 static const struct config_enum_entry track_options[] =
 {
@@ -169,6 +171,9 @@ static MemoryContext pg_tracing_mem_ctx;
 /* trace context at the root level of parse/planning hook */
 static Traceparent parse_traceparent;
 
+/* trace context set through GUC */
+static Traceparent * guc_tracecontext;
+
 /* trace context used in nested levels or within executor hooks */
 static Traceparent executor_traceparent;
 
@@ -250,14 +255,17 @@ static void pg_tracing_ProcessUtility(PlannedStmt *pstmt, const char *queryStrin
 									  ParamListInfo params,
 									  QueryEnvironment *queryEnv,
 									  DestReceiver *dest, QueryCompletion *qc);
+static void pg_tracing_xact_callback(XactEvent event, void *arg);
 
 static void pg_tracing_shmem_request(void);
 static void pg_tracing_shmem_startup_hook(void);
-static void reset_traceparent(Traceparent * traceparent);
+
+static void initialize_trace_level(void);
+
 static bool check_filter_query_ids(char **newval, void **extra, GucSource source);
 static void assign_filter_query_ids(const char *newval, void *extra);
-static void initialize_trace_level(void);
-static void pg_tracing_xact_callback(XactEvent event, void *arg);
+static bool check_guc_tracecontext(char **newval, void **extra, GucSource source);
+static void assign_guc_tracecontext_hook(const char *newval, void *extra);
 
 /*
  * Module load callback
@@ -447,6 +455,18 @@ _PG_init(void)
 							   NULL,
 							   NULL);
 
+	DefineCustomStringVariable("pg_tracing.trace_context",
+							   "Trace context propagated through GUC variable.",
+							   NULL,
+							   &guc_tracecontext_str,
+							   NULL,
+							   PGC_USERSET,
+							   0,
+							   check_guc_tracecontext,
+							   assign_guc_tracecontext_hook,
+							   NULL);
+
+
 	MarkGUCPrefixReserved("pg_tracing");
 
 	/* For jumble state */
@@ -651,6 +671,60 @@ assign_filter_query_ids(const char *newval, void *extra)
 	query_id_filter = (QueryIdFilter *) extra;
 }
 
+/*
+ * Check hook for trace context guc parameter
+ */
+static bool
+check_guc_tracecontext(char **newval, void **extra, GucSource source)
+{
+	char	   *input_trace_context = *newval;
+	ParseTraceparentErr err;
+	Traceparent *result;
+	Traceparent parsed_traceparent;
+
+	if (!input_trace_context || strcmp(input_trace_context, "") == 0)
+	{
+		/*
+		 * Empty tracecontext, extra should already have been set to NULL so
+		 * nothing to do
+		 */
+		return true;
+	}
+
+	/* Parse the parameter string */
+	err = parse_trace_context(&parsed_traceparent, input_trace_context, strlen(input_trace_context));
+	if (err != PARSE_OK)
+	{
+		GUC_check_errdetail("Error parsing tracecontext: %s", parse_code_to_err(err));
+		return false;
+	}
+
+	/* We have a valid traceaprent, store it in a guc_malloced variable */
+	result = (Traceparent *) guc_malloc(LOG, sizeof(Traceparent));
+	if (result == NULL)
+		return false;
+
+	*result = parsed_traceparent;
+	*extra = result;
+	return true;
+}
+
+/*
+ * Assign hook for trace context guc
+ */
+static void
+assign_guc_tracecontext_hook(const char *newval, void *extra)
+{
+	/*
+	 * We don't want to trace the SET query modifying the trace context so
+	 * cancel ongoing trace.
+	 */
+	cleanup_tracing();
+
+	/* Copy the parsed traceparent to our global guc_tracecontext */
+	guc_tracecontext = (Traceparent *) extra;
+}
+
 /*
  * Check whether trace context is filtered based on query id
  */
@@ -969,6 +1043,13 @@ extract_trace_context(Traceparent * traceparent, ParseState *pstate,
 	if (nested_level > 0)
 		goto cleanup;
 
+	/* We have a traceparent propagated through GUC, copy it */
+	if (guc_tracecontext && guc_tracecontext->sampled)
+	{
+		*traceparent = *guc_tracecontext;
+		goto cleanup;
+	}
+
 	/* Both sampling rate are set to 0, no tracing will happen */
 	if (pg_tracing_sample_rate == 0 && pg_tracing_caller_sample_rate == 0)
 		goto cleanup;