Use ILambdaContext.RemainingTime to create a valid CancellationToken for AspNetCoreServer Requests

[brianfeucht]

Describe the feature

Currently when a Lambda execution times out, a AspNetCore request continues to execute even if it is checking a CancellationToken parameter for cancellation.

This feature would create a new CancellationTokenSource using ILambdaContext.RemainingTime and request cancellation when remaining time reaches or nears zero. This would allow consumer code to cancel long running operations in a way that is more controlled. This would create similar functionality to how timeouts are handled in Kestrel and IIS

Use Case

Due to upstream concurrency limitations, we sometimes have a lambda AspNetCore request timeout. Since the CancellationToken passed into the request is never cancelled, we continue to try to execute retry requests in the background if a lambda is reused.

Proposed Solution

No response

Other Information

No response

Acknowledgements

• [X] I may be able to implement this feature request
• [ ] This feature might incur a breaking change

AWS .NET SDK and/or Package version used

Amazon.Lambda.AspNetCoreServer.Hosting 1.5.0

Targeted .NET Platform

net6.0

Operating System and version

AmazonLinux (Lambda)

[ashishdhingra]

Needs review with the team.

[brianfeucht]

Sample implementation:

  builder.Services.Configure<MvcOptions>(options =>
  {
      var existingBinders = options.ModelBinderProviders
          .Where(x => x.GetType() == typeof(CancellationTokenModelBinderProvider)).ToArray();
      foreach (var binder in existingBinders)
      {
          options.ModelBinderProviders.Remove(binder);
      }

      options.ModelBinderProviders.Insert(0, new LambdaRemainingTimeCancellationTokenModelBinderProvider());
  });

  public class LambdaRemainingTimeCancellationTokenModelBinderProvider : IModelBinderProvider
  {
      public IModelBinder? GetBinder(ModelBinderProviderContext context)
      {
          ArgumentNullException.ThrowIfNull(context);

          if (context.Metadata.ModelType == typeof(CancellationToken))
          {
              return new LambdaRemainingTimeCancellationTokenModelBinder();
          }

          return null;
      }
  }
    
  public class LambdaRemainingTimeCancellationTokenModelBinder : IModelBinder
  {
      private const string HttpContextItemKey = "LambdaContextRemainingTimeCancellationTokenSource";

      private static readonly CancellationTokenModelBinder fallbackBinder = new();

      public Task BindModelAsync(ModelBindingContext bindingContext)
      {
          ArgumentNullException.ThrowIfNull(bindingContext);

          if (bindingContext.HttpContext.Items.ContainsKey(AbstractAspNetCoreFunction.LAMBDA_CONTEXT))
          {
              if (bindingContext.HttpContext.Items[AbstractAspNetCoreFunction.LAMBDA_CONTEXT] is ILambdaContext lambdaContext)
              {
                  var cancellationTokenSource = new CancellationTokenSource(lambdaContext.RemainingTime);

                  // Add the CTS to the Request Context, so GC doesn't clean it up
                  bindingContext.HttpContext.Items.Add(HttpContextItemKey, cancellationTokenSource);

                  bindingContext.ValidationState.Add(cancellationTokenSource.Token,
                      new ValidationStateEntry() { SuppressValidation = true });
                  bindingContext.Result = ModelBindingResult.Success(cancellationTokenSource.Token);

                  return Task.CompletedTask;
              }
          }

          // The Lambda Context is missing.  We must not be running in Lambda, so fallback to the default behavior.
          return fallbackBinder.BindModelAsync(bindingContext);
      }
    }

[ryanwilliams83]

Here is my implementation of this. 1 second buffer to shutdown gracefully. Defaults to 15 minutes when LambdaContext?.RemainingTime is null

WARNING - Might have a bug with GC, see message from brianfeucht

public static class ExtensionMethods
{
    public static CancellationTokenSource CreateCancellationTokenSource(this HttpContext instance, TimeSpan? timeoutBuffer = null)
    {
        var lambdaContext = instance.GetLambdaContext();

        return lambdaContext.CreateCancellationTokenSource(timeoutBuffer ?? TimeSpan.FromSeconds(1));
    }
    
    private static CancellationTokenSource CreateCancellationTokenSource(this ILambdaContext instance, TimeSpan buffer)
    {
        if (instance.RemainingTime == TimeSpan.MaxValue)
            throw new ArgumentOutOfRangeException(nameof(instance.RemainingTime));

        if (buffer.TotalSeconds < 0)
            throw new ArgumentOutOfRangeException(nameof(buffer));

        var remainingTime = instance.RemainingTime.Subtract(buffer);

        return remainingTime.TotalSeconds > 0
            ? new CancellationTokenSource(remainingTime)
            : throw new InvalidOperationException(@"Insufficient Time Remaining");
    }

    private static ILambdaContext GetLambdaContext(this HttpContext instance)
    {
        const string key = @"LambdaContext";

        if (!AWSXRayRecorder.IsLambda())
            return new LocalLambdaContext(TimeSpan.FromMinutes(15));

        instance.Items.TryGetValue(key, out var lambdaContext);

        return lambdaContext as ILambdaContext
            ?? throw new InvalidOperationException($@"{nameof(HttpContext)}.{nameof(HttpContext.Items)}[{key}] is not {nameof(ILambdaContext)}");
    }
}