fix: fix trace context pattern, remove trace id and respect logging span id field. (#667)

liuyunnnn · web-flow · commit 0fb00a5a5485 · 2024-12-26T14:12:25.000-08:00
diff --git a/docs/generated/api.json b/docs/generated/api.json
@@ -1952,33 +1952,6 @@
                 "startIndex": 1,
                 "endIndex": 2
               }
-            },
-            {
-              "kind": "PropertySignature",
-              "canonicalReference": "@google-cloud/functions-framework!Request_2#traceId:member",
-              "docComment": "/**\n * Cloud Trace trace ID.\n */\n",
-              "excerptTokens": [
-                {
-                  "kind": "Content",
-                  "text": "traceId?: "
-                },
-                {
-                  "kind": "Content",
-                  "text": "string"
-                },
-                {
-                  "kind": "Content",
-                  "text": ";"
-                }
-              ],
-              "isReadonly": false,
-              "isOptional": true,
-              "releaseTag": "Public",
-              "name": "traceId",
-              "propertyTypeTokenRange": {
-                "startIndex": 1,
-                "endIndex": 2
-              }
             }
           ],
           "extendsTokenRanges": [
diff --git a/docs/generated/api.md.api.md b/docs/generated/api.md.api.md
@@ -116,7 +116,6 @@ interface Request_2 extends Request_3 {
     executionId?: string;
     rawBody?: Buffer;
     spanId?: string;
-    traceId?: string;
 }
 export { Request_2 as Request }
 
diff --git a/src/async_local_storage.ts b/src/async_local_storage.ts
@@ -6,7 +6,6 @@ import type {AsyncLocalStorage} from 'node:async_hooks';
 
 export interface ExecutionContext {
   executionId?: string;
-  traceId?: string;
   spanId?: string;
 }
 
@@ -32,7 +31,6 @@ export async function asyncLocalStorageMiddleware(
   asyncLocalStorage.run(
     {
       executionId: req.executionId,
-      traceId: req.traceId,
       spanId: req.spanId,
     },
     () => {
diff --git a/src/execution_context.ts b/src/execution_context.ts
@@ -5,7 +5,7 @@ const FUNCTION_EXECUTION_ID_HEADER_KEY = 'function-execution-id';
 const TRACE_CONTEXT_HEADER_KEY = 'X-Cloud-Trace-Context';
 
 const TRACE_CONTEXT_PATTERN =
-  /^(?<traceId>\w+)\/(?<spanId>\d+);o=(?<options>.+)$/;
+  /^(?<traceId>\w+)\/(?<spanId>\d+)(?:;o=(?<options>.+))?$/;
 
 function generateExecutionId() {
   const timestampPart = Date.now().toString(36).slice(-6);
@@ -28,9 +28,7 @@ export const executionContextMiddleware = (
   if (cloudTraceContext) {
     const match = cloudTraceContext.match(TRACE_CONTEXT_PATTERN);
     if (match?.groups) {
-      const {traceId, spanId} = match.groups;
-      req.traceId = traceId;
-      req.spanId = spanId;
+      req.spanId = match.groups.spanId;
     }
   }
   next();
diff --git a/src/functions.ts b/src/functions.ts
@@ -36,10 +36,6 @@ export interface Request extends ExpressRequest {
    * Request-specified execution ID.
    */
   executionId?: string;
-  /**
-   * Cloud Trace trace ID.
-   */
-  traceId?: string;
   /**
    * Cloud Trace span ID.
    */
diff --git a/src/logger.ts b/src/logger.ts
@@ -18,7 +18,6 @@ import {getCurrentContext, ExecutionContext} from './async_local_storage';
 import {Buffer} from 'buffer';
 
 export const EXECUTION_CONTEXT_LABELS_KEY = 'logging.googleapis.com/labels';
-export const EXECUTION_CONTEXT_TRACE_KEY = 'logging.googleapis.com/trace';
 export const EXECUTION_CONTEXT_SPAN_ID_KEY = 'logging.googleapis.com/spanId';
 const SEVERITY = 'severity';
 
@@ -132,7 +131,6 @@ export function getModifiedData(
   let dataWithContext: {
     message: string | Uint8Array;
     'logging.googleapis.com/labels': {execution_id: string | undefined};
-    'logging.googleapis.com/trace': string | undefined;
     'logging.googleapis.com/spanId': string | undefined;
     severity?: string | undefined;
   };
@@ -155,7 +153,6 @@ function getTextWithContext(
   return {
     message: data,
     [EXECUTION_CONTEXT_LABELS_KEY]: {execution_id: context.executionId},
-    [EXECUTION_CONTEXT_TRACE_KEY]: context.traceId,
     [EXECUTION_CONTEXT_SPAN_ID_KEY]: context.spanId,
   };
 }
@@ -167,11 +164,10 @@ function getJSONWithContext(json: any, context: ExecutionContext) {
   } else {
     json[EXECUTION_CONTEXT_LABELS_KEY] = {execution_id: context.executionId};
   }
-  return {
-    ...json,
-    [EXECUTION_CONTEXT_TRACE_KEY]: context.traceId,
-    [EXECUTION_CONTEXT_SPAN_ID_KEY]: context.spanId,
-  };
+  if (!(EXECUTION_CONTEXT_SPAN_ID_KEY in json)) {
+    json[EXECUTION_CONTEXT_SPAN_ID_KEY] = context.spanId;
+  }
+  return json;
 }
 
 function processData(data: Uint8Array | string, encoding?: BufferEncoding) {
diff --git a/test/async_local_storage.ts b/test/async_local_storage.ts
@@ -15,7 +15,6 @@ describe('asyncLocalStorageMiddleware', () => {
     const req = {
       body: 'test body',
       executionId: 'testExecutionId',
-      traceId: 'testtrace',
       spanId: 'testSpanId',
     };
     let executionContext;
@@ -25,7 +24,6 @@ describe('asyncLocalStorageMiddleware', () => {
       assert(executionContext);
       assert.strictEqual(executionContext.executionId, req.executionId);
       assert.strictEqual(executionContext.spanId, req.spanId);
-      assert.strictEqual(executionContext.traceId, req.traceId);
     };
 
     await asyncLocalStorageMiddleware(
diff --git a/test/execution_context.ts b/test/execution_context.ts
@@ -19,29 +19,32 @@ describe('executionContextMiddleware', () => {
   const testSpanId = '123';
   const testTrace = 'testtrace';
   const validExecutionId = 'xn1h9xdgv6zw';
-  const headers = {
-    'X-Cloud-Trace-Context': `${testTrace}/${testSpanId};o=1`,
-    'function-execution-id': validExecutionId,
-  };
 
   it('uses execution ID in header', () => {
-    const req = createRequest({}, headers);
+    const req = createRequest(
+      {},
+      {
+        'X-Cloud-Trace-Context': `${testTrace}/${testSpanId};o=1`,
+        'function-execution-id': validExecutionId,
+      }
+    );
 
     executionContextMiddleware(req as Request, {} as Response, next);
 
     assert.strictEqual(req.executionId, validExecutionId);
     assert.strictEqual(req.spanId, testSpanId);
-    assert.strictEqual(req.traceId, testTrace);
   });
 
   it('generates execution ID if not in header', () => {
-    const req = createRequest({}, headers);
+    const req = createRequest(
+      {},
+      {'X-Cloud-Trace-Context': `${testTrace}/${testSpanId}`}
+    );
 
     executionContextMiddleware(req as Request, {} as Response, next);
 
     assert(req.executionId);
     assert.strictEqual(req.spanId, testSpanId);
-    assert.strictEqual(req.traceId, testTrace);
   });
 
   it('req trace undefined if not in header', () => {
@@ -51,6 +54,5 @@ describe('executionContextMiddleware', () => {
 
     assert(req.executionId);
     assert.strictEqual(req.spanId, undefined);
-    assert.strictEqual(req.traceId, undefined);
   });
 });
diff --git a/test/logger.ts b/test/logger.ts
@@ -47,14 +47,12 @@ describe('getModifiedData', () => {
   const sampleUint8Arr = new Uint8Array(Buffer.from(sampleText));
   const expectedExecutionContext = {
     executionId: 'testExecutionId',
-    traceId: 'testTraceId',
     spanId: 'testSpanId',
   };
   const expectedMetadata = {
     'logging.googleapis.com/labels': {
       execution_id: 'testExecutionId',
     },
-    'logging.googleapis.com/trace': 'testTraceId',
     'logging.googleapis.com/spanId': 'testSpanId',
   };
   const expectedTextOutput =
@@ -102,13 +100,31 @@ describe('getModifiedData', () => {
           user_label_1: 'value_1',
           execution_id: 'testExecutionId',
         },
-        'logging.googleapis.com/trace': 'testTraceId',
         'logging.googleapis.com/spanId': 'testSpanId',
       }) + '\n';
     const modifiedData = getModifiedData(data);
     assert.equal(modifiedData, expectedOutput);
   });
 
+  it('json with user span id', () => {
+    const data = JSON.stringify({
+      text: 'default text.',
+      component: 'arbitrary-property',
+      'logging.googleapis.com/spanId': 'mySpanId',
+    });
+    const expectedOutput =
+      JSON.stringify({
+        text: 'default text.',
+        component: 'arbitrary-property',
+        'logging.googleapis.com/spanId': 'mySpanId',
+        'logging.googleapis.com/labels': {
+          execution_id: 'testExecutionId',
+        },
+      }) + '\n';
+    const modifiedData = getModifiedData(data);
+    assert.equal(modifiedData, expectedOutput);
+  });
+
   it('uint8array', () => {
     const modifiedData = getModifiedData(sampleUint8Arr);
     assert.equal(modifiedData, expectedTextOutput);

Original file line number	Diff line number	Diff line change
`@@ -116,7 +116,6 @@ interface Request_2 extends Request_3 {`
`116`	`116`	`executionId?: string;`
`117`	`117`	`rawBody?: Buffer;`
`118`	`118`	`spanId?: string;`
`119`		`- traceId?: string;`
`120`	`119`	`}`
`121`	`120`	`export { Request_2 as Request }`
`122`	`121`
Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,6 @@ import type {AsyncLocalStorage} from 'node:async_hooks';`
`6`	`6`
`7`	`7`	`export interface ExecutionContext {`
`8`	`8`	`executionId?: string;`
`9`		`- traceId?: string;`
`10`	`9`	`spanId?: string;`
`11`	`10`	`}`
`12`	`11`
`@@ -32,7 +31,6 @@ export async function asyncLocalStorageMiddleware(`
`32`	`31`	`asyncLocalStorage.run(`
`33`	`32`	`{`
`34`	`33`	`executionId: req.executionId,`
`35`		`- traceId: req.traceId,`
`36`	`34`	`spanId: req.spanId,`
`37`	`35`	`},`
`38`	`36`	`() => {`