chore: Add cancellation support for PdfBuilder operation (#10460)

* chore: add cancellation support for PdfBuilder

* chore: add CancellableCommand

Author: filzrev

src/Docfx.App/PdfBuilder.cs

@@ -51,17 +51,17 @@ class Outline
         public string? pdfFooterTemplate { get; init; }
     }
 
-    public static Task Run(BuildJsonConfig config, string configDirectory, string? outputDirectory = null)
+    public static Task Run(BuildJsonConfig config, string configDirectory, string? outputDirectory = null, CancellationToken cancellationToken = default)
     {
         var outputFolder = Path.GetFullPath(Path.Combine(
             string.IsNullOrEmpty(outputDirectory) ? Path.Combine(configDirectory, config.Output ?? "") : outputDirectory,
             config.Dest ?? ""));
 
         Logger.LogInfo($"Searching for manifest in {outputFolder}");
-        return CreatePdf(outputFolder);
+        return CreatePdf(outputFolder, cancellationToken);
     }
 
-    public static async Task CreatePdf(string outputFolder)
+    public static async Task CreatePdf(string outputFolder, CancellationToken cancellationToken = default)
     {
         var stopwatch = Stopwatch.StartNew();
         var pdfTocs = GetPdfTocs().ToDictionary(p => p.url, p => p.toc);
@@ -82,7 +82,7 @@ public static async Task CreatePdf(string outputFolder)
         using var app = builder.Build();
         app.UseServe(outputFolder);
         app.MapGet("/_pdftoc/{*url}", TocPage);
-        await app.StartAsync();
+        await app.StartAsync(cancellationToken);
 
         baseUrl = new Uri(app.Urls.First());
 
@@ -100,25 +100,51 @@ public static async Task CreatePdf(string outputFolder)
         var headerFooterTemplateCache = new ConcurrentDictionary<string, string>();
         var headerFooterPageCache = new ConcurrentDictionary<(string, string), Task<byte[]>>();
 
-        await AnsiConsole.Progress().StartAsync(async progress =>
+        var pdfBuildTask = AnsiConsole.Progress().StartAsync(async progress =>
         {
-            await Parallel.ForEachAsync(pdfTocs, async (item, _) =>
+            await Parallel.ForEachAsync(pdfTocs, new ParallelOptions { CancellationToken = cancellationToken }, async (item, _) =>
             {
                 var (url, toc) = item;
                 var outputName = Path.Combine(Path.GetDirectoryName(url) ?? "", toc.pdfFileName ?? Path.ChangeExtension(Path.GetFileName(url), ".pdf"));
                 var task = progress.AddTask(outputName);
-                var outputPath = Path.Combine(outputFolder, outputName);
+                var pdfOutputPath = Path.Combine(outputFolder, outputName);
 
                 await CreatePdf(
-                    PrintPdf, PrintHeaderFooter, task, new(baseUrl, url), toc, outputFolder, outputPath,
-                    pageNumbers => pdfPageNumbers[url] = pageNumbers);
+                    PrintPdf, PrintHeaderFooter, task, new(baseUrl, url), toc, outputFolder, pdfOutputPath,
+                    pageNumbers => pdfPageNumbers[url] = pageNumbers,
+                    cancellationToken);
 
                 task.Value = task.MaxValue;
                 task.StopTask();
             });
         });
 
+        try
+        {
+            await pdfBuildTask.WaitAsync(cancellationToken);
+        }
+        catch (OperationCanceledException)
+        {
+            if (!pdfBuildTask.IsCompleted)
+            {
+                // If pdf generation task is not completed.
+                // Manually close playwright context/browser to immediately shutdown remaining tasks.
+                await context.CloseAsync();
+                await browser.CloseAsync();
+                try
+                {
+                    await pdfBuildTask; // Wait AnsiConsole.Progress operation completed to output logs.
+                }
+                catch
+                {
+                    Logger.LogError($"PDF file generation is canceled by user interaction.");
+                    return;
+                }
+            }
+        }
+
         Logger.LogVerbose($"PDF done in {stopwatch.Elapsed}");
+        return;
 
         IEnumerable<(string url, Outline toc)> GetPdfTocs()
         {
@@ -150,7 +176,7 @@ IResult TocPage(string url)
 
         async Task<byte[]?> PrintPdf(Outline outline, Uri url)
         {
-            await pageLimiter.WaitAsync();
+            await pageLimiter.WaitAsync(cancellationToken);
             var page = pagePool.TryTake(out var pooled) ? pooled : await context.NewPageAsync();
 
             try
@@ -273,7 +299,7 @@ static string ExpandTemplate(string? pdfTemplate, int pageNumber, int totalPages
 
     static async Task CreatePdf(
         Func<Outline, Uri, Task<byte[]?>> printPdf, Func<Outline, int, int, Page, Task<byte[]>> printHeaderFooter, ProgressTask task,
-        Uri outlineUrl, Outline outline, string outputFolder, string outputPath, Action<Dictionary<Outline, int>> updatePageNumbers)
+        Uri outlineUrl, Outline outline, string outputFolder, string pdfOutputPath, Action<Dictionary<Outline, int>> updatePageNumbers, CancellationToken cancellationToken)
     {
         var pages = GetPages(outline).ToArray();
         if (pages.Length == 0)
@@ -284,7 +310,7 @@ static async Task CreatePdf(
         // Make progress at 99% before merge PDF
         task.MaxValue = pages.Length + (pages.Length / 99.0);
 
-        await Parallel.ForEachAsync(pages, async (item, _) =>
+        await Parallel.ForEachAsync(pages, new ParallelOptions { CancellationToken = cancellationToken }, async (item, _) =>
         {
             var (url, node) = item;
             if (await printPdf(outline, url) is { } bytes)
@@ -302,6 +328,8 @@ await Parallel.ForEachAsync(pages, async (item, _) =>
 
         foreach (var (url, node) in pages)
         {
+            cancellationToken.ThrowIfCancellationRequested();
+
             if (!pageBytes.TryGetValue(node, out var bytes))
                 continue;
 
@@ -324,13 +352,14 @@ await Parallel.ForEachAsync(pages, async (item, _) =>
 
         var producer = $"docfx ({typeof(PdfBuilder).Assembly.GetCustomAttribute<AssemblyFileVersionAttribute>()?.Version})";
 
-        using var output = File.Create(outputPath);
+        using var output = File.Create(pdfOutputPath);
         using var builder = new PdfDocumentBuilder(output);
 
         builder.DocumentInformation = new() { Producer = producer };
         builder.Bookmarks = CreateBookmarks(outline.items);
 
         await MergePdf();
+        return;
 
         IEnumerable<(Uri url, Outline node)> GetPages(Outline outline)
         {
@@ -368,6 +397,8 @@ async Task MergePdf()
 
             foreach (var (url, node) in pages)
             {
+                cancellationToken.ThrowIfCancellationRequested();
+
                 if (!pageBytes.TryGetValue(node, out var bytes))
                     continue;
 
@@ -387,6 +418,8 @@ async Task MergePdf()
                 using var document = PdfDocument.Open(bytes);
                 for (var i = 1; i <= document.NumberOfPages; i++)
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
+
                     pageNumber++;
 
                     var pageBuilder = builder.AddPage(document, i, x => CopyLink(node, x));

src/docfx/Models/CancellableCommandBase.cs

@@ -0,0 +1,33 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+#nullable enable
+
+using System.Runtime.InteropServices;
+using Spectre.Console.Cli;
+
+namespace Docfx;
+
+public abstract class CancellableCommandBase<TSettings> : Command<TSettings>
+     where TSettings : CommandSettings
+{
+    public abstract int Execute(CommandContext context, TSettings settings, CancellationToken cancellation);
+
+    public sealed override int Execute(CommandContext context, TSettings settings)
+    {
+        using var cancellationSource = new CancellationTokenSource();
+
+        using var sigInt = PosixSignalRegistration.Create(PosixSignal.SIGINT, onSignal);
+        using var sigQuit = PosixSignalRegistration.Create(PosixSignal.SIGQUIT, onSignal);
+        using var sigTerm = PosixSignalRegistration.Create(PosixSignal.SIGTERM, onSignal);
+
+        var exitCode = Execute(context, settings, cancellationSource.Token);
+        return exitCode;
+
+        void onSignal(PosixSignalContext context)
+        {
+            context.Cancel = true;
+            cancellationSource.Cancel();
+        }
+    }
+}

src/docfx/Models/DefaultCommand.cs

@@ -10,7 +10,7 @@
 
 namespace Docfx;
 
-class DefaultCommand : Command<DefaultCommand.Options>
+class DefaultCommand : CancellableCommandBase<DefaultCommand.Options>
 {
     [Description("Runs metadata, build and pdf commands")]
     internal class Options : BuildCommandOptions
@@ -20,7 +20,7 @@ internal class Options : BuildCommandOptions
         public bool Version { get; set; }
     }
 
-    public override int Execute(CommandContext context, Options options)
+    public override int Execute(CommandContext context, Options options, CancellationToken cancellationToken)
     {
         if (options.Version)
         {
@@ -48,9 +48,9 @@ public override int Execute(CommandContext context, Options options)
             if (config.build is not null)
             {
                 BuildCommand.MergeOptionsToConfig(options, config.build, configDirectory);
-                serveDirectory = RunBuild.Exec(config.build, new(), configDirectory, outputFolder);
+                serveDirectory = RunBuild.Exec(config.build, new(), configDirectory, outputFolder, cancellationToken);
 
-                PdfBuilder.CreatePdf(serveDirectory).GetAwaiter().GetResult();
+                PdfBuilder.CreatePdf(serveDirectory, cancellationToken).GetAwaiter().GetResult();
             }
 
             if (options.Serve && serveDirectory is not null)

src/docfx/Models/PdfCommand.cs

@@ -1,22 +1,23 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Diagnostics.CodeAnalysis;
 using Docfx.Pdf;
 using Spectre.Console.Cli;
 
+#nullable enable
+
 namespace Docfx;
 
-internal class PdfCommand : Command<PdfCommandOptions>
+internal class PdfCommand : CancellableCommandBase<PdfCommandOptions>
 {
-    public override int Execute([NotNull] CommandContext context, [NotNull] PdfCommandOptions options)
+    public override int Execute(CommandContext context, PdfCommandOptions options, CancellationToken cancellationToken)
     {
         return CommandHelper.Run(options, () =>
         {
             var (config, configDirectory) = Docset.GetConfig(options.ConfigFile);
 
             if (config.build is not null)
-                PdfBuilder.Run(config.build, configDirectory, options.OutputFolder).GetAwaiter().GetResult();
+                PdfBuilder.Run(config.build, configDirectory, options.OutputFolder, cancellationToken).GetAwaiter().GetResult();
         });
     }
 }