Implement tracing of google cloud requests (#2981)

Co-authored-by: Kamil Ogórek <[email protected]>

packages/serverless/package.json

@@ -23,12 +23,16 @@
     "tslib": "^1.9.3"
   },
   "devDependencies": {
+    "@google-cloud/common": "^3.4.1",
+    "@google-cloud/pubsub": "^2.5.0",
+    "@google-cloud/bigquery": "^5.3.0",
     "@google-cloud/functions-framework": "^1.7.1",
     "@sentry-internal/eslint-config-sdk": "5.26.0",
     "@types/aws-lambda": "^8.10.62",
     "@types/node": "^14.6.4",
     "aws-sdk": "^2.765.0",
     "eslint": "7.6.0",
+    "google-gax": "^2.9.0",
     "jest": "^24.7.1",
     "nock": "^13.0.4",
     "npm-run-all": "^4.1.2",

packages/serverless/src/awsservices.ts

@@ -50,7 +50,7 @@ export class AWSServices implements Integration {
             if (transaction) {
               span = transaction.startChild({
                 description: describe(this, operation, params),
-                op: 'request',
+                op: 'aws.request',
               });
             }
           });

packages/serverless/src/gcpfunction/index.ts

@@ -1,15 +1,28 @@
 import * as Sentry from '@sentry/node';
+import { Integration } from '@sentry/types';
+
+import { GoogleCloudGrpc } from '../google-cloud-grpc';
+import { GoogleCloudHttp } from '../google-cloud-http';
 
 import { serverlessEventProcessor } from '../utils';
 
 export * from './http';
 export * from './events';
 export * from './cloud_events';
 
+export const defaultIntegrations: Integration[] = [
+  ...Sentry.defaultIntegrations,
+  new GoogleCloudHttp({ optional: true }), // We mark this integration optional since '@google-cloud/common' module could be missing.
+  new GoogleCloudGrpc({ optional: true }), // We mark this integration optional since 'google-gax' module could be missing.
+];
+
 /**
  * @see {@link Sentry.init}
  */
 export function init(options: Sentry.NodeOptions = {}): void {
+  if (options.defaultIntegrations === undefined) {
+    options.defaultIntegrations = defaultIntegrations;
+  }
   Sentry.init(options);
   Sentry.addGlobalEventProcessor(serverlessEventProcessor('GCPFunction'));
 }

packages/serverless/src/google-cloud-grpc.ts

@@ -0,0 +1,132 @@
+import { getCurrentHub } from '@sentry/node';
+import { Integration, Span, Transaction } from '@sentry/types';
+import { fill } from '@sentry/utils';
+import { EventEmitter } from 'events';
+
+interface GrpcFunction extends CallableFunction {
+  (...args: unknown[]): EventEmitter;
+}
+
+interface GrpcFunctionObject extends GrpcFunction {
+  requestStream: boolean;
+  responseStream: boolean;
+  originalName: string;
+}
+
+interface StubOptions {
+  servicePath?: string;
+}
+
+interface CreateStubFunc extends CallableFunction {
+  (createStub: unknown, options: StubOptions): PromiseLike<Stub>;
+}
+
+interface Stub {
+  [key: string]: GrpcFunctionObject;
+}
+
+/** Google Cloud Platform service requests tracking for GRPC APIs */
+export class GoogleCloudGrpc implements Integration {
+  /**
+   * @inheritDoc
+   */
+  public static id: string = 'GoogleCloudGrpc';
+
+  /**
+   * @inheritDoc
+   */
+  public name: string = GoogleCloudGrpc.id;
+
+  private readonly _optional: boolean;
+
+  public constructor(options: { optional?: boolean } = {}) {
+    this._optional = options.optional || false;
+  }
+
+  /**
+   * @inheritDoc
+   */
+  public setupOnce(): void {
+    try {
+      const gaxModule = require('google-gax');
+      fill(
+        gaxModule.GrpcClient.prototype, // eslint-disable-line @typescript-eslint/no-unsafe-member-access
+        'createStub',
+        wrapCreateStub,
+      );
+    } catch (e) {
+      if (!this._optional) {
+        throw e;
+      }
+    }
+  }
+}
+
+/** Returns a wrapped function that returns a stub with tracing enabled */
+function wrapCreateStub(origCreate: CreateStubFunc): CreateStubFunc {
+  return async function(this: unknown, ...args: Parameters<CreateStubFunc>) {
+    const servicePath = args[1]?.servicePath;
+    if (servicePath == null || servicePath == undefined) {
+      return origCreate.apply(this, args);
+    }
+    const serviceIdentifier = identifyService(servicePath);
+    const stub = await origCreate.apply(this, args);
+    for (const methodName of Object.keys(Object.getPrototypeOf(stub))) {
+      fillGrpcFunction(stub, serviceIdentifier, methodName);
+    }
+    return stub;
+  };
+}
+
+/** Patches the function in grpc stub to enable tracing */
+function fillGrpcFunction(stub: Stub, serviceIdentifier: string, methodName: string): void {
+  const funcObj = stub[methodName];
+  if (typeof funcObj !== 'function') {
+    return;
+  }
+  const callType =
+    !funcObj.requestStream && !funcObj.responseStream
+      ? 'unary call'
+      : funcObj.requestStream && !funcObj.responseStream
+      ? 'client stream'
+      : !funcObj.requestStream && funcObj.responseStream
+      ? 'server stream'
+      : 'bidi stream';
+  if (callType != 'unary call') {
+    return;
+  }
+  fill(
+    stub,
+    methodName,
+    (orig: GrpcFunction): GrpcFunction => (...args) => {
+      const ret = orig.apply(stub, args);
+      if (typeof ret?.on !== 'function') {
+        return ret;
+      }
+      let transaction: Transaction | undefined;
+      let span: Span | undefined;
+      const scope = getCurrentHub().getScope();
+      if (scope) {
+        transaction = scope.getTransaction();
+      }
+      if (transaction) {
+        span = transaction.startChild({
+          description: `${callType} ${methodName}`,
+          op: `gcloud.grpc.${serviceIdentifier}`,
+        });
+      }
+      ret.on('status', () => {
+        if (span) {
+          span.finish();
+        }
+      });
+      return ret;
+    },
+  );
+}
+
+/** Identifies service by its address */
+function identifyService(servicePath: string): string {
+  const match = servicePath.match(/^(\w+)\.googleapis.com$/);
+  return match ? match[1] : servicePath;
+}

packages/serverless/src/google-cloud-http.ts

@@ -0,0 +1,77 @@
+// '@google-cloud/common' import is expected to be type-only so it's erased in the final .js file.
+// When TypeScript compiler is upgraded, use `import type` syntax to explicitly assert that we don't want to load a module here.
+import * as common from '@google-cloud/common';
+import { getCurrentHub } from '@sentry/node';
+import { Integration, Span, Transaction } from '@sentry/types';
+import { fill } from '@sentry/utils';
+
+type RequestOptions = common.DecorateRequestOptions;
+type ResponseCallback = common.BodyResponseCallback;
+// This interace could be replaced with just type alias once the `strictBindCallApply` mode is enabled.
+interface RequestFunction extends CallableFunction {
+  (reqOpts: RequestOptions, callback: ResponseCallback): void;
+}
+
+/** Google Cloud Platform service requests tracking for RESTful APIs */
+export class GoogleCloudHttp implements Integration {
+  /**
+   * @inheritDoc
+   */
+  public static id: string = 'GoogleCloudHttp';
+
+  /**
+   * @inheritDoc
+   */
+  public name: string = GoogleCloudHttp.id;
+
+  private readonly _optional: boolean;
+
+  public constructor(options: { optional?: boolean } = {}) {
+    this._optional = options.optional || false;
+  }
+
+  /**
+   * @inheritDoc
+   */
+  public setupOnce(): void {
+    try {
+      const commonModule = require('@google-cloud/common') as typeof common;
+      fill(commonModule.Service.prototype, 'request', wrapRequestFunction);
+    } catch (e) {
+      if (!this._optional) {
+        throw e;
+      }
+    }
+  }
+}
+
+/** Returns a wrapped function that makes a request with tracing enabled */
+function wrapRequestFunction(orig: RequestFunction): RequestFunction {
+  return function(this: common.Service, reqOpts: RequestOptions, callback: ResponseCallback): void {
+    let transaction: Transaction | undefined;
+    let span: Span | undefined;
+    const scope = getCurrentHub().getScope();
+    if (scope) {
+      transaction = scope.getTransaction();
+    }
+    if (transaction) {
+      const httpMethod = reqOpts.method || 'GET';
+      span = transaction.startChild({
+        description: `${httpMethod} ${reqOpts.uri}`,
+        op: `gcloud.http.${identifyService(this.apiEndpoint)}`,
+      });
+    }
+    orig.call(this, reqOpts, (...args: Parameters<ResponseCallback>) => {
+      if (span) {
+        span.finish();
+      }
+      callback(...args);
+    });
+  };
+}
+
+/** Identifies service by its base url */
+function identifyService(apiEndpoint: string): string {
+  const match = apiEndpoint.match(/^https:\/\/(\w+)\.googleapis.com$/);
+  return match ? match[1] : apiEndpoint.replace(/^(http|https)?:\/\//, '');
+}

packages/serverless/test/__mocks__/dns.ts

@@ -0,0 +1,2 @@
+export const lookup = jest.fn();
+export const resolveTxt = jest.fn();

packages/serverless/test/awsservices.test.ts

@@ -33,7 +33,10 @@ describe('AWSServices', () => {
       const data = await s3.getObject({ Bucket: 'foo', Key: 'bar' }).promise();
       expect(data.Body?.toString('utf-8')).toEqual('contents');
       // @ts-ignore see "Why @ts-ignore" note
-      expect(Sentry.fakeTransaction.startChild).toBeCalledWith({ op: 'request', description: 'aws.s3.getObject foo' });
+      expect(Sentry.fakeTransaction.startChild).toBeCalledWith({
+        op: 'aws.request',
+        description: 'aws.s3.getObject foo',
+      });
       // @ts-ignore see "Why @ts-ignore" note
       expect(Sentry.fakeSpan.finish).toBeCalled();
     });
@@ -49,7 +52,10 @@ describe('AWSServices', () => {
         done();
       });
       // @ts-ignore see "Why @ts-ignore" note
-      expect(Sentry.fakeTransaction.startChild).toBeCalledWith({ op: 'request', description: 'aws.s3.getObject foo' });
+      expect(Sentry.fakeTransaction.startChild).toBeCalledWith({
+        op: 'aws.request',
+        description: 'aws.s3.getObject foo',
+      });
     });
   });
 
@@ -63,7 +69,10 @@ describe('AWSServices', () => {
       const data = await lambda.invoke({ FunctionName: 'foo' }).promise();
       expect(data.Payload?.toString('utf-8')).toEqual('reply');
       // @ts-ignore see "Why @ts-ignore" note
-      expect(Sentry.fakeTransaction.startChild).toBeCalledWith({ op: 'request', description: 'aws.lambda.invoke foo' });
+      expect(Sentry.fakeTransaction.startChild).toBeCalledWith({
+        op: 'aws.request',
+        description: 'aws.lambda.invoke foo',
+      });
     });
   });
 });