Docs: explain how to cleanup with middy early return

[quantumew]

Expected Behaviour

Instructions on how to avoid this segment leak scenario.

Current Behaviour

When middy returns early, the segment is not closed (middy does not run after middleware), eventually this results in a RangeError via the incrementCounter call due to how many open segments exist.

Code snippet

const tracer = new Tracer({ serviceName: 'xxx' })

...

export const handler = middy(lambdaHandler).use(captureLambdaHandler(tracer)).use(cacheMiddleware);

Steps to Reproduce

1. Create a lambda with provisioned concurrency
2. Hook up the captureLambdaHandler middleware
3. Create middleware that runs after the capture mw, that returns early in its before
4. Put this lambda under high load (this issue keeps arising 40 mins into our load tests)

Possible Solution

Before returning early in middy, I grab the current segment and close it. That does not seem to solve the problem, surprisingly. Do I have to also close all potential sub segments? Is the handler segment not the issue, is it actually the parent lambda segment?

AWS Lambda Powertools for TypeScript version

latest

AWS Lambda function runtime

16.x

Packaging format used

npm

Execution logs

2023-04-10T14:16:56.326Z	d94677bc-a252-4830-a678-2c5ece7e0aff	ERROR	Invoke Error 	

{""errorType"":""RangeError"",""errorMessage"":""Maximum call stack size exceeded"",""stack"":[""RangeError: Maximum call stack size exceeded"",""    

at Ae.incrementCounter (/var/task/index.js:25:9545)"",""    

at Ae.incrementCounter (/var/task/index.js:25:9621)"",""    at Ae.incrementCounter (/var/task/index.js:25:9621)"",""

at Ae.incrementCounter (/var/task/index.js:25:9621)"",""    at Ae.incrementCounter (/var/task/index.js:25:9621)"","" 

at Ae.incrementCounter (/var/task/index.js:25:9621)"",""    at Ae.incrementCounter (/var/task/index.js:25:9621)"","" 

at Ae.incrementCounter (/var/task/index.js:25:9621)"",""    at Ae.incrementCounter (/var/task/index.js:25:9621)"",""

at Ae.incrementCounter (/var/task/index.js:25:9621)""]}

[dreamorosi]

Hi @quantumew, thank you for taking the time to open this issue.

Based on the info you have provided above I don't have enough info to reproduce the issue. Also, asking this to better understand the issue, how did you pinpoint the issue to either Middy or the Tracer?

The log that you have provided contains only a method called incrementCounter which is not a method present in our codebase, nor in Middy's (based on this search) so I'm not sure what I should be looking for.

When middy returns early, the segment is not closed (middy does not run after middleware)

Are you referring to this feature of Middy? Does your second custom middleware returns early like this?

Before returning early in middy, I grab the current segment and close it. That does not seem to solve the problem, surprisingly. Do I have to also close all potential sub segments? Is the handler segment not the issue, is it actually the parent lambda segment?

Let me break this down:

Before returning early in middy, I grab the current segment and close it. That does not seem to solve the problem, surprisingly.

That's a sensible solution, any chance you could share the bit of code that you use to do this operation? I'd like to see how you're closing the segment.

Do I have to also close all potential sub segments?

Not necessarily. Once you close the parent all its children should also be flushed.

Is the handler segment not the issue, is it actually the parent lambda segment?

I'm not 100% sure, but I doubt that. The function cannot mutate this segment and it's handled by the X-Ray SDK and refreshed at each invocation based on the invocation it/Trace ID present in the environment.

---

I'm not 100% sure I can pinpoint the problem with the current info, but what I can say is that based on the current implementation of the captureLambdaHandler middleware, both handler and root segments should be garbage collected at each invocation given that we are reassigning them at the beginning of each request.

Basically even if you didn't close any of the segment in your early return, the next function invocation should reassign the values of the segments and overwrite its previous values. In this case you might see trace data loss or incorrect segment duration but not stack overflow.

[quantumew]

Thank you for such a thorough response.

Here is how I attempt to close the segment:

if (tracer && tracer.isTracingEnabled()) {
  try {
    // Close lambda handler segment since we are circuit breaking here
    tracer.getSegment().close();
  } catch (err) {
    log.warn('Could not close the current segment before returning early', { err });
  }
}

The increment call is coming from the xray sdk.

Based on that stack trace, it looks like the problem is deeply nested segments, eventually causing a stack overflow.

I came to this conclusion purely through inference. The Tracer is created on lambda cold start. I assumed that when you called getSegment here, it returned the previous invocation's handler segment, since it was never closed, since we returned early. Resulting in handler segments nested in handler segments.

It just occurred to me that I am not doing this step. Could that be the problem? I am closing it but it is still the current segment?

[dreamorosi]

I assumed that when you called getSegment here, it returned the previous invocation's handler segment, since it was never closed, since we returned early.

That's a fair assumption and one that right off the bat, I'm also not able to rule out.

Now that I have a clearer idea of what your code is doing I can do more tests and give you a more precise answer.

It just occurred to me that I am not doing this step. Could that be the problem? I am closing it but it is still the current segment?

It could. The issue here is that you won't have access to that segment since it's not stored anywhere in the Tracer class.

I need to spend some time thinking about those two points before weighing in, but I see that there's a chance we might not be taking in account this edge case of users returning early in one of the downstream middleware.

I'll spend some time looking at this tomorrow and get back to you here.

Thanks for the patience and for reporting this interesting case!

[dreamorosi]

If the above is correct, I'm thinking this might be a good use case for using Middy's internal storage and allow downstream middleware to bail out early while also cleaning up traces, metrics, and other Powertools related actions.

[dreamorosi]

I've been talking with @willfarrell (maintainer of Middy) and another option that came up would be to place any middleware that you know might return early as first in the middleware chain:

const handler = middy()
  .use(myCacheMiddleware)
  .use(captureLambdaHandler(tracer))

This way if the first middleware returns, the captureLambdaHandler won't enter the before phase and potentially leaving a dangling segment.

If this is the problem, this should fix your issue. However I'll still continue investigating to see if we can do something for this edge case as I would prefer not having to rely on ordering.

[quantumew]

Definitely, great suggestion, it is one of our workarounds on one of our endpoints where disabling tracing was not an option. However, then we never get any trace data for our cache hits. My hope is to have the best of both worlds. Although, you could argue that a combination of ordering plus manual instrumentation would suffice here. Perhaps that is the tactical path I will pursue for now.