From ebc9ea8ecf86d9c27c7ee3aa59a3c09952d96c6f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Serkan=20=C3=96ZAL?= <sozal@catchpoint.com>
Date: Sat, 16 Nov 2024 11:50:00 +0300
Subject: [PATCH] In AWS Lambda instrumentation, extend invocation span start
 time back to Lambda environment startup time on coldstart

---
 .../src/instrumentation.ts                    |  3 ++
 .../test/integrations/lambda-handler.test.ts  | 40 ++++++++++++++++---
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts
index 23755bb264..ef5c87d2e1 100644
--- a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts
+++ b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts
@@ -232,6 +232,9 @@ export class AwsLambdaInstrumentation extends InstrumentationBase<AwsLambdaInstr
       const span = plugin.tracer.startSpan(
         name,
         {
+          // For coldstarts,
+          // use Lambda start time as invocation span start time to take initialization time into account.
+          startTime: requestIsColdStart ? lambdaStartTime : Date.now(),
           kind: SpanKind.SERVER,
           attributes: {
             [SEMATTRS_FAAS_EXECUTION]: context.awsRequestId,
diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts
index 7b338cbd6d..8bf176fbeb 100644
--- a/plugins/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts
+++ b/plugins/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts
@@ -51,7 +51,7 @@ import {
   defaultTextMapGetter,
 } from '@opentelemetry/api';
 import { AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
-import { W3CTraceContextPropagator } from '@opentelemetry/core';
+import { hrTimeToMilliseconds, W3CTraceContextPropagator } from '@opentelemetry/core';
 import { AWSXRayLambdaPropagator } from '@opentelemetry/propagator-aws-xray-lambda';
 
 const memoryExporter = new InMemorySpanExporter();
@@ -266,10 +266,15 @@ describe('lambda handler', () => {
     });
 
     it('should record coldstart', async () => {
-      initializeHandler('lambda-test/sync.handler');
+      const time0 = Date.now();
+      initializeHandler('lambda-test/sync.handler', { lambdaStartTime: time0 });
+
+      // Simulate initialization time
+      await new Promise(resolve => setTimeout(resolve, 10));
 
       const handlerModule = lambdaRequire('lambda-test/sync');
 
+      const time1 = Date.now();
       const result1 = await new Promise((resolve, reject) => {
         handlerModule.handler('arg', ctx, (err: Error, res: any) => {
           if (err) {
@@ -280,6 +285,7 @@ describe('lambda handler', () => {
         });
       });
 
+      const time2 = Date.now();
       const result2 = await new Promise((resolve, reject) => {
         handlerModule.handler('arg', ctx, (err: Error, res: any) => {
           if (err) {
@@ -298,18 +304,28 @@ describe('lambda handler', () => {
       assertSpanSuccess(span1);
       assert.strictEqual(span1.parentSpanId, undefined);
       assert.strictEqual(span1.attributes[SEMATTRS_FAAS_COLDSTART], true);
+      // Since it is coldstart, invocation span start time should be equal to lambda start time
+      assert.strictEqual(hrTimeToMilliseconds(span1.startTime), time0);
+      // Since it is coldstart, invocation span start time should be before actual handler call time
+      assert.ok(hrTimeToMilliseconds(span1.startTime) < time1);
 
       assert.strictEqual(result2, 'ok');
       assertSpanSuccess(span2);
       assert.strictEqual(span2.parentSpanId, undefined);
       assert.strictEqual(span2.attributes[SEMATTRS_FAAS_COLDSTART], false);
+      // Since it is warm invocation, invocation span start time should be after than lambda start time
+      assert.ok(hrTimeToMilliseconds(span2.startTime) > time0);
+      // Since it is warm invocation, invocation span start time should be equal or after than handler call time
+      assert.ok(hrTimeToMilliseconds(span2.startTime) >= time2);
     });
 
     it('should record coldstart with provisioned concurrency', async () => {
       process.env.AWS_LAMBDA_INITIALIZATION_TYPE = 'provisioned-concurrency';
 
-      initializeHandler('lambda-test/sync.handler');
+      const time0 = Date.now() - 1;
+      initializeHandler('lambda-test/sync.handler', { lambdaStartTime: time0 });
 
+      const time1 = Date.now();
       const result = await new Promise((resolve, reject) => {
         lambdaRequire('lambda-test/sync').handler(
           'arg',
@@ -323,20 +339,26 @@ describe('lambda handler', () => {
           }
         );
       });
+
       assert.strictEqual(result, 'ok');
       const spans = memoryExporter.getFinishedSpans();
       const [span] = spans;
+
       assert.strictEqual(spans.length, 1);
       assertSpanSuccess(span);
       assert.strictEqual(span.parentSpanId, undefined);
       assert.strictEqual(span.attributes[SEMATTRS_FAAS_COLDSTART], false);
+      // Since it is warm invocation, invocation span start time should be after than lambda start time
+      assert.ok(hrTimeToMilliseconds(span.startTime) > time0);
+      // Since it is warm invocation, invocation span start time should be equal or after than handler call time
+      assert.ok(hrTimeToMilliseconds(span.startTime) >= time1);
     });
 
     it('should record coldstart with proactive initialization', async () => {
-      initializeHandler('lambda-test/sync.handler', {
-        lambdaStartTime: Date.now() - 2 * lambdaMaxInitInMilliseconds,
-      });
+      const time0 = Date.now() - 2 * lambdaMaxInitInMilliseconds;
+      initializeHandler('lambda-test/sync.handler', { lambdaStartTime: time0 });
 
+      const time1 = Date.now();
       const result = await new Promise((resolve, reject) => {
         lambdaRequire('lambda-test/sync').handler(
           'arg',
@@ -350,13 +372,19 @@ describe('lambda handler', () => {
           }
         );
       });
+
       assert.strictEqual(result, 'ok');
       const spans = memoryExporter.getFinishedSpans();
       const [span] = spans;
+
       assert.strictEqual(spans.length, 1);
       assertSpanSuccess(span);
       assert.strictEqual(span.parentSpanId, undefined);
       assert.strictEqual(span.attributes[SEMATTRS_FAAS_COLDSTART], false);
+      // Since it is warm invocation, invocation span start time should be after than lambda start time
+      assert.ok(hrTimeToMilliseconds(span.startTime) > time0);
+      // Since it is warm invocation, invocation span start time should be equal or after than handler call time
+      assert.ok(hrTimeToMilliseconds(span.startTime) >= time1);
     });
 
     it('should record error', async () => {