feat: add execution id support (#592)

When set LOG_EXECUTION_ID=True in environment variables or set --log-execution-id, execution id will be added to log along with trace and span id.
Nodejs version 13.0.0 is required for this feature.

README.md

@@ -190,6 +190,7 @@ ignored.
 | `--target`         | `FUNCTION_TARGET`         | The name of the exported function to be invoked in response to requests. Default: `function`                                                                                                                     |
 | `--signature-type` | `FUNCTION_SIGNATURE_TYPE` | The signature used when writing your function. Controls unmarshalling rules and determines which arguments are used to invoke your function. Default: `http`; accepted values: `http` or `event` or `cloudevent` |
 | `--source`         | `FUNCTION_SOURCE`         | The path to the directory of your function. Default: `cwd` (the current working directory)                                                                                                                       |
+| `--log-execution-id`| `LOG_EXECUTION_ID`        | Enables execution IDs in logs, either `true` or `false`. When not specified, default to disable. Requires Node.js 13.0.0 or later.                                                                                |
 
 You can set command-line flags in your `package.json` via the `start` script.
 For example:

docs/generated/api.json

@@ -1757,6 +1757,33 @@
           "name": "Request_2",
           "preserveMemberOrder": false,
           "members": [
+            {
+              "kind": "PropertySignature",
+              "canonicalReference": "@google-cloud/functions-framework!Request_2#executionId:member",
+              "docComment": "/**\n * Request-specified execution ID.\n */\n",
+              "excerptTokens": [
+                {
+                  "kind": "Content",
+                  "text": "executionId?: "
+                },
+                {
+                  "kind": "Content",
+                  "text": "string"
+                },
+                {
+                  "kind": "Content",
+                  "text": ";"
+                }
+              ],
+              "isReadonly": false,
+              "isOptional": true,
+              "releaseTag": "Public",
+              "name": "executionId",
+              "propertyTypeTokenRange": {
+                "startIndex": 1,
+                "endIndex": 2
+              }
+            },
             {
               "kind": "PropertySignature",
               "canonicalReference": "@google-cloud/functions-framework!Request_2#rawBody:member",
@@ -1784,6 +1811,60 @@
                 "startIndex": 1,
                 "endIndex": 2
               }
+            },
+            {
+              "kind": "PropertySignature",
+              "canonicalReference": "@google-cloud/functions-framework!Request_2#spanId:member",
+              "docComment": "/**\n * Cloud Trace span ID.\n */\n",
+              "excerptTokens": [
+                {
+                  "kind": "Content",
+                  "text": "spanId?: "
+                },
+                {
+                  "kind": "Content",
+                  "text": "string"
+                },
+                {
+                  "kind": "Content",
+                  "text": ";"
+                }
+              ],
+              "isReadonly": false,
+              "isOptional": true,
+              "releaseTag": "Public",
+              "name": "spanId",
+              "propertyTypeTokenRange": {
+                "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": [

docs/generated/api.md

@@ -112,7 +112,10 @@ export interface LegacyEvent {
 
 // @public (undocumented)
 interface Request_2 extends Request_3 {
+    executionId?: string;
     rawBody?: Buffer;
+    spanId?: string;
+    traceId?: string;
 }
 export { Request_2 as Request }
 

package.json

@@ -56,7 +56,7 @@
     "@types/body-parser": "1.19.5",
     "@types/minimist": "1.2.5",
     "@types/mocha": "9.1.1",
-    "@types/node": "20.11.15",
+    "@types/node": "^20.11.24",
     "@types/on-finished": "2.3.4",
     "@types/semver": "^7.3.6",
     "@types/sinon": "^10.0.0",

src/async_local_storage.ts

@@ -0,0 +1,49 @@
+import * as semver from 'semver';
+import {Request, Response} from './functions';
+import {NextFunction} from 'express';
+import {requiredNodeJsVersionForLogExecutionID} from './options';
+import type {AsyncLocalStorage} from 'node:async_hooks';
+
+export interface ExecutionContext {
+  executionId?: string;
+  traceId?: string;
+  spanId?: string;
+}
+
+let asyncLocalStorage: AsyncLocalStorage<ExecutionContext> | undefined;
+
+export async function asyncLocalStorageMiddleware(
+  req: Request,
+  res: Response,
+  next: NextFunction
+) {
+  if (
+    semver.lt(process.versions.node, requiredNodeJsVersionForLogExecutionID)
+  ) {
+    // Skip for unsupported Node.js version.
+    next();
+    return;
+  }
+  if (!asyncLocalStorage) {
+    const asyncHooks = await import('node:async_hooks');
+    asyncLocalStorage = new asyncHooks.AsyncLocalStorage();
+  }
+
+  asyncLocalStorage.run(
+    {
+      executionId: req.executionId,
+      traceId: req.traceId,
+      spanId: req.spanId,
+    },
+    () => {
+      next();
+    }
+  );
+}
+
+export function getCurrentContext(): ExecutionContext | undefined {
+  if (!asyncLocalStorage) {
+    return undefined;
+  }
+  return asyncLocalStorage.getStore();
+}

src/execution_context.ts

@@ -0,0 +1,37 @@
+import {Request, Response} from './functions';
+import {NextFunction} from 'express';
+
+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>.+)$/;
+
+function generateExecutionId() {
+  const timestampPart = Date.now().toString(36).slice(-6);
+  const randomPart = Math.random().toString(36).slice(-6);
+  return timestampPart + randomPart;
+}
+
+export const executionContextMiddleware = (
+  req: Request,
+  res: Response,
+  next: NextFunction
+) => {
+  let executionId = req.header(FUNCTION_EXECUTION_ID_HEADER_KEY);
+  if (!executionId) {
+    executionId = generateExecutionId();
+  }
+  req.executionId = executionId;
+
+  const cloudTraceContext = req.header(TRACE_CONTEXT_HEADER_KEY);
+  if (cloudTraceContext) {
+    const match = cloudTraceContext.match(TRACE_CONTEXT_PATTERN);
+    if (match?.groups) {
+      const {traceId, spanId} = match.groups;
+      req.traceId = traceId;
+      req.spanId = spanId;
+    }
+  }
+  next();
+};

src/functions.ts

@@ -32,6 +32,18 @@ export interface Request extends ExpressRequest {
    * A buffer which provides access to the request's raw HTTP body.
    */
   rawBody?: Buffer;
+  /**
+   * Request-specified execution ID.
+   */
+  executionId?: string;
+  /**
+   * Cloud Trace trace ID.
+   */
+  traceId?: string;
+  /**
+   * Cloud Trace span ID.
+   */
+  spanId?: string;
 }
 
 /**

src/logger.ts

@@ -14,6 +14,13 @@
 
 import * as express from 'express';
 import {FUNCTION_STATUS_HEADER_FIELD} from './types';
+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';
 
 /**
  * Logs an error message and sends back an error response to the incoming
@@ -60,3 +67,120 @@ export function sendCrashResponse({
     callback();
   }
 }
+
+export function loggingHandlerAddExecutionContext() {
+  interceptStdoutWrite();
+  interceptStderrWrite();
+}
+
+function interceptStdoutWrite() {
+  const originalStdoutWrite = process.stdout.write;
+  process.stdout.write = (data, ...args) => {
+    const {encoding, cb} = splitArgs(args);
+    const modifiedData = getModifiedData(data, encoding);
+    return originalStdoutWrite.apply(process.stdout, [modifiedData, cb]);
+  };
+}
+
+function interceptStderrWrite() {
+  const originalStderrWrite = process.stderr.write;
+  process.stderr.write = (data, ...args) => {
+    const {encoding, cb} = splitArgs(args);
+    const modifiedData = getModifiedData(data, encoding, true);
+    return originalStderrWrite.apply(process.stderr, [modifiedData, cb]);
+  };
+}
+
+export const errorHandler = (
+  err: Error | any,
+  req: express.Request,
+  res: express.Response,
+  next: express.NextFunction
+) => {
+  interceptStderrWrite();
+  res.status(500);
+  res.render('error', {error: err});
+};
+
+export function splitArgs(args: any[]) {
+  let encoding, cb;
+  if (
+    args.length > 0 &&
+    (Buffer.isEncoding(args[0]) || typeof args[0] === 'undefined')
+  ) {
+    encoding = args[0];
+    args.shift();
+  }
+  if (args.length > 0 && typeof args[0] === 'function') {
+    cb = args[0];
+  }
+  return {encoding: encoding, cb: cb};
+}
+
+export function getModifiedData(
+  data: Uint8Array | string,
+  encoding?: BufferEncoding,
+  stderr = false
+) {
+  const currentContext = getCurrentContext();
+  if (!currentContext) {
+    return data;
+  }
+  const {isJSON, processedData} = processData(data, encoding);
+  let dataWithContext;
+  if (isJSON) {
+    dataWithContext = getJSONWithContext(processedData, currentContext);
+  } else {
+    dataWithContext = getTextWithContext(processedData, currentContext);
+  }
+  if (stderr) {
+    dataWithContext[SEVERITY] = 'ERROR';
+  }
+
+  return JSON.stringify(dataWithContext) + '\n';
+}
+
+function getTextWithContext(
+  data: Uint8Array | string,
+  context: ExecutionContext
+) {
+  return {
+    message: data,
+    [EXECUTION_CONTEXT_LABELS_KEY]: {execution_id: context.executionId},
+    [EXECUTION_CONTEXT_TRACE_KEY]: context.traceId,
+    [EXECUTION_CONTEXT_SPAN_ID_KEY]: context.spanId,
+  };
+}
+
+function getJSONWithContext(json: any, context: ExecutionContext) {
+  if (EXECUTION_CONTEXT_LABELS_KEY in json) {
+    json[EXECUTION_CONTEXT_LABELS_KEY]['execution_id'] = context.executionId;
+  } 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,
+  };
+}
+
+function processData(data: Uint8Array | string, encoding?: BufferEncoding) {
+  let decodedData;
+  try {
+    if (data instanceof Uint8Array) {
+      decodedData = Buffer.from(data.buffer).toString();
+    } else {
+      decodedData = Buffer.from(data, encoding).toString();
+    }
+  } catch (e) {
+    // Failed to decode, treat it as simple text.
+    return {isJSON: false, processedData: data};
+  }
+
+  try {
+    return {isJSON: true, processedData: JSON.parse(decodedData)};
+  } catch (e) {
+    return {isJSON: false, processedData: decodedData};
+  }
+}