feat(tracer): add try/catch logic to decorator and middleware close (#…

…1716)
aws-powertools · Sep 28, 2023 · be16b59 · be16b59
1 parent 793f12f
commit be16b59
Show file tree

Hide file tree

Showing 5 changed files with 284 additions and 131 deletions.
diff --git a/packages/tracer/src/Tracer.ts b/packages/tracer/src/Tracer.ts
@@ -406,8 +406,14 @@ class Tracer extends Utility implements TracerInterface {
               tracerRef.addErrorAsMetadata(error as Error);
               throw error;
             } finally {
-              subsegment?.close();
-              subsegment?.flush();
+              try {
+                subsegment?.close();
+              } catch (error) {
+                console.warn(
+                  `Failed to close or serialize segment, ${subsegment?.name}. We are catching the error but data might be lost.`,
+                  error
+                );
+              }
             }
 
             return result;
@@ -489,8 +495,14 @@ class Tracer extends Utility implements TracerInterface {
 
               throw error;
             } finally {
-              subsegment?.close();
-              subsegment?.flush();
+              try {
+                subsegment?.close();
+              } catch (error) {
+                console.warn(
+                  `Failed to close or serialize segment, ${subsegment?.name}. We are catching the error but data might be lost.`,
+                  error
+                );
+              }
             }
 
             return result;

diff --git a/packages/tracer/src/middleware/middy.ts b/packages/tracer/src/middleware/middy.ts
@@ -70,7 +70,14 @@ const captureLambdaHandler = (
     if (handlerSegment === undefined || lambdaSegment === null) {
       return;
     }
-    handlerSegment.close();
+    try {
+      handlerSegment.close();
+    } catch (error) {
+      console.warn(
+        `Failed to close or serialize segment, ${handlerSegment.name}. We are catching the error but data might be lost.`,
+        error
+      );
+    }
     target.setSegment(lambdaSegment);
   };
 

diff --git a/packages/tracer/src/provider/ProviderService.ts b/packages/tracer/src/provider/ProviderService.ts
@@ -103,6 +103,7 @@ class ProviderService implements ProviderServiceInterface {
 
       return;
     }
+
     segment.addMetadata(key, value, namespace);
   }
 

diff --git a/packages/tracer/tests/unit/Tracer.test.ts b/packages/tracer/tests/unit/Tracer.test.ts
@@ -914,6 +914,49 @@ describe('Class: Tracer', () => {
         otherDummyMethodSpy.mock.invocationCallOrder[0];
       expect(otherDummyCallOrder).toBeLessThan(dummyCallOrder);
     });
+
+    it('catches the error and logs a warning when a segment fails to close/serialize', async () => {
+      // Prepare
+      const tracer: Tracer = new Tracer();
+      const handlerSubsegment: Segment | Subsegment | undefined =
+        new Subsegment('### dummyMethod');
+      jest
+        .spyOn(tracer.provider, 'getSegment')
+        .mockImplementation(() => handlerSubsegment);
+      setContextMissingStrategy(() => null);
+      jest
+        .spyOn(tracer.provider, 'captureAsyncFunc')
+        .mockImplementation(async (methodName, callBackFn) => {
+          await callBackFn(handlerSubsegment);
+        });
+      const logWarningSpy = jest.spyOn(console, 'warn');
+      const closeSpy = jest
+        .spyOn(handlerSubsegment, 'close')
+        .mockImplementation(() => {
+          throw new Error('dummy error');
+        });
+
+      class Lambda implements LambdaInterface {
+        @tracer.captureLambdaHandler()
+        public async handler(
+          _event: unknown,
+          _context: Context
+        ): Promise<string> {
+          return 'foo bar';
+        }
+      }
+
+      // Act
+      await new Lambda().handler(event, context);
+
+      // Assess
+      expect(closeSpy).toHaveBeenCalledTimes(1);
+      expect(logWarningSpy).toHaveBeenNthCalledWith(
+        1,
+        `Failed to close or serialize segment, ${handlerSubsegment.name}. We are catching the error but data might be lost.`,
+        new Error('dummy error')
+      );
+    });
   });
 
   describe('Method: captureMethod', () => {
@@ -1328,6 +1371,53 @@ describe('Class: Tracer', () => {
         expect.anything()
       );
     });
+
+    it('catches the error and logs a warning when a segment fails to close/serialize', async () => {
+      // Prepare
+      const tracer: Tracer = new Tracer();
+      const handlerSubsegment: Segment | Subsegment | undefined =
+        new Subsegment('### dummyMethod');
+      jest
+        .spyOn(tracer.provider, 'getSegment')
+        .mockImplementation(() => handlerSubsegment);
+      setContextMissingStrategy(() => null);
+      jest
+        .spyOn(tracer.provider, 'captureAsyncFunc')
+        .mockImplementation(async (methodName, callBackFn) => {
+          await callBackFn(handlerSubsegment);
+        });
+      const logWarningSpy = jest.spyOn(console, 'warn');
+      const closeSpy = jest
+        .spyOn(handlerSubsegment, 'close')
+        .mockImplementation(() => {
+          throw new Error('dummy error');
+        });
+
+      class Lambda implements LambdaInterface {
+        @tracer.captureMethod()
+        public async dummyMethod(some: string): Promise<string> {
+          return some;
+        }
+
+        public async handler(
+          _event: unknown,
+          _context: Context
+        ): Promise<string> {
+          return await this.dummyMethod('foo bar');
+        }
+      }
+
+      // Act
+      await new Lambda().handler(event, context);
+
+      // Assess
+      expect(closeSpy).toHaveBeenCalledTimes(1);
+      expect(logWarningSpy).toHaveBeenNthCalledWith(
+        1,
+        `Failed to close or serialize segment, ${handlerSubsegment.name}. We are catching the error but data might be lost.`,
+        new Error('dummy error')
+      );
+    });
   });
 
   describe('Method: captureAWS', () => {