[v7] Make all serverless tracing wrappers use provided client directly

kamilogorek · kamilogorek · commit 27f8b4359652 · 2021-03-23T15:46:07.000+01:00
diff --git a/packages/serverless/src/awsservices.ts b/packages/serverless/src/awsservices.ts
@@ -1,4 +1,3 @@
-import { getTransaction } from '@sentry/minimal';
 import { ClientLike, Integration, Span } from '@sentry/types';
 import { fill } from '@sentry/utils';
 // 'aws-sdk/global' import is expected to be type-only so it's erased in the final .js file.
@@ -19,6 +18,8 @@ interface AWSService {
 export class AWSServices implements Integration {
   public name = this.constructor.name;
 
+  private _client!: ClientLike;
+
   private readonly _optional: boolean;
 
   public constructor(options: { optional?: boolean } = {}) {
@@ -28,47 +29,54 @@ export class AWSServices implements Integration {
   /**
    * @inheritDoc
    */
-  public install(_client: ClientLike): void {
+  public install(client: ClientLike): void {
+    this._client = client;
+
     try {
       // eslint-disable-next-line @typescript-eslint/no-var-requires
       const awsModule = require('aws-sdk/global') as typeof AWS;
-      fill(awsModule.Service.prototype, 'makeRequest', wrapMakeRequest);
+      fill(awsModule.Service.prototype, 'makeRequest', this._wrapMakeRequest.bind(this));
     } catch (e) {
       if (!this._optional) {
         throw e;
       }
     }
   }
-}
 
-/** */
-function wrapMakeRequest<TService extends AWSService, TResult>(
-  orig: MakeRequestFunction<GenericParams, TResult>,
-): MakeRequestFunction<GenericParams, TResult> {
-  return function(this: TService, operation: string, params?: GenericParams, callback?: MakeRequestCallback<TResult>) {
-    // TODO: Use clients scope instead of global call
-    const transaction = getTransaction();
-    let span: Span | undefined;
-    const req = orig.call(this, operation, params);
-    req.on('afterBuild', () => {
-      if (transaction) {
-        span = transaction.startChild({
-          description: describe(this, operation, params),
-          op: 'aws.request',
-        });
-      }
-    });
-    req.on('complete', () => {
-      if (span) {
-        span.finish();
-      }
-    });
+  private _wrapMakeRequest<TService extends AWSService, TResult>(
+    orig: MakeRequestFunction<GenericParams, TResult>,
+  ): MakeRequestFunction<GenericParams, TResult> {
+    const client = this._client;
 
-    if (callback) {
-      req.send(callback);
-    }
-    return req;
-  };
+    return function(
+      this: TService,
+      operation: string,
+      params?: GenericParams,
+      callback?: MakeRequestCallback<TResult>,
+    ) {
+      const transaction = client.getScope().getTransaction();
+      let span: Span | undefined;
+      const req = orig.call(this, operation, params);
+      req.on('afterBuild', () => {
+        if (transaction) {
+          span = transaction.startChild({
+            description: describe(this, operation, params),
+            op: 'aws.request',
+          });
+        }
+      });
+      req.on('complete', () => {
+        if (span) {
+          span.finish();
+        }
+      });
+
+      if (callback) {
+        req.send(callback);
+      }
+      return req;
+    };
+  }
 }
 
 /** Describes an operation on generic AWS service */
diff --git a/packages/serverless/src/google-cloud-grpc.ts b/packages/serverless/src/google-cloud-grpc.ts
@@ -2,7 +2,6 @@ import { EventEmitter } from 'events';
 
 import { ClientLike, Integration, Span } from '@sentry/types';
 import { fill } from '@sentry/utils';
-import { getTransaction } from '@sentry/minimal';
 
 interface GrpcFunction extends CallableFunction {
   (...args: unknown[]): EventEmitter;
@@ -30,6 +29,8 @@ interface Stub {
 export class GoogleCloudGrpc implements Integration {
   public name = this.constructor.name;
 
+  private _client!: ClientLike;
+
   private readonly _optional: boolean;
 
   public constructor(options: { optional?: boolean } = {}) {
@@ -39,41 +40,54 @@ export class GoogleCloudGrpc implements Integration {
   /**
    * @inheritDoc
    */
-  public install(_client: ClientLike): void {
+  public install(client: ClientLike): void {
+    this._client = client;
     try {
       // eslint-disable-next-line @typescript-eslint/no-var-requires
       const gaxModule = require('google-gax');
       fill(
         gaxModule.GrpcClient.prototype, // eslint-disable-line @typescript-eslint/no-unsafe-member-access
         'createStub',
-        wrapCreateStub,
+        this._wrapCreateStub.bind(this),
       );
     } 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;
-  };
+  /** Returns a wrapped function that returns a stub with tracing enabled */
+  private _wrapCreateStub(origCreate: CreateStubFunc): CreateStubFunc {
+    const client = this._client;
+
+    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({ client, stub, serviceIdentifier, methodName });
+      }
+      return stub;
+    };
+  }
 }
 
 /** Patches the function in grpc stub to enable tracing */
-function fillGrpcFunction(stub: Stub, serviceIdentifier: string, methodName: string): void {
+function fillGrpcFunction({
+  client,
+  stub,
+  serviceIdentifier,
+  methodName,
+}: {
+  client: ClientLike;
+  stub: Stub;
+  serviceIdentifier: string;
+  methodName: string;
+}): void {
   const funcObj = stub[methodName];
   if (typeof funcObj !== 'function') {
     return;
@@ -97,8 +111,7 @@ function fillGrpcFunction(stub: Stub, serviceIdentifier: string, methodName: str
       if (typeof ret?.on !== 'function') {
         return ret;
       }
-      // TODO: Use clients scope instead of global call
-      const transaction = getTransaction();
+      const transaction = client.getScope().getTransaction();
       let span: Span | undefined;
       if (transaction) {
         span = transaction.startChild({
diff --git a/packages/serverless/src/google-cloud-http.ts b/packages/serverless/src/google-cloud-http.ts
@@ -1,7 +1,6 @@
 // '@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 { getTransaction } from '@sentry/minimal';
 import { ClientLike, Integration, Span } from '@sentry/types';
 import { fill } from '@sentry/utils';
 
@@ -16,6 +15,8 @@ interface RequestFunction extends CallableFunction {
 export class GoogleCloudHttp implements Integration {
   public name = this.constructor.name;
 
+  private _client!: ClientLike;
+
   private readonly _optional: boolean;
 
   public constructor(options: { optional?: boolean } = {}) {
@@ -25,39 +26,42 @@ export class GoogleCloudHttp implements Integration {
   /**
    * @inheritDoc
    */
-  public install(_client: ClientLike): void {
+  public install(client: ClientLike): void {
+    this._client = client;
+
     try {
       // eslint-disable-next-line @typescript-eslint/no-var-requires
       const commonModule = require('@google-cloud/common') as typeof common;
-      fill(commonModule.Service.prototype, 'request', wrapRequestFunction);
+      fill(commonModule.Service.prototype, 'request', this._wrapRequestFunction.bind(this));
     } 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 {
-    // TODO: Use clients scope instead of global call
-    const transaction = getTransaction();
-    let span: Span | undefined;
-    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();
+  /** Returns a wrapped function that makes a request with tracing enabled */
+  private _wrapRequestFunction(orig: RequestFunction): RequestFunction {
+    const client = this._client;
+
+    return function(this: common.Service, reqOpts: RequestOptions, callback: ResponseCallback): void {
+      const transaction = client.getScope().getTransaction();
+      let span: Span | undefined;
+      if (transaction) {
+        const httpMethod = reqOpts.method || 'GET';
+        span = transaction.startChild({
+          description: `${httpMethod} ${reqOpts.uri}`,
+          op: `gcloud.http.${identifyService(this.apiEndpoint)}`,
+        });
       }
-      callback(...args);
-    });
-  };
+      orig.call(this, reqOpts, (...args: Parameters<ResponseCallback>) => {
+        if (span) {
+          span.finish();
+        }
+        callback(...args);
+      });
+    };
+  }
 }
 
 /** Identifies service by its base url */
