Skip to content

Commit 92984d5

Browse files
committed
fix: forward template cancellation tokens
1 parent b4c0e50 commit 92984d5

2 files changed

Lines changed: 158 additions & 6 deletions

File tree

src/PatternKit.Generators/TemplateGenerator.cs

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -566,7 +566,9 @@ private string GenerateTemplateMethod(
566566
}
567567
else
568568
{
569-
sb.AppendLine($" {hook.Method.Name}(ctx);");
569+
var hasCt = hook.Method.Parameters.Any(IsCancellationToken);
570+
var args = hasCt ? "ctx, ct" : "ctx";
571+
sb.AppendLine($" {hook.Method.Name}({args});");
570572
}
571573
}
572574

@@ -588,7 +590,9 @@ private string GenerateTemplateMethod(
588590
}
589591
else
590592
{
591-
sb.AppendLine($" {step.Method.Name}(ctx);");
593+
var hasCt = step.Method.Parameters.Any(IsCancellationToken);
594+
var args = hasCt ? "ctx, ct" : "ctx";
595+
sb.AppendLine($" {step.Method.Name}({args});");
592596
}
593597
}
594598

@@ -604,7 +608,9 @@ private string GenerateTemplateMethod(
604608
}
605609
else
606610
{
607-
sb.AppendLine($" {hook.Method.Name}(ctx);");
611+
var hasCt = hook.Method.Parameters.Any(IsCancellationToken);
612+
var args = hasCt ? "ctx, ct" : "ctx";
613+
sb.AppendLine($" {hook.Method.Name}({args});");
608614
}
609615
}
610616

@@ -624,7 +630,9 @@ private string GenerateTemplateMethod(
624630
}
625631
else
626632
{
627-
sb.AppendLine($" {hook.Method.Name}(ctx, ex);");
633+
var hasCt = hook.Method.Parameters.Any(IsCancellationToken);
634+
var args = hasCt ? "ctx, ex, ct" : "ctx, ex";
635+
sb.AppendLine($" {hook.Method.Name}({args});");
628636
}
629637
}
630638

@@ -649,7 +657,9 @@ private string GenerateTemplateMethod(
649657
}
650658
else
651659
{
652-
sb.AppendLine($" {step.Method.Name}(ctx);");
660+
var hasCt = step.Method.Parameters.Any(IsCancellationToken);
661+
var args = hasCt ? "ctx, ct" : "ctx";
662+
sb.AppendLine($" {step.Method.Name}({args});");
653663
}
654664
}
655665

@@ -665,7 +675,9 @@ private string GenerateTemplateMethod(
665675
}
666676
else
667677
{
668-
sb.AppendLine($" {hook.Method.Name}(ctx);");
678+
var hasCt = hook.Method.Parameters.Any(IsCancellationToken);
679+
var args = hasCt ? "ctx, ct" : "ctx";
680+
sb.AppendLine($" {hook.Method.Name}({args});");
669681
}
670682
}
671683
}

test/PatternKit.Generators.Tests/TemplateGeneratorTests.cs

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1068,5 +1068,145 @@ public partial class ImportWorkflow
10681068
ScenarioExpect.True(emit.Success, string.Join("\n", emit.Diagnostics));
10691069
}
10701070

1071+
[Scenario("CancellationToken Sync Template Members AreForwardedInAsyncWrapper")]
1072+
[Fact]
1073+
public void CancellationToken_Sync_Template_Members_AreForwardedInAsyncWrapper()
1074+
{
1075+
var source = """
1076+
using PatternKit.Generators.Template;
1077+
using System;
1078+
using System.Threading;
1079+
1080+
namespace PatternKit.Examples;
1081+
1082+
public sealed class ImportContext
1083+
{
1084+
public int Count { get; set; }
1085+
}
1086+
1087+
[Template(ErrorPolicy = TemplateErrorPolicy.Rethrow)]
1088+
public partial class ImportWorkflow
1089+
{
1090+
[TemplateHook(HookPoint.BeforeAll)]
1091+
private void Open(ImportContext ctx, CancellationToken ct) => ctx.Count++;
1092+
1093+
[TemplateStep(0)]
1094+
private void Validate(ImportContext ctx, CancellationToken ct) => ctx.Count++;
1095+
1096+
[TemplateHook(HookPoint.AfterAll)]
1097+
private void Close(ImportContext ctx, CancellationToken ct) => ctx.Count++;
1098+
1099+
[TemplateHook(HookPoint.OnError)]
1100+
private void CaptureError(ImportContext ctx, Exception ex, CancellationToken ct) => ctx.Count--;
1101+
}
1102+
""";
1103+
1104+
var comp = RoslynTestHelpers.CreateCompilation(
1105+
source,
1106+
assemblyName: nameof(CancellationToken_Sync_Template_Members_AreForwardedInAsyncWrapper));
1107+
1108+
var gen = new TemplateGenerator();
1109+
_ = RoslynTestHelpers.Run(comp, gen, out var run, out var updated);
1110+
1111+
ScenarioExpect.All(run.Results, r => ScenarioExpect.Empty(r.Diagnostics));
1112+
1113+
var generated = run.Results
1114+
.SelectMany(r => r.GeneratedSources)
1115+
.Single(gs => gs.HintName == "ImportWorkflow.Template.g.cs")
1116+
.SourceText.ToString();
1117+
1118+
ScenarioExpect.Contains("public async System.Threading.Tasks.ValueTask ExecuteAsync", generated);
1119+
ScenarioExpect.Contains("Open(ctx, ct);", generated);
1120+
ScenarioExpect.Contains("Validate(ctx, ct);", generated);
1121+
ScenarioExpect.Contains("Close(ctx, ct);", generated);
1122+
ScenarioExpect.Contains("CaptureError(ctx, ex, ct);", generated);
1123+
1124+
var emit = updated.Emit(Stream.Null);
1125+
ScenarioExpect.True(emit.Success, string.Join("\n", emit.Diagnostics));
1126+
}
1127+
1128+
[Scenario("CancellationToken Async Template WithoutErrorHooks ForwardsAfterAllTokens")]
1129+
[Fact]
1130+
public void CancellationToken_Async_Template_WithoutErrorHooks_ForwardsAfterAllTokens()
1131+
{
1132+
var source = """
1133+
using PatternKit.Generators.Template;
1134+
using System.Threading;
1135+
using System.Threading.Tasks;
1136+
1137+
namespace PatternKit.Examples;
1138+
1139+
public sealed class ImportContext
1140+
{
1141+
public int Count { get; set; }
1142+
}
1143+
1144+
[Template]
1145+
public partial class ImportWorkflow
1146+
{
1147+
[TemplateStep(0)]
1148+
private void Validate(ImportContext ctx, CancellationToken ct) => ctx.Count++;
1149+
1150+
[TemplateHook(HookPoint.AfterAll)]
1151+
private ValueTask FlushAsync(ImportContext ctx, CancellationToken ct) => ValueTask.CompletedTask;
1152+
1153+
[TemplateHook(HookPoint.AfterAll)]
1154+
private void Close(ImportContext ctx, CancellationToken ct) => ctx.Count++;
1155+
}
1156+
""";
1157+
1158+
var comp = RoslynTestHelpers.CreateCompilation(
1159+
source,
1160+
assemblyName: nameof(CancellationToken_Async_Template_WithoutErrorHooks_ForwardsAfterAllTokens));
1161+
1162+
var gen = new TemplateGenerator();
1163+
_ = RoslynTestHelpers.Run(comp, gen, out var run, out var updated);
1164+
1165+
ScenarioExpect.All(run.Results, r => ScenarioExpect.Empty(r.Diagnostics));
1166+
1167+
var generated = run.Results
1168+
.SelectMany(r => r.GeneratedSources)
1169+
.Single(gs => gs.HintName == "ImportWorkflow.Template.g.cs")
1170+
.SourceText.ToString();
1171+
1172+
ScenarioExpect.Contains("Validate(ctx, ct);", generated);
1173+
ScenarioExpect.Contains("await FlushAsync(ctx, ct).ConfigureAwait(false);", generated);
1174+
ScenarioExpect.Contains("Close(ctx, ct);", generated);
1175+
ScenarioExpect.DoesNotContain("catch (System.Exception ex)", generated);
1176+
1177+
var emit = updated.Emit(Stream.Null);
1178+
ScenarioExpect.True(emit.Success, string.Join("\n", emit.Diagnostics));
1179+
}
1180+
1181+
[Scenario("Reports Error When Template Step HasNoContextParameter")]
1182+
[Fact]
1183+
public void Reports_Error_When_Template_Step_HasNoContextParameter()
1184+
{
1185+
var source = """
1186+
using PatternKit.Generators.Template;
1187+
1188+
namespace PatternKit.Examples;
1189+
1190+
[Template]
1191+
public partial class ImportWorkflow
1192+
{
1193+
[TemplateStep(0)]
1194+
private void Validate()
1195+
{
1196+
}
1197+
}
1198+
""";
1199+
1200+
var comp = RoslynTestHelpers.CreateCompilation(
1201+
source,
1202+
assemblyName: nameof(Reports_Error_When_Template_Step_HasNoContextParameter));
1203+
1204+
var gen = new TemplateGenerator();
1205+
_ = RoslynTestHelpers.Run(comp, gen, out var run, out _);
1206+
1207+
var diagnostics = run.Results.SelectMany(r => r.Diagnostics).ToArray();
1208+
ScenarioExpect.Contains(diagnostics, d => d.Id == "PKTMP004" && d.GetMessage().Contains("Validate"));
1209+
}
1210+
10711211
#endregion
10721212
}

0 commit comments

Comments
 (0)