Add caching for Kroki diagrams

Directory.Packages.props

@@ -16,6 +16,7 @@
     <PackageVersion Include="AWSSDK.Core" Version="4.0.0.2" />
     <PackageVersion Include="AWSSDK.SQS" Version="4.0.0.1" />
     <PackageVersion Include="AWSSDK.S3" Version="4.0.0.1" />
+    <PackageVersion Include="Fake.Core.Context" Version="6.1.3" />
     <PackageVersion Include="FakeItEasy" Version="8.3.0" />
     <PackageVersion Include="Elastic.Ingest.Elasticsearch" Version="0.11.3" />
     <PackageVersion Include="Microsoft.OpenApi" Version="2.0.0-preview9" />

docs/_docset.yml

@@ -89,7 +89,7 @@ toc:
       - file: code.md
       - file: comments.md
       - file: conditionals.md
-      - hidden: diagrams.md
+      - file: diagrams.md
       - file: dropdowns.md
       - file: definition-lists.md
       - file: example_blocks.md

docs/syntax/diagrams.md

@@ -2,6 +2,10 @@
 
 The `diagram` directive allows you to render various types of diagrams using the [Kroki](https://kroki.io/) service. Kroki supports many diagram types including Mermaid, D2, Graphviz, PlantUML, and more.
 
+::::{warning}
+This is an experimental feature. It may change in the future.
+::::
+
 ## Basic usage
 
 The basic syntax for the diagram directive is:
@@ -84,7 +88,7 @@ sequenceDiagram
 :::::{tab-item} Rendered
 ::::{diagram} mermaid
 sequenceDiagram
-    participant A as Alice
+    participant A as Ada
     participant B as Bob
     A->>B: Hello Bob, how are you?
     B-->>A: Great!

src/Elastic.Documentation.Configuration/BuildContext.cs

@@ -6,6 +6,7 @@
 using System.Reflection;
 using Elastic.Documentation.Configuration.Assembler;
 using Elastic.Documentation.Configuration.Builder;
+using Elastic.Documentation.Configuration.Diagram;
 using Elastic.Documentation.Configuration.Versions;
 using Elastic.Documentation.Diagnostics;
 

src/Elastic.Documentation.Configuration/Diagram/DiagramRegistry.cs

@@ -0,0 +1,210 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+using System.Collections.Concurrent;
+using System.IO.Abstractions;
+using Elastic.Documentation.Diagnostics;
+using Elastic.Documentation.Extensions;
+using Microsoft.Extensions.Logging;
+
+namespace Elastic.Documentation.Configuration.Diagram;
+
+/// <summary>
+/// Information about a diagram that needs to be cached
+/// </summary>
+/// <param name="OutputFile">The intended cache output file location</param>
+/// <param name="EncodedUrl">Encoded Kroki URL for downloading</param>
+public record DiagramCacheInfo(IFileInfo OutputFile, string EncodedUrl);
+
+/// Registry to track active diagrams and manage cleanup of outdated cached files
+public class DiagramRegistry(ILoggerFactory logFactory, BuildContext context) : IDisposable
+{
+	private readonly ILogger<DiagramRegistry> _logger = logFactory.CreateLogger<DiagramRegistry>();
+	private readonly ConcurrentDictionary<string, DiagramCacheInfo> _diagramsToCache = new();
+	private readonly IFileSystem _writeFileSystem = context.WriteFileSystem;
+	private readonly IFileSystem _readFileSystem = context.ReadFileSystem;
+	private readonly HttpClient _httpClient = new() { Timeout = TimeSpan.FromSeconds(30) };
+
+	/// <summary>
+	/// Register a diagram for caching (collects info for later batch processing)
+	/// </summary>
+	/// <param name="localSvgPath">The local SVG path relative to the output directory</param>
+	/// <param name="encodedUrl">The encoded Kroki URL for downloading</param>
+	/// <param name="outputDirectory">The full path to the output directory</param>
+	public void RegisterDiagramForCaching(IFileInfo outputFile, string encodedUrl)
+	{
+		if (string.IsNullOrEmpty(encodedUrl))
+			return;
+
+		if (!outputFile.IsSubPathOf(context.DocumentationOutputDirectory))
+			return;
+
+		_ = _diagramsToCache.TryAdd(outputFile.FullName, new DiagramCacheInfo(outputFile, encodedUrl));
+	}
+
+	/// <summary>
+	/// Create cached diagram files by downloading from Kroki in parallel
+	/// </summary>
+	/// <returns>Number of diagrams downloaded</returns>
+	public async Task<int> CreateDiagramCachedFiles(Cancel ctx)
+	{
+		if (_diagramsToCache.IsEmpty)
+			return 0;
+
+		var downloadCount = 0;
+
+		await Parallel.ForEachAsync(_diagramsToCache.Values, new ParallelOptions
+		{
+			MaxDegreeOfParallelism = Environment.ProcessorCount,
+			CancellationToken = ctx
+		}, async (diagramInfo, ct) =>
+		{
+			var localPath = _readFileSystem.Path.GetRelativePath(context.DocumentationOutputDirectory.FullName, diagramInfo.OutputFile.FullName);
+
+			try
+			{
+				if (!diagramInfo.OutputFile.IsSubPathOf(context.DocumentationOutputDirectory))
+					return;
+
+				// Skip if the file already exists
+				if (_readFileSystem.File.Exists(diagramInfo.OutputFile.FullName))
+					return;
+
+				// If we are running on CI, and we are creating cached files we should fail the build and alert the user to create them
+				if (!string.IsNullOrEmpty(Environment.GetEnvironmentVariable("CI")))
+				{
+					context.Collector.EmitGlobalError($"Discovered new diagram SVG '{localPath}' please run `docs-builder --force` to ensure a cached version is generated");
+					return;
+				}
+
+				// Create the directory if needed
+				var directory = _writeFileSystem.Path.GetDirectoryName(diagramInfo.OutputFile.FullName);
+				if (directory != null && !_writeFileSystem.Directory.Exists(directory))
+					_ = _writeFileSystem.Directory.CreateDirectory(directory);
+
+				_logger.LogWarning("Creating local diagram:  {LocalPath}", localPath);
+				// Download SVG content
+				var svgContent = await _httpClient.GetStringAsync(diagramInfo.EncodedUrl, ct);
+
+				// Validate SVG content
+				if (string.IsNullOrWhiteSpace(svgContent) || !svgContent.Contains("<svg", StringComparison.OrdinalIgnoreCase))
+				{
+					_logger.LogWarning("Invalid SVG content received for diagram {LocalPath}", localPath);
+					return;
+				}
+
+				// Write atomically using a temp file
+				var tempPath = $"{diagramInfo.OutputFile.FullName}.tmp";
+				await _writeFileSystem.File.WriteAllTextAsync(tempPath, svgContent, ct);
+				_writeFileSystem.File.Move(tempPath, diagramInfo.OutputFile.FullName);
+
+				_ = Interlocked.Increment(ref downloadCount);
+				_logger.LogDebug("Downloaded diagram: {LocalPath}", localPath);
+			}
+			catch (HttpRequestException ex)
+			{
+				_logger.LogWarning("Failed to download diagram {LocalPath}: {Error}", localPath, ex.Message);
+			}
+			catch (TaskCanceledException ex) when (ex.InnerException is TimeoutException)
+			{
+				_logger.LogWarning("Timeout downloading diagram {LocalPath}", localPath);
+			}
+			catch (Exception ex)
+			{
+				_logger.LogWarning("Unexpected error downloading diagram {LocalPath}: {Error}", localPath, ex.Message);
+			}
+		});
+
+		if (downloadCount > 0)
+			_logger.LogInformation("Downloaded {DownloadCount} diagram files from Kroki", downloadCount);
+
+		return downloadCount;
+	}
+
+	/// <summary>
+	/// Clean up unused diagram files from the cache directory
+	/// </summary>
+	/// <returns>Number of files cleaned up</returns>
+	public int CleanupUnusedDiagrams()
+	{
+		if (!_readFileSystem.Directory.Exists(context.DocumentationOutputDirectory.FullName))
+			return 0;
+		var folders = _writeFileSystem.Directory.GetDirectories(context.DocumentationOutputDirectory.FullName, "generated-graphs", SearchOption.AllDirectories);
+		var existingFiles = folders
+			.Select(f => (Folder: f, Files: _writeFileSystem.Directory.GetFiles(f, "*.svg", SearchOption.TopDirectoryOnly)))
+			.ToArray();
+		if (existingFiles.Length == 0)
+			return 0;
+		var cleanedCount = 0;
+
+		try
+		{
+			foreach (var (folder, files) in existingFiles)
+			{
+				foreach (var file in files)
+				{
+					if (_diagramsToCache.ContainsKey(file))
+						continue;
+					try
+					{
+						_writeFileSystem.File.Delete(file);
+						cleanedCount++;
+					}
+					catch
+					{
+						// Silent failure - cleanup is opportunistic
+					}
+				}
+				// Clean up empty directories
+				CleanupEmptyDirectories(folder);
+			}
+		}
+		catch
+		{
+			// Silent failure - cleanup is opportunistic
+		}
+
+		return cleanedCount;
+	}
+
+	private void CleanupEmptyDirectories(string directory)
+	{
+		try
+		{
+			var folder = _writeFileSystem.DirectoryInfo.New(directory);
+			if (!folder.IsSubPathOf(context.DocumentationOutputDirectory))
+				return;
+
+			if (folder.Name != "generated-graphs")
+				return;
+
+			if (_writeFileSystem.Directory.EnumerateFileSystemEntries(folder.FullName).Any())
+				return;
+
+			_writeFileSystem.Directory.Delete(folder.FullName);
+
+			var parentFolder = folder.Parent;
+			if (parentFolder is null || parentFolder.Name != "images")
+				return;
+
+			if (_writeFileSystem.Directory.EnumerateFileSystemEntries(parentFolder.FullName).Any())
+				return;
+
+			_writeFileSystem.Directory.Delete(folder.FullName);
+		}
+		catch
+		{
+			// Silent failure - cleanup is opportunistic
+		}
+	}
+
+	/// <summary>
+	/// Dispose of resources, including the HttpClient
+	/// </summary>
+	public void Dispose()
+	{
+		_httpClient.Dispose();
+		GC.SuppressFinalize(this);
+	}
+}

src/Elastic.Documentation/Diagnostics/IDiagnosticsCollector.cs

@@ -53,5 +53,3 @@ public static void EmitGlobalWarning(this IDiagnosticsCollector collector, strin
 	public static void EmitGlobalHint(this IDiagnosticsCollector collector, string message) =>
 		collector.EmitHint(string.Empty, message);
 }
-
-

src/Elastic.Documentation/Extensions/IFileInfoExtensions.cs

@@ -18,4 +18,29 @@ public static string ReadToEnd(this IFileInfo fileInfo)
 		using var reader = new StreamReader(stream);
 		return reader.ReadToEnd();
 	}
+
+	/// Validates <paramref name="file"/> is in a subdirectory of <paramref name="parentDirectory"/>
+	public static bool IsSubPathOf(this IFileInfo file, IDirectoryInfo parentDirectory)
+	{
+		var parent = file.Directory;
+		return parent is not null && parent.IsSubPathOf(parentDirectory);
+	}
+}
+
+public static class IDirectoryInfoExtensions
+{
+	/// Validates <paramref name="directory"/> is subdirectory of <paramref name="parentDirectory"/>
+	public static bool IsSubPathOf(this IDirectoryInfo directory, IDirectoryInfo parentDirectory)
+	{
+		var parent = directory;
+		do
+		{
+			if (parent.FullName == parentDirectory.FullName)
+				return true;
+			parent = parent.Parent;
+		}
+		while (parent != null);
+
+		return false;
+	}
 }

src/Elastic.Markdown/DocumentationGenerator.cs

@@ -6,6 +6,7 @@
 using System.Text.Json;
 using Elastic.Documentation;
 using Elastic.Documentation.Configuration;
+using Elastic.Documentation.Configuration.Diagram;
 using Elastic.Documentation.Legacy;
 using Elastic.Documentation.Links;
 using Elastic.Documentation.Serialization;
@@ -16,6 +17,7 @@
 using Elastic.Markdown.Helpers;
 using Elastic.Markdown.IO;
 using Elastic.Markdown.Links.CrossLinks;
+using Elastic.Markdown.Myst.Directives.Diagram;
 using Elastic.Markdown.Myst.Renderers;
 using Elastic.Markdown.Myst.Renderers.LlmMarkdown;
 using Markdig.Syntax;
@@ -142,13 +144,36 @@ public async Task<GenerationResult> GenerateAll(Cancel ctx)
 		_logger.LogInformation($"Generating links.json");
 		var linkReference = await GenerateLinkReference(ctx);
 
+		await CreateDiagramCachedFiles(ctx);
+		CleanupUnusedDiagrams();
+
 		// ReSharper disable once WithExpressionModifiesAllMembers
 		return result with
 		{
 			Redirects = linkReference.Redirects ?? []
 		};
 	}
 
+	/// <summary>
+	/// Downloads diagram files in parallel from Kroki
+	/// </summary>
+	public async Task CreateDiagramCachedFiles(Cancel ctx)
+	{
+		var downloadedCount = await DocumentationSet.DiagramRegistry.CreateDiagramCachedFiles(ctx);
+		if (downloadedCount > 0)
+			_logger.LogInformation("Downloaded {DownloadedCount} diagram files from Kroki", downloadedCount);
+	}
+
+	/// <summary>
+	/// Cleans up unused diagram files from the output directory
+	/// </summary>
+	public void CleanupUnusedDiagrams()
+	{
+		var cleanedCount = DocumentationSet.DiagramRegistry.CleanupUnusedDiagrams();
+		if (cleanedCount > 0)
+			_logger.LogInformation("Cleaned up {CleanedCount} unused diagram files", cleanedCount);
+	}
+
 	private async Task ProcessDocumentationFiles(HashSet<string> offendingFiles, DateTimeOffset outputSeenChanges, Cancel ctx)
 	{
 		var processedFileCount = 0;

src/Elastic.Markdown/Extensions/DetectionRules/DetectionRuleFile.cs

@@ -29,14 +29,14 @@ protected override Task<MarkdownDocument> GetMinimalParseDocumentAsync(Cancel ct
 	{
 		Title = "Prebuilt detection rules reference";
 		var markdown = GetMarkdown();
-		var document = MarkdownParser.MinimalParseStringAsync(markdown, SourceFile, null);
+		var document = MarkdownParser.MinimalParseString(markdown, SourceFile, null);
 		return Task.FromResult(document);
 	}
 
 	protected override Task<MarkdownDocument> GetParseDocumentAsync(Cancel ctx)
 	{
 		var markdown = GetMarkdown();
-		var document = MarkdownParser.ParseStringAsync(markdown, SourceFile, null);
+		var document = MarkdownParser.ParseString(markdown, SourceFile, null);
 		return Task.FromResult(document);
 	}
 
@@ -127,14 +127,14 @@ protected override Task<MarkdownDocument> GetMinimalParseDocumentAsync(Cancel ct
 	{
 		Title = Rule?.Name;
 		var markdown = GetMarkdown();
-		var document = MarkdownParser.MinimalParseStringAsync(markdown, RuleSourceMarkdownPath, null);
+		var document = MarkdownParser.MinimalParseString(markdown, RuleSourceMarkdownPath, null);
 		return Task.FromResult(document);
 	}
 
 	protected override Task<MarkdownDocument> GetParseDocumentAsync(Cancel ctx)
 	{
 		var markdown = GetMarkdown();
-		var document = MarkdownParser.ParseStringAsync(markdown, RuleSourceMarkdownPath, null);
+		var document = MarkdownParser.ParseString(markdown, RuleSourceMarkdownPath, null);
 		return Task.FromResult(document);
 	}
 

src/Elastic.Markdown/HtmlWriter.cs

@@ -43,7 +43,7 @@ public class HtmlWriter(
 	public string Render(string markdown, IFileInfo? source)
 	{
 		source ??= DocumentationSet.Context.ConfigurationPath;
-		var parsed = DocumentationSet.MarkdownParser.ParseStringAsync(markdown, source, null);
+		var parsed = DocumentationSet.MarkdownParser.ParseString(markdown, source, null);
 		return MarkdownFile.CreateHtml(parsed);
 	}