Skip to content

Commit 4721abf

Browse files
talenfishertalenfisher
authored
LAMBJ-19 Task Cancellation (#96)
1 parent 11c57fc commit 4721abf

File tree

26 files changed

+82
-39
lines changed

26 files changed

+82
-39
lines changed

‎.github/releases/v0.6.0-beta4.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
This release introduces the following:
22

3-
-
3+
- **Breaking Change**: Lambdas now take a CancellationToken as a second argument and will gracefully stop execution when cancellation is requested.

‎examples/AwsClientFactories/Handler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using System;
2+
using System.Threading;
23
using System.Threading.Tasks;
34

45
using Amazon.Lambda.Core;
@@ -24,7 +25,7 @@ public Handler(IAwsFactory<IAmazonS3> s3Factory)
2425
this.s3Factory = s3Factory;
2526
}
2627

27-
public async Task<string> Handle(Request request)
28+
public async Task<string> Handle(Request request, CancellationToken cancellationToken = default)
2829
{
2930
s3Client = await s3Factory.Create(request.RoleArn);
3031

‎examples/CustomConfiguration/Handler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
using Amazon.Lambda.Core;
@@ -18,7 +19,7 @@ public Handler(IOptions<Config> config)
1819
this.config = config.Value;
1920
}
2021

21-
public Task<string> Handle(object request)
22+
public Task<string> Handle(object request, CancellationToken cancellationToken = default)
2223
{
2324
return Task.FromResult(config.Foo);
2425
}

‎examples/CustomRuntime/Handler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
using Amazon.Lambda.Core;
@@ -9,7 +10,7 @@ namespace Lambdajection.Examples.CustomRuntime
910
[Lambda(typeof(Startup))]
1011
public partial class Handler
1112
{
12-
public async Task<string> Handle(object request)
13+
public async Task<string> Handle(object request, CancellationToken cancellationToken = default)
1314
{
1415
return await Task.FromResult("Hello World!");
1516
}

‎examples/CustomSerializer/Handler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
using System.Collections.Generic;
33
using System.Net;
44
using System.Text.RegularExpressions;
5+
using System.Threading;
56
using System.Threading.Tasks;
67

78
using Amazon.Lambda.ApplicationLoadBalancerEvents;
@@ -24,7 +25,7 @@ public Handler(EmbeddedResourceReader reader)
2425
this.reader = reader;
2526
}
2627

27-
public async Task<ApplicationLoadBalancerResponse> Handle(ApplicationLoadBalancerRequest request)
28+
public async Task<ApplicationLoadBalancerResponse> Handle(ApplicationLoadBalancerRequest request, CancellationToken cancellationToken = default)
2829
{
2930
var path = Regex.Replace(request.Path, @"^\/", "");
3031
var contents = await reader.ReadAsString(path);

‎examples/EncryptedOptions/Handler.cs

+3-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
using System.Threading.Tasks;
1+
using System.Threading;
2+
using System.Threading.Tasks;
23

34
using Amazon.Lambda.Core;
45

@@ -18,7 +19,7 @@ public Handler(IOptions<Options> options)
1819
this.options = options.Value;
1920
}
2021

21-
public Task<string> Handle(object request)
22+
public Task<string> Handle(object request, CancellationToken cancellationToken = default)
2223
{
2324
return Task.FromResult(options.EncryptedValue); // despite the name, this value will have already been decrypted for you.
2425
}

‎src/Core/ILambda.cs

+3-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
namespace Lambdajection.Core
@@ -13,7 +14,8 @@ public interface ILambda<TLambdaParameter, TLambdaOutput>
1314
/// The lambda's entrypoint.
1415
/// </summary>
1516
/// <param name="parameter">The lambda's input parameter.</param>
17+
/// <param name="cancellationToken">Token used to cancel the operation.</param>
1618
/// <returns>The lambda's return value.</returns>
17-
Task<TLambdaOutput> Handle(TLambdaParameter parameter);
19+
Task<TLambdaOutput> Handle(TLambdaParameter parameter, CancellationToken cancellationToken = default);
1820
}
1921
}

‎src/Core/ILambdaInitializationService.cs

+3-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
namespace Lambdajection.Core
@@ -10,7 +11,8 @@ public interface ILambdaInitializationService
1011
/// <summary>
1112
/// Runs the initialization service's instructions.
1213
/// </summary>
14+
/// <param name="cancellationToken">Token used to cancel operation.</param>
1315
/// <returns>The initialization task.</returns>
14-
Task Initialize();
16+
Task Initialize(CancellationToken cancellationToken = default);
1517
}
1618
}

‎src/Core/LambdaHost.cs

+9-5
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using System;
22
using System.Linq;
3+
using System.Threading;
34
using System.Threading.Tasks;
45

56
using Amazon.Lambda.Core;
@@ -76,22 +77,24 @@ internal LambdaHost(Action<LambdaHost<TLambda, TLambdaParameter, TLambdaOutput,
7677
/// </summary>
7778
/// <param name="parameter">The input parameter to pass to the lambda.</param>
7879
/// <param name="context">The context object to pass to the lambda.</param>
80+
/// <param name="cancellationToken">Token used to cancel the operation.</param>
7981
/// <returns>The return value of the lambda.</returns>
80-
public async Task<TLambdaOutput> Run(TLambdaParameter parameter, ILambdaContext context)
82+
public async Task<TLambdaOutput> Run(TLambdaParameter parameter, ILambdaContext context, CancellationToken cancellationToken = default)
8183
{
8284
if (RunInitializationServices)
8385
{
84-
await Initialize();
86+
await Initialize(cancellationToken);
8587
}
8688

89+
cancellationToken.ThrowIfCancellationRequested();
8790
scope = ServiceProvider.CreateScope();
8891

8992
var provider = scope.ServiceProvider;
9093
var scopeContext = provider.GetRequiredService<LambdaScope>();
9194
scopeContext.LambdaContext = context;
9295

9396
lambda = provider.GetRequiredService<TLambda>();
94-
return await lambda.Handle(parameter);
97+
return await lambda.Handle(parameter, cancellationToken);
9598
}
9699

97100
/// <summary>
@@ -120,11 +123,12 @@ private static async Task MaybeDispose(object? obj)
120123
(obj as IDisposable)?.Dispose();
121124
}
122125

123-
private async Task Initialize()
126+
private async Task Initialize(CancellationToken cancellationToken = default)
124127
{
128+
cancellationToken.ThrowIfCancellationRequested();
125129
var services = ServiceProvider.GetServices<ILambdaInitializationService>();
126130

127-
var initializeTasks = services.Select(service => service.Initialize());
131+
var initializeTasks = services.Select(service => service.Initialize(cancellationToken));
128132
await Task.WhenAll(initializeTasks);
129133

130134
var disposeTasks = services.Select(MaybeDispose);

‎src/Encryption/DefaultDecryptionService.cs

+10-8
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using System;
22
using System.IO;
3+
using System.Threading;
34
using System.Threading.Tasks;
45

56
using Amazon.KeyManagementService;
@@ -23,19 +24,20 @@ public DefaultDecryptionService(IAmazonKeyManagementService kmsClient)
2324
this.kmsClient = kmsClient;
2425
}
2526

26-
/// <summary>
27-
/// Decrypts a value and returns it as a plaintext string.
28-
/// </summary>
29-
/// <param name="ciphertext">The ciphertext to decrypt.</param>
30-
/// <returns>The plaintext decrypted value.</returns>
31-
public virtual async Task<string> Decrypt(string ciphertext)
27+
/// <inheritdoc />
28+
public virtual async Task<string> Decrypt(string ciphertext, CancellationToken cancellationToken = default)
3229
{
30+
cancellationToken.ThrowIfCancellationRequested();
31+
3332
using var stream = new MemoryStream();
3433
var byteArray = Convert.FromBase64String(ciphertext);
35-
await stream.WriteAsync(byteArray);
34+
35+
await stream.WriteAsync(byteArray, cancellationToken);
36+
cancellationToken.ThrowIfCancellationRequested();
3637

3738
var request = new DecryptRequest { CiphertextBlob = stream };
38-
var response = await kmsClient.DecryptAsync(request);
39+
var response = await kmsClient.DecryptAsync(request, cancellationToken);
40+
cancellationToken.ThrowIfCancellationRequested();
3941

4042
using var reader = new StreamReader(response.Plaintext);
4143
return await reader.ReadToEndAsync();

‎src/Encryption/IDecryptionService.cs

+3-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
namespace Lambdajection.Encryption
@@ -11,7 +12,8 @@ public interface IDecryptionService
1112
/// Decrypts a value.
1213
/// </summary>
1314
/// <param name="ciphertext">The value to be decrypted.</param>
15+
/// <param name="cancellationToken">Token to cancel the operation with.</param>
1416
/// <returns>The decrypted value.</returns>
15-
Task<string> Decrypt(string ciphertext);
17+
Task<string> Decrypt(string ciphertext, CancellationToken cancellationToken = default);
1618
}
1719
}

‎src/Generator/LambdaCompilationScanner.cs

+2-2
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@ public class LambdaCompilationScanner
1919
private readonly Compilation compilation;
2020
private readonly string lambdaTypeName;
2121
private readonly string startupDisplayName;
22-
private readonly HashSet<OptionClass> optionClasses = new HashSet<OptionClass>();
23-
private readonly HashSet<AwsServiceMetadata> awsServices = new HashSet<AwsServiceMetadata>();
22+
private readonly HashSet<OptionClass> optionClasses = new();
23+
private readonly HashSet<AwsServiceMetadata> awsServices = new();
2424
private bool includeDecryptionFacade;
2525

2626
public LambdaCompilationScanner(Compilation compilation, IEnumerable<SyntaxTree> syntaxTrees, string lambdaTypeName, string startupDisplayName)

‎src/Generator/OptionsDecryptorGenerator.cs

+16-2
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ static IEnumerable<ExpressionSyntax> GenerateDecryptPropertyCalls(IEnumerable<st
7070
{
7171
foreach (var prop in properties)
7272
{
73-
yield return ParseExpression($"Decrypt{prop}()");
73+
yield return ParseExpression($"Decrypt{prop}(cancellationToken)");
7474
}
7575
}
7676

@@ -89,8 +89,15 @@ IEnumerable<StatementSyntax> GenerateBody()
8989
yield return ExpressionStatement(awaitExpression);
9090
}
9191

92+
var defaultValue = EqualsValueClause(ParseToken("="), ParseExpression("default"));
93+
var parameters = SeparatedList(new ParameterSyntax[]
94+
{
95+
Parameter(List<AttributeListSyntax>(), TokenList(), ParseTypeName("CancellationToken"), ParseToken("cancellationToken"), defaultValue),
96+
});
97+
9298
return MethodDeclaration(ParseTypeName("Task"), "Initialize")
9399
.WithModifiers(TokenList(Token(PublicKeyword), Token(AsyncKeyword)))
100+
.WithParameterList(ParameterList(parameters))
94101
.WithBody(Block(GenerateBody()));
95102
}
96103

@@ -111,11 +118,18 @@ public static MemberDeclarationSyntax GenerateDecryptPropertyMethod(string prop)
111118
{
112119
IEnumerable<StatementSyntax> GenerateBody()
113120
{
114-
yield return ParseStatement($"options.{prop} = await decryptionService.Decrypt(options.{prop});");
121+
yield return ParseStatement($"options.{prop} = await decryptionService.Decrypt(options.{prop}, cancellationToken);");
115122
}
116123

124+
var defaultValue = EqualsValueClause(ParseToken("="), ParseExpression("default"));
125+
var parameters = SeparatedList(new ParameterSyntax[]
126+
{
127+
Parameter(List<AttributeListSyntax>(), TokenList(), ParseTypeName("CancellationToken"), ParseToken("cancellationToken"), defaultValue),
128+
});
129+
117130
return MethodDeclaration(ParseTypeName("Task"), $"Decrypt{prop}")
118131
.WithModifiers(TokenList(Token(PrivateKeyword), Token(AsyncKeyword)))
132+
.WithParameterList(ParameterList(parameters))
119133
.WithBody(Block(GenerateBody()));
120134
}
121135
}

‎src/Generator/UnitGenerator.cs

+1-1
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ namespace Lambdajection.Generator
2222
[Generator]
2323
public class UnitGenerator : ISourceGenerator
2424
{
25-
private readonly UsingsGenerator usingsGenerator = new UsingsGenerator();
25+
private readonly UsingsGenerator usingsGenerator = new();
2626

2727
public void Initialize(GeneratorInitializationContext context)
2828
{

‎src/Generator/UsingsGenerator.cs

+1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ public class UsingsGenerator
1313
private static readonly string[] DefaultUsings = new string[]
1414
{
1515
"System",
16+
"System.Threading",
1617
"System.Threading.Tasks",
1718
"System.IO",
1819
"Microsoft.Extensions.DependencyInjection",

‎templates/Project/Handler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
using Amazon.Lambda.Core;
@@ -17,7 +18,7 @@ public partial class Handler
1718
private bool disposed;
1819

1920
#endif
20-
public async Task<object> Handle(object request)
21+
public async Task<object> Handle(object request, CancellationToken cancellationToken = default)
2122
{
2223
return await Task.FromResult(new { });
2324
}

‎tests/Common/TestLambda.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
using Lambdajection.Core;
@@ -8,7 +9,7 @@ public class TestLambda : ILambda<object, object>
89
{
910
public object Request { get; set; } = null!;
1011

11-
public virtual Task<object> Handle(object request)
12+
public virtual Task<object> Handle(object request, CancellationToken cancellationToken = default)
1213
{
1314
Request = request;
1415

‎tests/Compilation/Projects/AmazonFactories/Handler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
using Amazon.Lambda.Core;
@@ -22,7 +23,7 @@ public Handler(S3Utility utility)
2223
this.utility = utility;
2324
}
2425

25-
public Task<IAwsFactory<IAmazonS3>> Handle(string request)
26+
public Task<IAwsFactory<IAmazonS3>> Handle(string request, CancellationToken cancellationToken = default)
2627
{
2728
return Task.FromResult(utility.Factory);
2829
}

‎tests/Compilation/Projects/ConfigFactory/Handler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
using Amazon.Lambda.Core;
@@ -20,7 +21,7 @@ public Handler(IConfiguration config)
2021
this.config = config;
2122
}
2223

23-
public Task<string> Handle(string _)
24+
public Task<string> Handle(string _, CancellationToken cancellationToken = default)
2425
{
2526
var value = config.GetValue<string>("TestKey");
2627
return Task.FromResult(value);

‎tests/Compilation/Projects/Configuration/Handler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
using Amazon.Lambda.Core;
@@ -21,7 +22,7 @@ public Handler(IOptions<Options> exampleOptions)
2122
this.exampleOptions = exampleOptions.Value;
2223
}
2324

24-
public Task<Options> Handle(string request)
25+
public Task<Options> Handle(string request, CancellationToken cancellationToken = default)
2526
{
2627
return Task.FromResult(exampleOptions);
2728
}

‎tests/Compilation/Projects/CustomRunnerMethod/Handler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using System.Threading;
12
using System.Threading.Tasks;
23

34
using Amazon.Lambda.Core;
@@ -13,7 +14,7 @@ namespace Lambdajection.CompilationTests.CustomRunnerMethod
1314
[Lambda(typeof(Startup), RunnerMethod = "Process")]
1415
public partial class Handler
1516
{
16-
public Task<string> Handle(string request)
17+
public Task<string> Handle(string request, CancellationToken cancellationToken = default)
1718
{
1819
return Task.FromResult("Success");
1920
}

‎tests/Compilation/Projects/Disposables/AsyncDisposableHandler.cs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using System;
2+
using System.Threading;
23
using System.Threading.Tasks;
34

45
using Amazon.Lambda.Core;
@@ -17,7 +18,7 @@ public partial class AsyncDisposableHandler : IAsyncDisposable
1718
{
1819
public bool DisposeAsyncWasCalled = false;
1920

20-
public Task<AsyncDisposableHandler> Handle(string request)
21+
public Task<AsyncDisposableHandler> Handle(string request, CancellationToken cancellationToken = default)
2122
{
2223
return Task.FromResult(this);
2324
}

0 commit comments

Comments
 (0)