feat: Enforce 20% sampling on successful invocations in prod environment

Author: medikoo

node/packages/aws-lambda-sdk/instrument/index.js

@@ -18,6 +18,13 @@ const serverlessSdk = require('./lib/sdk');
 const objHasOwnProperty = Object.prototype.hasOwnProperty;
 const unresolvedPromise = new Promise(() => {});
 
+const coreTraceSpanNames = new Set([
+  'aws.lambda',
+  'aws.lambda.initialization',
+  'aws.lambda.invocation',
+]);
+const alertEventNames = new Set(['telemetry.error.generated.v1', 'telemetry.warning.generated.v1']);
+
 const capturedEvents = [];
 serverlessSdk._eventEmitter.on('captured-event', (capturedEvent) =>
   capturedEvents.push(capturedEvent)
@@ -91,25 +98,38 @@ const reportResponse = async (response, context, endTime) => {
   await sendTelemetry('request-response', payloadBuffer);
 };
 
-const reportTrace = () => {
+const reportTrace = ({ isErrorOutcome }) => {
+  // Sample out 80% of traces in production mode if no warning or error event was reported
+  const isSampledOut =
+    (!isErrorOutcome &&
+      !serverlessSdk._isDebugMode &&
+      !serverlessSdk._isDevMode &&
+      !capturedEvents.some(({ name }) => alertEventNames.has(name)) &&
+      Math.random() > 0.2) ||
+    undefined;
   const payload = (serverlessSdk._lastTrace = {
+    isSampledOut,
     slsTags: {
       orgId: serverlessSdk.orgId,
       service: process.env.AWS_LAMBDA_FUNCTION_NAME,
       sdk: { name: pkgJson.name, version: pkgJson.version },
     },
-    spans: Array.from(awsLambdaSpan.spans).map((span) => {
+    spans: Array.from(awsLambdaSpan.spans, (span) => {
+      if (isSampledOut && !coreTraceSpanNames.has(span.name)) return null;
       const spanPayload = span.toProtobufObject();
       delete spanPayload.input;
       delete spanPayload.output;
       return spanPayload;
-    }),
-    events: capturedEvents
-      .filter(filterCapturedEvent)
-      .map((capturedEvent) => capturedEvent.toProtobufObject()),
-    customTags: objHasOwnProperty.call(serverlessSdk, '_customTags')
-      ? JSON.stringify(serverlessSdk._customTags)
-      : undefined,
+    }).filter(Boolean),
+    events: isSampledOut
+      ? []
+      : capturedEvents
+          .filter(filterCapturedEvent)
+          .map((capturedEvent) => capturedEvent.toProtobufObject()),
+    customTags:
+      !isSampledOut && objHasOwnProperty.call(serverlessSdk, '_customTags')
+        ? JSON.stringify(serverlessSdk._customTags)
+        : undefined,
   });
   const payloadBuffer = (serverlessSdk._lastTraceBuffer =
     traceProto.TracePayload.encode(payload).finish());
@@ -172,7 +192,7 @@ const closeTrace = async (outcome, outcomeResult) => {
     awsLambdaSpan.close({ endTime });
     // Root span comes with "aws.lambda.*" tags, which require unconditionally requestId
     // which we don't have if handler crashed at initialization
-    if (invocationContextAccessor.value) reportTrace();
+    if (invocationContextAccessor.value) reportTrace({ isErrorOutcome });
     flushSpans();
     clearRootSpan();
 

node/packages/aws-lambda-sdk/package.json

@@ -5,7 +5,7 @@
   "author": "Serverless, Inc.",
   "dependencies": {
     "@serverless/sdk": "^0.5.2",
-    "@serverless/sdk-schema": "^0.15.1",
+    "@serverless/sdk-schema": "^0.15.2",
     "d": "^1.0.1",
     "ext": "^1.7.0",
     "long": "^5.2.1",

node/packages/aws-lambda-sdk/test/integration/index.test.js

@@ -552,6 +552,20 @@ describe('integration', function () {
           ['v14', { configuration: { Runtime: 'nodejs14.x' } }],
           ['v16', { configuration: { Runtime: 'nodejs16.x' } }],
           ['v18', { configuration: { Runtime: 'nodejs18.x' } }],
+          [
+            'sampled',
+            {
+              configuration: {
+                Environment: {
+                  Variables: {
+                    SLS_ORG_ID: process.env.SLS_ORG_ID,
+                    SLS_CRASH_ON_SDK_ERROR: '1',
+                    AWS_LAMBDA_EXEC_WRAPPER: '/opt/sls-sdk-node/exec-wrapper.sh',
+                  },
+                },
+              },
+            },
+          ],
           [
             'sqs',
             {

node/packages/aws-lambda-sdk/test/unit/internal-extension/index.test.js

@@ -140,17 +140,32 @@ describe('internal-extension/index.test.js', () => {
     process.env.SLS_ORG_ID = 'dummy';
     process.env.SLS_UNIT_TEST_RUN = '1';
     process.env.SLS_CRASH_ON_SDK_ERROR = '1';
+    process.env.SLS_SDK_DEBUG = '1';
   });
   afterEach(() => {
     delete process.env._HANDLER;
     delete process.env.AWS_LAMBDA_FUNCTION_NAME;
   });
 
+  describe('sampling', () => {
+    before(() => {
+      delete process.env.SLS_SDK_DEBUG;
+    });
+    after(() => {
+      process.env.SLS_SDK_DEBUG = '1';
+    });
+    it('should produce complete basic trace when sampled', async () => {
+      // There's 80% chance that trace will be sampled
+      return handleInvocation('callback');
+    });
+  });
+
   it('should handle "ESM callback"', async () => handleInvocation('esm-callback/index'));
   it('should handle "ESM thenable"', async () => handleInvocation('esm-thenable/index'));
   it('should handle "ESM nested module"', async () =>
     handleInvocation('esm-nested/nested/within/index'));
   it('should handle "callback"', async () => handleInvocation('callback'));
+
   it('should handle "thenable"', async () => handleInvocation('thenable'));
   it('should handle "esbuild from ESM callback', async () =>
     handleInvocation('esbuild-from-esm-callback'));