refactor(tracer): log warning instead of throwing when segment is not found (#1370)

* improv: moved logic to provider + changed middleware behavior around absent segment

* tests: fixed tests for middleware

* tests: fixed tests for main class

* tests: fixed provider service tests

* chore: bump xray-sdk-core

* chore: update jsdoc for addErrorAsMetadata

* chore: update jsdoc for getSegment

* chore: update jsdoc for getSegment

Author: dreamorosi

package-lock.json

@@ -52,7 +52,7 @@
   },
   "dependencies": {
     "@aws-lambda-powertools/commons": "^1.6.0",
-    "aws-xray-sdk-core": "^3.4.0"
+    "aws-xray-sdk-core": "^3.4.1"
   },
   "keywords": [
     "aws",
@@ -63,4 +63,4 @@
     "serverless",
     "nodejs"
   ]
-}
+}

packages/tracer/package.json

@@ -38,7 +38,7 @@ import { Segment, Subsegment } from 'aws-xray-sdk-core';
  * 
  * const tracer = new Tracer({ serviceName: 'serverlessAirline' });
  * 
-* const lambdaHandler = async (_event: any, _context: any) => {
+ * const lambdaHandler = async (_event: any, _context: any) => {
  *   ...
  * };
  * 
@@ -153,20 +153,26 @@ class Tracer extends Utility implements TracerInterface {
     * @see https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-errors
     *
     * @param error - Error to serialize as metadata
+    * @param [remote] - Whether the error was thrown by a remote service. Defaults to `false`
     */
-  public addErrorAsMetadata(error: Error): void {
+  public addErrorAsMetadata(error: Error, remote?: boolean): void {
     if (!this.isTracingEnabled()) {
       return;
     }
 
     const subsegment = this.getSegment();
+    if (subsegment === undefined) {
+
+      return;
+    }
+
     if (!this.captureError) {
       subsegment.addErrorFlag();
 
       return;
     }
 
-    subsegment.addError(error, false);
+    subsegment.addError(error, remote || false);
   }
 
   /**
@@ -506,7 +512,7 @@ class Tracer extends Utility implements TracerInterface {
   }
 
   /**
-   * Get the active segment or subsegment in the current scope.
+   * Get the active segment or subsegment (if any) in the current scope.
    * 
    * Usually you won't need to call this method unless you are creating custom subsegments or using manual mode.
    *
@@ -525,15 +531,17 @@ class Tracer extends Utility implements TracerInterface {
    * }
    * ```
    * 
-   * @returns segment - The active segment or subsegment in the current scope.
+   * @returns The active segment or subsegment in the current scope. Will log a warning and return `undefined` if no segment is found.
    */
-  public getSegment(): Segment | Subsegment {
+  public getSegment(): Segment | Subsegment | undefined {
     if (!this.isTracingEnabled()) {
       return new Subsegment('## Dummy segment');
     }
     const segment = this.provider.getSegment();    
     if (segment === undefined) {
-      throw new Error('Failed to get the current sub/segment from the context.');
+      console.warn(
+        'Failed to get the current sub/segment from the context, this is likely because you are not using the Tracer in a Lambda function.'
+      );
     }
 
     return segment;
@@ -573,13 +581,7 @@ class Tracer extends Utility implements TracerInterface {
   public putAnnotation(key: string, value: string | number | boolean): void {
     if (!this.isTracingEnabled()) return;
 
-    const document = this.getSegment();
-    if (document instanceof Segment) {
-      console.warn('You cannot annotate the main segment in a Lambda execution environment');
-      
-      return;
-    }
-    document.addAnnotation(key, value);
+    this.provider.putAnnotation(key, value);
   }
 
   /**
@@ -606,15 +608,7 @@ class Tracer extends Utility implements TracerInterface {
   public putMetadata(key: string, value: unknown, namespace?: string | undefined): void {
     if (!this.isTracingEnabled()) return;
 
-    const document = this.getSegment();
-    if (document instanceof Segment) {
-      console.warn('You cannot add metadata to the main segment in a Lambda execution environment');
-      
-      return;
-    }
-    
-    namespace = namespace || this.serviceName;
-    document.addMetadata(key, value, namespace);
+    this.provider.putMetadata(key, value, namespace || this.serviceName);
   }
 
   /**

packages/tracer/src/Tracer.ts

@@ -2,7 +2,7 @@ import { CaptureLambdaHandlerOptions, CaptureMethodOptions, HandlerMethodDecorat
 import { Segment, Subsegment } from 'aws-xray-sdk-core';
 
 interface TracerInterface {
-  addErrorAsMetadata(error: Error): void
+  addErrorAsMetadata(error: Error, remote?: boolean): void
   addResponseAsMetadata(data?: unknown, methodName?: string): void
   addServiceNameAnnotation(): void
   annotateColdStart(): void
@@ -11,7 +11,7 @@ interface TracerInterface {
   captureAWSClient<T>(service: T): void | T
   captureLambdaHandler(options?: CaptureLambdaHandlerOptions): HandlerMethodDecorator
   captureMethod(options?: CaptureMethodOptions): MethodDecorator
-  getSegment(): Segment | Subsegment
+  getSegment(): Segment | Subsegment | undefined
   getRootXrayTraceId(): string | undefined
   isTracingEnabled(): boolean
   putAnnotation: (key: string, value: string | number | boolean) => void

packages/tracer/src/TracerInterface.ts

@@ -34,18 +34,26 @@ import type {
  * @returns middleware - The middy middleware object
  */
 const captureLambdaHandler = (target: Tracer, options?: CaptureLambdaHandlerOptions): MiddlewareLikeObj => {
-  let lambdaSegment: Subsegment | Segment;
+  let lambdaSegment: Segment;
+  let handlerSegment: Subsegment;
 
   const open = (): void => {
-    lambdaSegment = target.getSegment();
-    const handlerSegment = lambdaSegment.addNewSubsegment(`## ${process.env._HANDLER}`);
+    const segment = target.getSegment();
+    if (segment === undefined) {
+      return;
+    }
+    // If segment is defined, then it is a Segment as this middleware is only used for Lambda Handlers
+    lambdaSegment = segment as Segment;
+    handlerSegment = lambdaSegment.addNewSubsegment(`## ${process.env._HANDLER}`);
     target.setSegment(handlerSegment);
   };
 
   const close = (): void => {
-    const subsegment = target.getSegment();
-    subsegment.close();
-    target.setSegment(lambdaSegment as Segment);
+    if (handlerSegment === undefined || lambdaSegment === null) {
+      return;
+    }
+    handlerSegment.close();
+    target.setSegment(lambdaSegment);
   };
 
   const captureLambdaHandlerBefore = async (): Promise<void> => {

packages/tracer/src/middleware/middy.ts

@@ -1,10 +1,27 @@
-import { ContextMissingStrategy } from 'aws-xray-sdk-core/dist/lib/context_utils';
+import {
+  ContextMissingStrategy
+} from 'aws-xray-sdk-core/dist/lib/context_utils';
 import { Namespace } from 'cls-hooked';
 import { ProviderServiceInterface } from '.';
-import { captureAWS, captureAWSClient, captureAWSv3Client, captureAsyncFunc, captureFunc, captureHTTPsGlobal, getNamespace, getSegment, setSegment, Segment, Subsegment, setContextMissingStrategy, setDaemonAddress, setLogger, Logger } from 'aws-xray-sdk-core';
+import {
+  captureAWS,
+  captureAWSClient,
+  captureAWSv3Client,
+  captureAsyncFunc,
+  captureFunc,
+  captureHTTPsGlobal,
+  getNamespace,
+  getSegment,
+  setSegment,
+  Segment,
+  Subsegment,
+  setContextMissingStrategy,
+  setDaemonAddress,
+  setLogger,
+  Logger
+} from 'aws-xray-sdk-core';
 
 class ProviderService implements ProviderServiceInterface {
-  
   public captureAWS<T>(awssdk: T): T {
     return captureAWS(awssdk);
   }
@@ -42,6 +59,36 @@ class ProviderService implements ProviderServiceInterface {
     return getSegment();
   }
 
+  public putAnnotation(key: string, value: string | number | boolean): void {
+    const segment = this.getSegment();
+    if (segment === undefined) {
+      console.warn('No active segment or subsegment found, skipping annotation');
+
+      return;
+    }
+    if (segment instanceof Segment) {
+      console.warn('You cannot annotate the main segment in a Lambda execution environment');
+      
+      return;
+    }
+    segment.addAnnotation(key, value);
+  }
+
+  public putMetadata(key: string, value: unknown, namespace?: string): void {
+    const segment = this.getSegment();
+    if (segment === undefined) {
+      console.warn('No active segment or subsegment found, skipping metadata addition');
+
+      return;
+    }
+    if (segment instanceof Segment) {
+      console.warn('You cannot add metadata to the main segment in a Lambda execution environment');
+      
+      return;
+    }
+    segment.addMetadata(key, value, namespace);
+  }
+
   public setContextMissingStrategy(strategy: unknown): void {
     setContextMissingStrategy(strategy as ContextMissingStrategy);
   }

packages/tracer/src/provider/ProviderService.ts

@@ -25,6 +25,10 @@ interface ProviderServiceInterface {
   captureFunc(name: string, fcn: (subsegment?: Subsegment) => unknown, parent?: Segment | Subsegment): unknown
 
   captureHTTPsGlobal(): void
+
+  putAnnotation(key: string, value: string | number | boolean): void
+
+  putMetadata(key: string, value: unknown, namespace?: string): void
 }
 
 export {