Lambda mode shares one CLS context and facade segment across concurrent invocations (Lambda Managed Instances)

[dreamorosi]

Summary

In Lambda mode, the SDK creates a single process-wide CLS context at initialization and never creates another (lib/env/aws_lambda.js):

var namespace = contextUtils.getNamespace();
namespace.enter(namespace.createContext());
contextUtils.setSegment(facadeSegment());

With Lambda Managed Instances (perExecutionEnvironmentMaxConcurrency > 1), multiple invocations are multiplexed into the same execution environment concurrently. Two pieces of state are then shared by all in-flight invocations:

1. The active-segment slot. setSegment() writes into the single shared context, so any code that manages subsegments via getSegment()/setSegment() (rather than captureAsyncFunc) resolves another invocation's subsegment once invocations interleave. Annotations, metadata, and instrumented-call subsegments attach to the wrong invocation - potentially another tenant's trace.

2. The facade segment itself. There is one facade per process, and resolveLambdaTraceData() mutates its trace_id/id/notTraced in place whenever getSegment() observes a new trace header. Under interleaving, subsegment flush() reads this.segment.trace_id at emission time, so even captureAsyncFunc users can emit subsegments onto whichever trace the facade was last resolved to.

The runtime already provides the primitive needed to fix this: globalThis.awslambda.InvokeStore, an AsyncLocalStorage that the runtime scopes around each invocation - and which resolveLambdaTraceData() already consults for the trace id.

Proposal

In Lambda mode, when awslambda.InvokeStore is available:

• scope the current-segment storage per invocation (e.g. key the segment slot by InvokeStore identity instead of the single CLS slot), and
• maintain a per-invocation facade (or at minimum snapshot the resolved trace data onto subsegments at creation time rather than reading the shared facade at flush time),

falling back to today's single-context behavior on runtimes without InvokeStore. This would make both captureAsyncFunc and getSegment()/setSegment()-based instrumentation safe under LMI concurrency without API changes.

Reproduction

Deploy any function using getSegment()/addNewSubsegment()/setSegment() per invocation on a capacity provider with perExecutionEnvironmentMaxConcurrency > 1, add a distinctive annotation per invocation, and force multiplexing: annotations land on other invocations' subsegments, and subsegment documents are emitted with the wrong trace_id.

Downstream report with full trace evidence: aws-powertools/powertools-lambda-typescript#5434.

[svozza]

This is a very severe concurrency bug, we are losing or corrupting customers telemetry data in Lambda Managed Instance environments. In order to aid a quick resolution, I've put together a minimal, deployable reproduction using only aws-xray-sdk-core — no Powertools or other instrumentation layers: https://github.com/svozza/xray-lmi-shared-context-repro

It deploys a single function on a minimum-size capacity provider (CDK app included), fires 30 simultaneous invocations to force multiplexing, and asserts attribution from in-process ground truth: each invocation records which subsegment it opened (by object identity, at creation time), and the checker verifies every annotation sits on the subsegment opened by the invocation that wrote it. Overlap is deterministic — each invocation blocks inside the handler until a peer invocation is observed in flight in the same execution environment, then annotates inside the clobber window — so the result doesn't depend on timing luck, and nothing depends on trace delivery.

Three commands, three outcomes, verified on real Lambda Managed Instances (eu-west-1, Node.js 22, perExecutionEnvironmentMaxConcurrency: 10):

• npm run repro — the getSegment()/addNewSubsegment()/setSegment() lifecycle. Fails every run with both symptoms of the shared slot, which pair up in the output: invocation A's annotation lands on B's subsegment (cross-attribution), and B's own annotation then resolves to the restored facade — whose mutators are no-op stubs in Lambda mode — and is silently dropped, with no warning:

  FAIL: annotations misattributed across concurrent invocations
    - cross-attribution: annotation written by inv-0 landed on the subsegment opened by inv-26
    - cross-attribution: annotation written by inv-10 landed on the subsegment opened by inv-15
    ...
    - loss: annotations from inv-15, inv-26, ... never landed on any subsegment
  

• npm run repro:gap — same lifecycle with one event-loop yield between annotate and close, matching how integrations with phased teardown (middleware after-hooks, finally behind an await) behave. Fails with pure cross-attribution — one invocation's data written into another invocation's trace with no loss and no signal at all.

• npm run control — identical barrier/annotation path wrapped in captureAsyncFunc. Passes: all 30 invocations attribute correctly, confirming the isolation mechanism exists in the SDK and the gap is that it isn't applied per-invocation in Lambda mode (and can't be by integrations that don't own the whole invocation callback).

This exercises the active-segment-slot half of this issue (point 1). It doesn't cover the shared-facade trace_id mutation (point 2), since it asserts on live subsegment objects rather than emitted documents — though the harness (barrier + prototype-patch ground truth) should extend to that by capturing flush() output if useful.

[svozza]

Following up on the deployable repro above: here are failing tests written against this repo's own test suite — mocha/chai/sinon, same conventions as test/integration/, run from dist/ per the package's test script. No LMI infrastructure required: both mechanisms reproduce with plain interleaved async in a single process, which is exactly why LMI surfaces them in production (the runtime just makes that interleaving happen for real invocations). Verified failing on current master (458fc17).

Note: the full test files are inside the collapsed ▸ lambda_…test.js sections below — click to expand.

Both files assert the desired isolation invariant, so they fail today and pass once fixed. Drop them into packages/core/test/integration/ and run each standalone (they set Lambda-mode env at module load, so they shouldn't be collected into a shared suite run):

cd packages/core && npm run compile
npx mocha ./dist/test/integration/lambda_context_isolated_across_concurrent_invocations.test.js -R spec
npx mocha ./dist/test/integration/lambda_facade_trace_data_isolated_across_concurrent_invocations.test.js -R spec

Bug 1 — shared CLS slot

Two interleaved "invocations" each do the documented manual-subsegment pattern (getSegment → addNewSubsegment → setSegment), await, then check they still see their own subsegment. Also asserts parenting: because B's initial getSegment() returns whatever A last setSegment()-ed, B's subsegment gets nested under A's subsegment rather than under the facade — the corruption starts before setSegment is even called.

lambda_context_isolated_across_concurrent_invocations.test.js

process.env.LAMBDA_TASK_ROOT = process.env.LAMBDA_TASK_ROOT || '/var/task';
process.env._X_AMZN_TRACE_ID = 'Root=1-5f000001-000000000000000000000001;Parent=0000000000000001;Sampled=1';

// Requires must come after the env setup above: lib/aws-xray.js initializes
// Lambda mode at module load based on LAMBDA_TASK_ROOT (entering the shared
// CLS context, creating the facade segment from _X_AMZN_TRACE_ID), and none
// of that is undoable in-process. Run this file standalone.
var assert = require('chai').assert;
var sinon = require('sinon');

var AWSXRay = require('../../lib');
var SegmentEmitter = require('../../lib/segment_emitter');

describe('lambda CLS context isolated across concurrent invocations', function() {
  var sendStub, facade;

  before(function() {
    sendStub = sinon.stub(SegmentEmitter, 'send');
    facade = AWSXRay.getSegment();
    assert.isTrue(facade.facade, 'expected the Lambda facade segment on the context');
  });

  after(function() {
    sendStub.restore();
  });

  afterEach(function() {
    AWSXRay.setSegment(facade);
  });

  // The documented Lambda pattern for custom subsegments:
  //   var segment = AWSXRay.getSegment();
  //   var subsegment = segment.addNewSubsegment(name);
  //   AWSXRay.setSegment(subsegment);
  //   ... async work, annotations, downstream calls ...
  //   subsegment.close(); AWSXRay.setSegment(segment);
  function invocation(name, delayMs) {
    var segment = AWSXRay.getSegment();
    var subsegment = segment.addNewSubsegment(name);
    AWSXRay.setSegment(subsegment);

    return new Promise(function(resolve) {
      setTimeout(resolve, delayMs);
    }).then(function() {
      return { created: subsegment, seen: AWSXRay.getSegment() };
    });
  }

  it('each concurrent invocation still sees its own subsegment after awaiting', function() {
    return Promise.all([
      invocation('invocation-A', 20),
      invocation('invocation-B', 10)
    ]).then(function(results) {
      var a = results[0], b = results[1];

      assert.equal(a.seen.id, a.created.id,
        'invocation A must get its own subsegment back from getSegment(), but saw "' + a.seen.name + '"');
      assert.equal(b.seen.id, b.created.id,
        'invocation B must get its own subsegment back from getSegment(), but saw "' + b.seen.name + '"');
    });
  });

  it('each concurrent invocation parents its subsegment on the facade, not on another invocation\'s subsegment', function() {
    return Promise.all([
      invocation('invocation-C', 20),
      invocation('invocation-D', 10)
    ]).then(function(results) {
      var c = results[0], d = results[1];

      assert.equal(c.created.parent.id, facade.id,
        'invocation C\'s subsegment must be parented on the facade, but its parent is "' + c.created.parent.name + '"');
      assert.equal(d.created.parent.id, facade.id,
        'invocation D\'s subsegment must be parented on the facade, but its parent is "' + d.created.parent.name + '"');
    });
  });
});

Failure output today:

1) each concurrent invocation still sees its own subsegment after awaiting:
   invocation A must get its own subsegment back from getSegment(), but saw "invocation-B"
2) each concurrent invocation parents its subsegment on the facade, not on another
   invocation's subsegment:
   invocation D's subsegment must be parented on the facade, but its parent is "invocation-C"

Bug 2 — facade trace data resolved at flush time

Invocation A (Sampled=1) opens a subsegment; invocation B's arrival is simulated by changing the current trace header and calling getSegment() (which triggers resolveLambdaTraceData()'s in-place mutation of the shared facade); A closes. SegmentEmitter.send is stubbed to capture what's emitted. Two header sources are exercised:

• _X_AMZN_TRACE_ID env var (header-source-agnostic proof of mechanism), and
• an InvokeStore variant: globalThis.awslambda.InvokeStore.getXRayTraceId() backed by a real AsyncLocalStorage keyed per invocation — faithful to InvokeStoreMulti in @aws/lambda-invoke-store, i.e. the actual LMI production path. (AWS_LAMBDA_INITIALIZATION_TYPE/AWS_LAMBDA_MAX_CONCURRENCY are set for environmental faithfulness; note the SDK's only LMI awareness is the presence of the global, in resolveLambdaTraceData().)

lambda_facade_trace_data_isolated_across_concurrent_invocations.test.js

var TRACE_A_SAMPLED = 'Root=1-5f0000aa-0000000000000000000000aa;Parent=00000000000000aa;Sampled=1';
var TRACE_B_UNSAMPLED = 'Root=1-5f0000bb-0000000000000000000000bb;Parent=00000000000000bb;Sampled=0';
var TRACE_B_SAMPLED = 'Root=1-5f0000cc-0000000000000000000000cc;Parent=00000000000000cc;Sampled=1';
var ROOT_A = '1-5f0000aa-0000000000000000000000aa';

process.env.LAMBDA_TASK_ROOT = process.env.LAMBDA_TASK_ROOT || '/var/task';
process.env._X_AMZN_TRACE_ID = TRACE_A_SAMPLED;

// Requires must come after the env setup above: lib/aws-xray.js initializes
// Lambda mode at module load based on LAMBDA_TASK_ROOT (entering the shared
// CLS context, creating the facade segment from _X_AMZN_TRACE_ID), and none
// of that is undoable in-process. Run this file standalone.
var assert = require('chai').assert;
var sinon = require('sinon');
var AsyncLocalStorage = require('async_hooks').AsyncLocalStorage;

var AWSXRay = require('../../lib');
var SegmentEmitter = require('../../lib/segment_emitter');

describe('lambda facade trace data isolated across concurrent invocations', function() {
  var sandbox, sendStub;

  beforeEach(function() {
    sandbox = sinon.createSandbox();
    sendStub = sandbox.stub(SegmentEmitter, 'send');
  });

  afterEach(function() {
    sandbox.restore();
    delete globalThis.awslambda;
    delete process.env.AWS_LAMBDA_INITIALIZATION_TYPE;
    delete process.env.AWS_LAMBDA_MAX_CONCURRENCY;
    process.env._X_AMZN_TRACE_ID = TRACE_A_SAMPLED;
  });

  describe('via _X_AMZN_TRACE_ID changing between open and close', function() {
    it('closed subsegment is emitted with the trace_id of the invocation that opened it', function() {
      process.env._X_AMZN_TRACE_ID = TRACE_A_SAMPLED;
      var facade = AWSXRay.getSegment();
      assert.isTrue(facade.facade, 'expected the Lambda facade segment on the context');
      var subsegment = facade.addNewSubsegment('invocation-A-work');

      // Invocation B arrives concurrently with a different (sampled) trace
      // header and touches the SDK.
      process.env._X_AMZN_TRACE_ID = TRACE_B_SAMPLED;
      AWSXRay.getSegment();

      // Invocation A finishes its work.
      subsegment.close();

      assert.equal(sendStub.callCount, 1, 'subsegment must be sent to the emitter');
      var emitted = JSON.parse(sendStub.firstCall.args[0].format());
      assert.equal(emitted.trace_id, ROOT_A,
        'subsegment opened under trace A must be emitted with trace A\'s trace_id, got ' + emitted.trace_id);
    });

    it('closed subsegment is still sent when a not-sampled invocation arrived in between', function() {
      process.env._X_AMZN_TRACE_ID = TRACE_A_SAMPLED;
      var facade = AWSXRay.getSegment();
      var subsegment = facade.addNewSubsegment('invocation-A-work');

      // Invocation B arrives concurrently with Sampled=0.
      process.env._X_AMZN_TRACE_ID = TRACE_B_UNSAMPLED;
      AWSXRay.getSegment();

      subsegment.close();

      assert.equal(sendStub.callCount, 1,
        'subsegment opened under a sampled trace must be sent even though another, ' +
        'not-sampled invocation resolved trace data before it closed');
    });
  });

  describe('via InvokeStore (LMI): per-invocation getXRayTraceId()', function() {
    var als;

    beforeEach(function() {
      // Set for environmental faithfulness only: LMI sets these on the
      // runtime, and @aws/lambda-invoke-store's getInstanceAsync() keys the
      // multi-concurrency (AsyncLocalStorage) implementation off
      // AWS_LAMBDA_MAX_CONCURRENCY. The SDK itself never reads them - its
      // only LMI awareness is the presence of globalThis.awslambda.InvokeStore,
      // which we install directly below.
      process.env.AWS_LAMBDA_INITIALIZATION_TYPE = 'lambda-managed-instances';
      process.env.AWS_LAMBDA_MAX_CONCURRENCY = '2';

      // Faithful stand-in for the LMI runtime's @aws/lambda-invoke-store
      // multi-concurrency store: getXRayTraceId() is backed by an
      // AsyncLocalStorage keyed per invocation.
      als = new AsyncLocalStorage();
      globalThis.awslambda = {
        InvokeStore: {
          getXRayTraceId: function() {
            var store = als.getStore();
            return store && store.xrayTraceId;
          }
        }
      };
    });

    function runAsInvocation(traceHeader, fn) {
      return als.run({ xrayTraceId: traceHeader }, fn);
    }

    it('closed subsegment carries the opening invocation\'s trace_id and is sent', function() {
      var subsegment;

      // Invocation A (sampled) opens a subsegment.
      runAsInvocation(TRACE_A_SAMPLED, function() {
        var facade = AWSXRay.getSegment();
        assert.isTrue(facade.facade, 'expected the Lambda facade segment on the context');
        subsegment = facade.addNewSubsegment('invocation-A-work');
      });

      // Invocation B (not sampled) runs concurrently and touches the SDK.
      runAsInvocation(TRACE_B_UNSAMPLED, function() {
        AWSXRay.getSegment();
      });

      // Invocation A resumes and closes its subsegment.
      runAsInvocation(TRACE_A_SAMPLED, function() {
        subsegment.close();
      });

      assert.equal(sendStub.callCount, 1,
        'subsegment opened under sampled invocation A must be sent to the emitter');
      var emitted = JSON.parse(sendStub.firstCall.args[0].format());
      assert.equal(emitted.trace_id, ROOT_A,
        'subsegment opened under trace A must be emitted with trace A\'s trace_id, got ' + emitted.trace_id);
    });
  });
});

Failure output today:

1) closed subsegment is emitted with the trace_id of the invocation that opened it:
   subsegment opened under trace A must be emitted with trace A's trace_id,
   got 1-5f0000cc-0000000000000000000000cc        <- B's trace
2) closed subsegment is still sent when a not-sampled invocation arrived in between:
   expected callCount 1, got 0                    <- silently dropped (B was Sampled=0)
3) [InvokeStore variant] subsegment opened under sampled invocation A must be sent:
   expected callCount 1, got 0

Failures 2/3 are the data-loss mode: one Sampled=0 invocation anywhere in flight flips the singleton facade's notTraced, and every other invocation's subsegments closing in that window are dropped by Subsegment#flush(). Failure 1 is the cross-trace pollution mode via Subsegment#format() reading this.segment.trace_id at emission time. Note this bug hits captureAsyncFunc/captureFunc and all captured clients too — their subsegments hang off the same shared facade — so its blast radius is wider than the CLS bug.