From eb44883da2ddb2d9b67470954ddde0c69ec0524e Mon Sep 17 00:00:00 2001
From: Lucas Trzesniewski <lucas.trzesniewski@gmail.com>
Date: Wed, 3 Apr 2024 23:17:26 +0200
Subject: [PATCH 1/6] Add FlushAsync() method

---
 src/RazorBlade.Library/HtmlLayout.cs    |   6 +-
 src/RazorBlade.Library/RazorTemplate.cs | 159 +++++++++++++++++++-----
 2 files changed, 128 insertions(+), 37 deletions(-)

diff --git a/src/RazorBlade.Library/HtmlLayout.cs b/src/RazorBlade.Library/HtmlLayout.cs
index d02b3f0..82cbf72 100644
--- a/src/RazorBlade.Library/HtmlLayout.cs
+++ b/src/RazorBlade.Library/HtmlLayout.cs
@@ -21,7 +21,7 @@ async Task<IRazorExecutionResult> IRazorLayout.ExecuteLayoutAsync(IRazorExecutio
         try
         {
             _layoutInput = input;
-            return await ExecuteAsyncCore(input.CancellationToken);
+            return await ExecuteAsyncCore(null, input.CancellationToken);
         }
         finally
         {
@@ -29,12 +29,12 @@ async Task<IRazorExecutionResult> IRazorLayout.ExecuteLayoutAsync(IRazorExecutio
         }
     }
 
-    private protected override Task<IRazorExecutionResult> ExecuteAsyncCore(CancellationToken cancellationToken)
+    private protected override Task<IRazorExecutionResult> ExecuteAsyncCore(TextWriter? targetOutput, CancellationToken cancellationToken)
     {
         if (_layoutInput is null)
             throw new InvalidOperationException(_contentsRequiredErrorMessage);
 
-        return base.ExecuteAsyncCore(cancellationToken);
+        return base.ExecuteAsyncCore(targetOutput, cancellationToken);
     }
 
     /// <summary>
diff --git a/src/RazorBlade.Library/RazorTemplate.cs b/src/RazorBlade.Library/RazorTemplate.cs
index 11faf74..3d37ba7 100644
--- a/src/RazorBlade.Library/RazorTemplate.cs
+++ b/src/RazorBlade.Library/RazorTemplate.cs
@@ -15,6 +15,8 @@ namespace RazorBlade;
 /// </summary>
 public abstract class RazorTemplate : IEncodedContent
 {
+    private ExecutionScope? _executionScope;
+    private IRazorLayout? _layout;
     private Dictionary<string, Func<Task>>? _sections;
 
     private Dictionary<string, Func<Task>> Sections => _sections ??= new(StringComparer.OrdinalIgnoreCase);
@@ -22,7 +24,7 @@ public abstract class RazorTemplate : IEncodedContent
     /// <summary>
     /// The <see cref="TextWriter"/> which receives the output.
     /// </summary>
-    protected internal TextWriter Output { get; internal set; } = new StreamWriter(Stream.Null);
+    protected internal TextWriter Output { get; internal set; } = TextWriter.Null;
 
     /// <summary>
     /// The cancellation token.
@@ -32,7 +34,18 @@ public abstract class RazorTemplate : IEncodedContent
     /// <summary>
     /// The layout to use.
     /// </summary>
-    private protected IRazorLayout? Layout { get; set; }
+    private protected IRazorLayout? Layout
+    {
+        get => _layout;
+        set
+        {
+            if (ReferenceEquals(value, _layout))
+                return;
+
+            _executionScope?.EnsureCanChangeLayout();
+            _layout = value;
+        }
+    }
 
     /// <summary>
     /// Renders the template synchronously and returns the result as a string.
@@ -46,11 +59,23 @@ public string Render(CancellationToken cancellationToken = default)
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        var renderTask = RenderAsyncCore(cancellationToken);
+        var textWriter = new StringWriter();
+        var renderTask = RenderAsyncCore(textWriter, cancellationToken);
+
         if (renderTask.IsCompleted)
-            return renderTask.GetAwaiter().GetResult().ToString();
+        {
+            renderTask.GetAwaiter().GetResult();
+            return textWriter.ToString();
+        }
 
-        return Task.Run(async () => await renderTask.ConfigureAwait(false), CancellationToken.None).GetAwaiter().GetResult().ToString();
+        return Task.Run(
+            async () =>
+            {
+                await renderTask.ConfigureAwait(false);
+                return textWriter.ToString();
+            },
+            CancellationToken.None
+        ).GetAwaiter().GetResult();
     }
 
     /// <summary>
@@ -66,7 +91,8 @@ public void Render(TextWriter textWriter, CancellationToken cancellationToken =
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        var renderTask = RenderAsync(textWriter, cancellationToken);
+        var renderTask = RenderAsyncCore(textWriter, cancellationToken);
+
         if (renderTask.IsCompleted)
         {
             renderTask.GetAwaiter().GetResult();
@@ -87,8 +113,9 @@ public async Task<string> RenderAsync(CancellationToken cancellationToken = defa
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        var stringBuilder = await RenderAsyncCore(cancellationToken).ConfigureAwait(false);
-        return stringBuilder.ToString();
+        var textWriter = new StringWriter();
+        await RenderAsyncCore(textWriter, cancellationToken).ConfigureAwait(false);
+        return textWriter.ToString();
     }
 
     /// <summary>
@@ -103,23 +130,17 @@ public async Task RenderAsync(TextWriter textWriter, CancellationToken cancellat
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        var stringBuilder = await RenderAsyncCore(cancellationToken).ConfigureAwait(false);
-
-#if NET6_0_OR_GREATER
-        await textWriter.WriteAsync(stringBuilder, cancellationToken).ConfigureAwait(false);
-#else
-        await textWriter.WriteAsync(stringBuilder.ToString()).ConfigureAwait(false);
-#endif
+        await RenderAsyncCore(textWriter, cancellationToken).ConfigureAwait(false);
     }
 
     /// <summary>
     /// Renders the template asynchronously including its layout and returns the result as a <see cref="StringBuilder"/>.
     /// </summary>
-    private async Task<StringBuilder> RenderAsyncCore(CancellationToken cancellationToken)
+    private async Task RenderAsyncCore(TextWriter targetOutput, CancellationToken cancellationToken)
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        var executionResult = await ExecuteAsyncCore(cancellationToken);
+        var executionResult = await ExecuteAsyncCore(targetOutput, cancellationToken);
 
         while (executionResult.Layout is { } layout)
         {
@@ -127,27 +148,39 @@ private async Task<StringBuilder> RenderAsyncCore(CancellationToken cancellation
             executionResult = await layout.ExecuteLayoutAsync(executionResult).ConfigureAwait(false);
         }
 
-        if (executionResult.Body is EncodedContent { Output: var outputStringBuilder })
-            return outputStringBuilder;
+        switch (executionResult.Body)
+        {
+            case EncodedContent { Output: var bufferedOutput }:
+                await WriteStringBuilderToOutputAndFlushAsync(bufferedOutput, targetOutput, cancellationToken).ConfigureAwait(false);
+                break;
 
-        // Fallback case, shouldn't happen
-        var outputStringWriter = new StringWriter();
-        executionResult.Body.WriteTo(outputStringWriter);
-        return outputStringWriter.GetStringBuilder();
+            case { } body: // Fallback case, shouldn't happen
+                body.WriteTo(targetOutput);
+                break;
+        }
     }
 
     /// <summary>
     /// Calls the <see cref="ExecuteAsync"/> method in a new <see cref="ExecutionScope"/>.
     /// </summary>
-    private protected virtual async Task<IRazorExecutionResult> ExecuteAsyncCore(CancellationToken cancellationToken)
+    private protected virtual async Task<IRazorExecutionResult> ExecuteAsyncCore(TextWriter? targetOutput, CancellationToken cancellationToken)
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        using var executionScope = new ExecutionScope(this, cancellationToken);
+        using var executionScope = ExecutionScope.Start(this, targetOutput, cancellationToken);
         await ExecuteAsync().ConfigureAwait(false);
         return new ExecutionResult(executionScope);
     }
 
+    /// <summary>
+    /// Writes the buffered output to the target output then flushes the output stream.
+    /// </summary>
+    /// <remarks>
+    /// This feature is not compatible with layouts.
+    /// </remarks>
+    protected Task FlushAsync()
+        => _executionScope?.FlushAsync() ?? Task.CompletedTask;
+
     /// <summary>
     /// Executes the template and appends the result to <see cref="Output"/>.
     /// </summary>
@@ -230,6 +263,27 @@ protected internal void DefineSection(string name, Func<Task> action)
 #endif
     }
 
+    /// <summary>
+    /// Writes the contents of a <see cref="StringBuilder"/> to a <see cref="TextWriter"/> asynchronuously.
+    /// </summary>
+    private static async Task WriteStringBuilderToOutputAndFlushAsync(StringBuilder input, TextWriter output, CancellationToken cancellationToken)
+    {
+        if (input.Length == 0)
+            return;
+
+#if NET6_0_OR_GREATER
+        await output.WriteAsync(input, cancellationToken).ConfigureAwait(false);
+#else
+        await output.WriteAsync(input.ToString()).ConfigureAwait(false);
+#endif
+
+#if NET8_0_OR_GREATER
+        await output.FlushAsync(cancellationToken).ConfigureAwait(false);
+#else
+        await output.FlushAsync().ConfigureAwait(false);
+#endif
+    }
+
     void IEncodedContent.WriteTo(TextWriter textWriter)
         => Render(textWriter, CancellationToken.None);
 
@@ -238,35 +292,69 @@ void IEncodedContent.WriteTo(TextWriter textWriter)
     /// </summary>
     private class ExecutionScope : IDisposable
     {
+        private readonly ExecutionScope? _previousExecutionScope;
         private readonly Dictionary<string, Func<Task>>? _previousSections;
         private readonly TextWriter _previousOutput;
         private readonly CancellationToken _previousCancellationToken;
         private readonly IRazorLayout? _previousLayout;
 
+        private readonly TextWriter? _targetOutput;
+        private bool _layoutFrozen;
+
         public RazorTemplate Page { get; }
-        public StringBuilder Output { get; } = new();
+        public StringBuilder BufferedOutput { get; } = new();
 
-        public ExecutionScope(RazorTemplate page, CancellationToken cancellationToken)
+        private ExecutionScope(RazorTemplate page, TextWriter? targetOutput, CancellationToken cancellationToken)
         {
             Page = page;
+            _targetOutput = targetOutput;
 
+            _previousExecutionScope = page._executionScope;
             _previousSections = page._sections;
             _previousOutput = page.Output;
             _previousCancellationToken = page.CancellationToken;
             _previousLayout = page.Layout;
 
+            page._executionScope = this;
+            page._layout = null;
             page._sections = null;
-            page.Output = new StringWriter(Output);
+            page.Output = new StringWriter(BufferedOutput);
             page.CancellationToken = cancellationToken;
-            page.Layout = null;
         }
 
+        public static ExecutionScope Start(RazorTemplate page, TextWriter? targetOutput, CancellationToken cancellationToken)
+            => new(page, targetOutput, cancellationToken);
+
         public void Dispose()
         {
+            Page._executionScope = _previousExecutionScope;
+            Page._layout = _previousLayout;
             Page._sections = _previousSections;
             Page.Output = _previousOutput;
             Page.CancellationToken = _previousCancellationToken;
-            Page.Layout = _previousLayout;
+        }
+
+        public void FreezeLayout()
+            => _layoutFrozen = true;
+
+        public void EnsureCanChangeLayout()
+        {
+            if (_layoutFrozen)
+                throw new InvalidOperationException("The layout can no longer be changed.");
+        }
+
+        public async Task FlushAsync()
+        {
+            if (Page.Layout is not null)
+                throw new InvalidOperationException("The output cannot be flushed when a layout is used.");
+
+            FreezeLayout();
+
+            if (_targetOutput is not null)
+            {
+                await WriteStringBuilderToOutputAndFlushAsync(BufferedOutput, _targetOutput, Page.CancellationToken).ConfigureAwait(false);
+                BufferedOutput.Clear();
+            }
         }
     }
 
@@ -286,7 +374,7 @@ public ExecutionResult(ExecutionScope executionScope)
         {
             _page = executionScope.Page;
             _sections = _page._sections;
-            Body = new EncodedContent(executionScope.Output);
+            Body = new EncodedContent(executionScope.BufferedOutput);
             Layout = _page.Layout;
             CancellationToken = _page.CancellationToken;
         }
@@ -299,10 +387,13 @@ public bool IsSectionDefined(string name)
             if (_sections is null || !_sections.TryGetValue(name, out var sectionAction))
                 return null;
 
-            using var executionScope = new ExecutionScope(_page, CancellationToken);
-            _page.Layout = Layout; // The section might reference this instance.
+            using var executionScope = ExecutionScope.Start(_page, null, CancellationToken);
+
+            _page._layout = Layout; // The section might reference this instance.
+            executionScope.FreezeLayout();
+
             await sectionAction().ConfigureAwait(false);
-            return new EncodedContent(executionScope.Output);
+            return new EncodedContent(executionScope.BufferedOutput);
         }
     }
 

From 7f50d7aeb957719ba7b5aff9e98e6388b3458831 Mon Sep 17 00:00:00 2001
From: Lucas Trzesniewski <lucas.trzesniewski@gmail.com>
Date: Fri, 5 Apr 2024 21:44:52 +0200
Subject: [PATCH 2/6] Refactor RazorTemplate

---
 src/RazorBlade.Library/RazorTemplate.cs       | 193 +++++++++---------
 src/RazorBlade.Tests/HtmlLayoutTests.cs       |  18 +-
 src/RazorBlade.Tests/HtmlTemplateTests.cs     |   9 +-
 .../PlainTextTemplateTests.cs                 |   9 +-
 src/RazorBlade.Tests/RazorTemplateTests.cs    |  36 ++--
 5 files changed, 121 insertions(+), 144 deletions(-)

diff --git a/src/RazorBlade.Library/RazorTemplate.cs b/src/RazorBlade.Library/RazorTemplate.cs
index 3d37ba7..a2aae76 100644
--- a/src/RazorBlade.Library/RazorTemplate.cs
+++ b/src/RazorBlade.Library/RazorTemplate.cs
@@ -16,35 +16,24 @@ namespace RazorBlade;
 public abstract class RazorTemplate : IEncodedContent
 {
     private ExecutionScope? _executionScope;
-    private IRazorLayout? _layout;
-    private Dictionary<string, Func<Task>>? _sections;
-
-    private Dictionary<string, Func<Task>> Sections => _sections ??= new(StringComparer.OrdinalIgnoreCase);
 
     /// <summary>
     /// The <see cref="TextWriter"/> which receives the output.
     /// </summary>
-    protected internal TextWriter Output { get; internal set; } = TextWriter.Null;
+    protected internal TextWriter Output => _executionScope?.BufferedOutput ?? TextWriter.Null;
 
     /// <summary>
     /// The cancellation token.
     /// </summary>
-    protected internal CancellationToken CancellationToken { get; private set; }
+    protected internal CancellationToken CancellationToken => _executionScope?.CancellationToken ?? CancellationToken.None;
 
     /// <summary>
     /// The layout to use.
     /// </summary>
     private protected IRazorLayout? Layout
     {
-        get => _layout;
-        set
-        {
-            if (ReferenceEquals(value, _layout))
-                return;
-
-            _executionScope?.EnsureCanChangeLayout();
-            _layout = value;
-        }
+        get => _executionScope?.Layout;
+        set => (_executionScope ?? throw new InvalidOperationException("The layout can only be set while the template is executing.")).SetLayout(value);
     }
 
     /// <summary>
@@ -150,8 +139,8 @@ private async Task RenderAsyncCore(TextWriter targetOutput, CancellationToken ca
 
         switch (executionResult.Body)
         {
-            case EncodedContent { Output: var bufferedOutput }:
-                await WriteStringBuilderToOutputAndFlushAsync(bufferedOutput, targetOutput, cancellationToken).ConfigureAwait(false);
+            case BufferedContent { Output: var bufferedOutput }:
+                await WriteStringBuilderToOutputAsync(bufferedOutput, targetOutput, cancellationToken).ConfigureAwait(false);
                 break;
 
             case { } body: // Fallback case, shouldn't happen
@@ -167,7 +156,7 @@ private protected virtual async Task<IRazorExecutionResult> ExecuteAsyncCore(Tex
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        using var executionScope = ExecutionScope.Start(this, targetOutput, cancellationToken);
+        using var executionScope = ExecutionScope.StartBody(this, targetOutput, cancellationToken);
         await ExecuteAsync().ConfigureAwait(false);
         return new ExecutionResult(executionScope);
     }
@@ -252,35 +241,24 @@ protected internal virtual void Write(IEncodedContent? content)
     [EditorBrowsable(EditorBrowsableState.Never)]
     protected internal void DefineSection(string name, Func<Task> action)
     {
-#if NET6_0_OR_GREATER
-        if (!Sections.TryAdd(name, action))
-            throw new InvalidOperationException($"Section '{name}' is already defined.");
-#else
-        if (Sections.ContainsKey(name))
-            throw new InvalidOperationException($"Section '{name}' is already defined.");
+        if (_executionScope is not { } executionScope)
+            throw new InvalidOperationException("Sections can only be defined while the template is executing.");
 
-        Sections[name] = action;
-#endif
+        executionScope.DefineSection(name, action);
     }
 
     /// <summary>
     /// Writes the contents of a <see cref="StringBuilder"/> to a <see cref="TextWriter"/> asynchronuously.
     /// </summary>
-    private static async Task WriteStringBuilderToOutputAndFlushAsync(StringBuilder input, TextWriter output, CancellationToken cancellationToken)
+    private static Task WriteStringBuilderToOutputAsync(StringBuilder input, TextWriter output, CancellationToken cancellationToken)
     {
         if (input.Length == 0)
-            return;
+            return Task.CompletedTask;
 
 #if NET6_0_OR_GREATER
-        await output.WriteAsync(input, cancellationToken).ConfigureAwait(false);
+        return output.WriteAsync(input, cancellationToken);
 #else
-        await output.WriteAsync(input.ToString()).ConfigureAwait(false);
-#endif
-
-#if NET8_0_OR_GREATER
-        await output.FlushAsync(cancellationToken).ConfigureAwait(false);
-#else
-        await output.FlushAsync().ConfigureAwait(false);
+        return output.WriteAsync(input.ToString());
 #endif
     }
 
@@ -288,74 +266,112 @@ void IEncodedContent.WriteTo(TextWriter textWriter)
         => Render(textWriter, CancellationToken.None);
 
     /// <summary>
-    /// Saves the current state, resets it, and restores it when disposed.
+    /// Stores the state of a template execution.
     /// </summary>
     private class ExecutionScope : IDisposable
     {
+        private readonly RazorTemplate _page;
+        private readonly TextWriter? _targetOutput;
         private readonly ExecutionScope? _previousExecutionScope;
-        private readonly Dictionary<string, Func<Task>>? _previousSections;
-        private readonly TextWriter _previousOutput;
-        private readonly CancellationToken _previousCancellationToken;
-        private readonly IRazorLayout? _previousLayout;
 
-        private readonly TextWriter? _targetOutput;
+        private IRazorLayout? _layout;
         private bool _layoutFrozen;
+        private Dictionary<string, Func<Task>>? _sections;
+
+        public StringWriter BufferedOutput { get; } = new();
+        public IRazorLayout? Layout => _layout;
+        public CancellationToken CancellationToken { get; }
+
+        public static ExecutionScope StartBody(RazorTemplate page, TextWriter? targetOutput, CancellationToken cancellationToken)
+            => new(page, targetOutput, cancellationToken);
 
-        public RazorTemplate Page { get; }
-        public StringBuilder BufferedOutput { get; } = new();
+        private static ExecutionScope StartSection(ExecutionScope parent)
+            => new(parent._page, null, parent.CancellationToken)
+            {
+                _layout = parent._layout, // The section might reference the layout instance.
+                _layoutFrozen = true
+            };
 
         private ExecutionScope(RazorTemplate page, TextWriter? targetOutput, CancellationToken cancellationToken)
         {
-            Page = page;
+            _page = page;
             _targetOutput = targetOutput;
+            CancellationToken = cancellationToken;
 
             _previousExecutionScope = page._executionScope;
-            _previousSections = page._sections;
-            _previousOutput = page.Output;
-            _previousCancellationToken = page.CancellationToken;
-            _previousLayout = page.Layout;
-
             page._executionScope = this;
-            page._layout = null;
-            page._sections = null;
-            page.Output = new StringWriter(BufferedOutput);
-            page.CancellationToken = cancellationToken;
         }
 
-        public static ExecutionScope Start(RazorTemplate page, TextWriter? targetOutput, CancellationToken cancellationToken)
-            => new(page, targetOutput, cancellationToken);
-
         public void Dispose()
         {
-            Page._executionScope = _previousExecutionScope;
-            Page._layout = _previousLayout;
-            Page._sections = _previousSections;
-            Page.Output = _previousOutput;
-            Page.CancellationToken = _previousCancellationToken;
+            if (ReferenceEquals(_page._executionScope, this))
+                _page._executionScope = _previousExecutionScope;
         }
 
-        public void FreezeLayout()
-            => _layoutFrozen = true;
-
-        public void EnsureCanChangeLayout()
+        public void SetLayout(IRazorLayout? layout)
         {
+            if (ReferenceEquals(layout, _layout))
+                return;
+
             if (_layoutFrozen)
                 throw new InvalidOperationException("The layout can no longer be changed.");
+
+            _layout = layout;
         }
 
         public async Task FlushAsync()
         {
-            if (Page.Layout is not null)
+            if (_layout is not null)
                 throw new InvalidOperationException("The output cannot be flushed when a layout is used.");
 
-            FreezeLayout();
+            // A part of the output will be written to the target output and discarded,
+            // so disallow setting a layout later on, as that would lead to inconsistent results.
+            _layoutFrozen = true;
 
             if (_targetOutput is not null)
             {
-                await WriteStringBuilderToOutputAndFlushAsync(BufferedOutput, _targetOutput, Page.CancellationToken).ConfigureAwait(false);
-                BufferedOutput.Clear();
+                var bufferedOutput = BufferedOutput.GetStringBuilder();
+                await WriteStringBuilderToOutputAsync(bufferedOutput, _targetOutput, CancellationToken).ConfigureAwait(false);
+                bufferedOutput.Clear();
+
+#if NET8_0_OR_GREATER
+                await _targetOutput.FlushAsync(CancellationToken).ConfigureAwait(false);
+#else
+                await _targetOutput.FlushAsync().ConfigureAwait(false);
+#endif
             }
         }
+
+        public BufferedContent ToBufferedContent()
+            => new(BufferedOutput.GetStringBuilder());
+
+        public bool IsSectionDefined(string name)
+            => _sections is { } sections && sections.ContainsKey(name);
+
+        public void DefineSection(string name, Func<Task> action)
+        {
+            var sections = _sections ??= new(StringComparer.OrdinalIgnoreCase);
+
+#if NET6_0_OR_GREATER
+            if (!sections.TryAdd(name, action))
+                throw new InvalidOperationException($"Section '{name}' is already defined.");
+#else
+            if (sections.ContainsKey(name))
+                throw new InvalidOperationException($"Section '{name}' is already defined.");
+
+            sections[name] = action;
+#endif
+        }
+
+        public async Task<IEncodedContent?> RenderSectionAsync(string name)
+        {
+            if (_sections is not { } sections || !sections.TryGetValue(name, out var sectionAction))
+                return null;
+
+            using var sectionScope = StartSection(this);
+            await sectionAction().ConfigureAwait(false);
+            return sectionScope.ToBufferedContent();
+        }
     }
 
     /// <summary>
@@ -363,38 +379,23 @@ public async Task FlushAsync()
     /// </summary>
     private class ExecutionResult : IRazorExecutionResult
     {
-        private readonly RazorTemplate _page;
-        private readonly IReadOnlyDictionary<string, Func<Task>>? _sections;
+        private readonly ExecutionScope _executionScope;
 
         public IEncodedContent Body { get; }
-        public IRazorLayout? Layout { get; }
-        public CancellationToken CancellationToken { get; }
+        public IRazorLayout? Layout => _executionScope.Layout;
+        public CancellationToken CancellationToken => _executionScope.CancellationToken;
 
         public ExecutionResult(ExecutionScope executionScope)
         {
-            _page = executionScope.Page;
-            _sections = _page._sections;
-            Body = new EncodedContent(executionScope.BufferedOutput);
-            Layout = _page.Layout;
-            CancellationToken = _page.CancellationToken;
+            _executionScope = executionScope;
+            Body = executionScope.ToBufferedContent();
         }
 
         public bool IsSectionDefined(string name)
-            => _sections?.ContainsKey(name) is true;
+            => _executionScope.IsSectionDefined(name);
 
-        public async Task<IEncodedContent?> RenderSectionAsync(string name)
-        {
-            if (_sections is null || !_sections.TryGetValue(name, out var sectionAction))
-                return null;
-
-            using var executionScope = ExecutionScope.Start(_page, null, CancellationToken);
-
-            _page._layout = Layout; // The section might reference this instance.
-            executionScope.FreezeLayout();
-
-            await sectionAction().ConfigureAwait(false);
-            return new EncodedContent(executionScope.BufferedOutput);
-        }
+        public Task<IEncodedContent?> RenderSectionAsync(string name)
+            => _executionScope.RenderSectionAsync(name);
     }
 
     /// <summary>
@@ -402,13 +403,13 @@ public bool IsSectionDefined(string name)
     /// </summary>
     /// <remarks>
     /// StringBuilders can be combined more efficiently than strings, which is useful for layouts.
-    /// <see cref="TextWriter"/> has a dedicated <c>Write</c> overload for <see cref="StringBuilder"/>.
+    /// <see cref="TextWriter"/> has a dedicated <c>Write</c> overload for <see cref="StringBuilder"/> in some frameworks.
     /// </remarks>
-    private class EncodedContent : IEncodedContent
+    private class BufferedContent : IEncodedContent
     {
         public StringBuilder Output { get; }
 
-        public EncodedContent(StringBuilder value)
+        public BufferedContent(StringBuilder value)
             => Output = value;
 
         public void WriteTo(TextWriter textWriter)
diff --git a/src/RazorBlade.Tests/HtmlLayoutTests.cs b/src/RazorBlade.Tests/HtmlLayoutTests.cs
index 44241fb..8aa3f04 100644
--- a/src/RazorBlade.Tests/HtmlLayoutTests.cs
+++ b/src/RazorBlade.Tests/HtmlLayoutTests.cs
@@ -184,16 +184,11 @@ public void should_throw_on_render()
         Assert.Throws<InvalidOperationException>(() => ((RazorTemplate)layout).Render(CancellationToken.None));
     }
 
-    private class Template : HtmlTemplate
+    private class Template(Action<Template> executeAction) : HtmlTemplate
     {
-        private readonly Action<Template> _executeAction;
-
-        public Template(Action<Template> executeAction)
-            => _executeAction = executeAction;
-
         protected internal override Task ExecuteAsync()
         {
-            _executeAction(this);
+            executeAction(this);
             return base.ExecuteAsync();
         }
 
@@ -210,18 +205,13 @@ public void SetSection(string name, string content)
         }
     }
 
-    private class Layout : HtmlLayout
+    private class Layout(Action<Layout> executeAction) : HtmlLayout
     {
-        private readonly Action<Layout> _executeAction;
-
         public bool WasExecuted { get; private set; }
 
-        public Layout(Action<Layout> executeAction)
-            => _executeAction = executeAction;
-
         protected internal override Task ExecuteAsync()
         {
-            _executeAction(this);
+            executeAction(this);
             WasExecuted = true;
             return base.ExecuteAsync();
         }
diff --git a/src/RazorBlade.Tests/HtmlTemplateTests.cs b/src/RazorBlade.Tests/HtmlTemplateTests.cs
index 55c889e..3daa7f5 100644
--- a/src/RazorBlade.Tests/HtmlTemplateTests.cs
+++ b/src/RazorBlade.Tests/HtmlTemplateTests.cs
@@ -103,16 +103,11 @@ public void should_write_raw_string()
         template.Render().ShouldEqual("&<>");
     }
 
-    private class Template : HtmlTemplate
+    private class Template(Action<Template> executeAction) : HtmlTemplate
     {
-        private readonly Action<Template> _executeAction;
-
-        public Template(Action<Template> executeAction)
-            => _executeAction = executeAction;
-
         protected internal override Task ExecuteAsync()
         {
-            _executeAction(this);
+            executeAction(this);
             return base.ExecuteAsync();
         }
     }
diff --git a/src/RazorBlade.Tests/PlainTextTemplateTests.cs b/src/RazorBlade.Tests/PlainTextTemplateTests.cs
index dd8f231..7ce0708 100644
--- a/src/RazorBlade.Tests/PlainTextTemplateTests.cs
+++ b/src/RazorBlade.Tests/PlainTextTemplateTests.cs
@@ -78,16 +78,11 @@ public void should_not_skip_prefixes_of_null_attribute_values()
         template.Render().ShouldEqual("foo=\" a True b False c 42 d  e bar\"");
     }
 
-    private class Template : PlainTextTemplate
+    private class Template(Action<Template> executeAction) : PlainTextTemplate
     {
-        private readonly Action<Template> _executeAction;
-
-        public Template(Action<Template> executeAction)
-            => _executeAction = executeAction;
-
         protected internal override Task ExecuteAsync()
         {
-            _executeAction(this);
+            executeAction(this);
             return base.ExecuteAsync();
         }
     }
diff --git a/src/RazorBlade.Tests/RazorTemplateTests.cs b/src/RazorBlade.Tests/RazorTemplateTests.cs
index 5b54e22..4708119 100644
--- a/src/RazorBlade.Tests/RazorTemplateTests.cs
+++ b/src/RazorBlade.Tests/RazorTemplateTests.cs
@@ -14,54 +14,58 @@ public class RazorTemplateTests
     public async Task should_write_literal()
     {
         var template = new Template(t => t.WriteLiteral("foo & bar < baz > foobar"));
-        await template.ExecuteAsync();
-        template.Output.ToString().ShouldEqual("foo & bar < baz > foobar");
+
+        var result = await template.RenderAsync();
+
+        result.ShouldEqual("foo & bar < baz > foobar");
     }
 
     [Test]
     public void should_render_to_local_output()
     {
         var template = new Template(t => t.WriteLiteral("foo"));
-        template.Output.Write("bar");
+        template.Output.ShouldEqual(TextWriter.Null);
 
         template.Render().ShouldEqual("foo");
-        template.Output.ToString().ShouldEqual("bar");
+
+        template.Output.ShouldEqual(TextWriter.Null);
     }
 
     [Test]
     public void should_render_to_text_writer()
     {
         var template = new Template(t => t.WriteLiteral("foo"));
-        template.Output.Write("bar");
+        template.Output.ShouldEqual(TextWriter.Null);
 
         var output = new StringWriter();
         template.Render(output);
         output.ToString().ShouldEqual("foo");
 
-        template.Output.ToString().ShouldEqual("bar");
+        template.Output.ShouldEqual(TextWriter.Null);
     }
 
     [Test]
     public async Task should_render_async_to_local_output()
     {
         var template = new Template(t => t.WriteLiteral("foo"));
-        await template.Output.WriteAsync("bar");
+        template.Output.ShouldEqual(TextWriter.Null);
 
         (await template.RenderAsync()).ShouldEqual("foo");
-        template.Output.ToString().ShouldEqual("bar");
+
+        template.Output.ShouldEqual(TextWriter.Null);
     }
 
     [Test]
     public async Task should_render_async_to_text_writer()
     {
         var template = new Template(t => t.WriteLiteral("foo"));
-        await template.Output.WriteAsync("bar");
+        template.Output.ShouldEqual(TextWriter.Null);
 
         var output = new StringWriter();
         await template.RenderAsync(output);
         output.ToString().ShouldEqual("foo");
 
-        template.Output.ToString().ShouldEqual("bar");
+        template.Output.ShouldEqual(TextWriter.Null);
     }
 
     [Test]
@@ -130,16 +134,8 @@ public void should_not_execute_template_when_already_cancelled()
         semaphore.CurrentCount.ShouldEqual(0);
     }
 
-    private class Template : RazorTemplate
+    private class Template(Func<Template, Task> executeAction) : RazorTemplate
     {
-        private readonly Func<Template, Task> _executeAction;
-
-        public Template(Func<Template, Task> executeAction)
-        {
-            _executeAction = executeAction;
-            Output = new StringWriter();
-        }
-
         public Template(Action<Template> executeAction)
             : this(t =>
             {
@@ -151,7 +147,7 @@ public Template(Action<Template> executeAction)
 
         protected internal override async Task ExecuteAsync()
         {
-            await _executeAction(this);
+            await executeAction(this);
             await base.ExecuteAsync();
         }
 

From ba94f7a4f77d2f5e3aada1a5dabe44cd60ffa53e Mon Sep 17 00:00:00 2001
From: Lucas Trzesniewski <lucas.trzesniewski@gmail.com>
Date: Fri, 5 Apr 2024 23:30:48 +0200
Subject: [PATCH 3/6] Add tests

---
 src/RazorBlade.Library/RazorTemplate.cs       |  2 +-
 src/RazorBlade.Tests/HtmlLayoutTests.cs       | 47 +++++++++-
 src/RazorBlade.Tests/RazorTemplateTests.cs    | 85 +++++++++++++++++++
 .../Support/AssertExtensions.cs               | 11 +++
 src/RazorBlade.Tests/Support/StepSemaphore.cs | 35 ++++++++
 5 files changed, 175 insertions(+), 5 deletions(-)
 create mode 100644 src/RazorBlade.Tests/Support/StepSemaphore.cs

diff --git a/src/RazorBlade.Library/RazorTemplate.cs b/src/RazorBlade.Library/RazorTemplate.cs
index a2aae76..a7d213e 100644
--- a/src/RazorBlade.Library/RazorTemplate.cs
+++ b/src/RazorBlade.Library/RazorTemplate.cs
@@ -167,7 +167,7 @@ private protected virtual async Task<IRazorExecutionResult> ExecuteAsyncCore(Tex
     /// <remarks>
     /// This feature is not compatible with layouts.
     /// </remarks>
-    protected Task FlushAsync()
+    protected internal Task FlushAsync()
         => _executionScope?.FlushAsync() ?? Task.CompletedTask;
 
     /// <summary>
diff --git a/src/RazorBlade.Tests/HtmlLayoutTests.cs b/src/RazorBlade.Tests/HtmlLayoutTests.cs
index 8aa3f04..d588c4c 100644
--- a/src/RazorBlade.Tests/HtmlLayoutTests.cs
+++ b/src/RazorBlade.Tests/HtmlLayoutTests.cs
@@ -184,12 +184,51 @@ public void should_throw_on_render()
         Assert.Throws<InvalidOperationException>(() => ((RazorTemplate)layout).Render(CancellationToken.None));
     }
 
-    private class Template(Action<Template> executeAction) : HtmlTemplate
+    [Test]
+    public void should_throw_when_setting_layout_after_flush()
     {
-        protected internal override Task ExecuteAsync()
+        var layout = new Layout(_ => { });
+
+        var page = new Template(async t =>
         {
-            executeAction(this);
-            return base.ExecuteAsync();
+            await t.FlushAsync();
+            t.Layout = layout;
+        });
+
+        Assert.Throws<InvalidOperationException>(() => page.Render())
+              .ShouldNotBeNull().Message.ShouldEqual("The layout can no longer be changed.");
+    }
+
+    [Test]
+    public void should_throw_when_flushing_with_layout()
+    {
+        var layout = new Layout(_ => { });
+
+        var page = new Template(async t =>
+        {
+            t.Layout = layout;
+            await t.FlushAsync();
+        });
+
+        Assert.Throws<InvalidOperationException>(() => page.Render())
+              .ShouldNotBeNull().Message.ShouldEqual("The output cannot be flushed when a layout is used.");
+    }
+
+    private class Template(Func<Template, Task> executeAction) : HtmlTemplate
+    {
+        public Template(Action<Template> executeAction)
+            : this(t =>
+            {
+                executeAction(t);
+                return Task.CompletedTask;
+            })
+        {
+        }
+
+        protected internal override async Task ExecuteAsync()
+        {
+            await executeAction(this);
+            await base.ExecuteAsync();
         }
 
         public void SetSection(string name, string content)
diff --git a/src/RazorBlade.Tests/RazorTemplateTests.cs b/src/RazorBlade.Tests/RazorTemplateTests.cs
index 4708119..8ab7c80 100644
--- a/src/RazorBlade.Tests/RazorTemplateTests.cs
+++ b/src/RazorBlade.Tests/RazorTemplateTests.cs
@@ -134,6 +134,91 @@ public void should_not_execute_template_when_already_cancelled()
         semaphore.CurrentCount.ShouldEqual(0);
     }
 
+    [Test]
+    public async Task should_flush_output()
+    {
+        var output = new StringWriter();
+        var stepSemaphore = new StepSemaphore();
+
+        var template = new Template(async t =>
+        {
+            await stepSemaphore.WaitForNextStepAsync();
+
+            t.WriteLiteral("foo");
+            await t.FlushAsync();
+
+            await stepSemaphore.WaitForNextStepAsync();
+
+            t.WriteLiteral(" bar");
+            await t.FlushAsync();
+
+            await stepSemaphore.WaitForNextStepAsync();
+
+            t.WriteLiteral(" baz");
+            await t.FlushAsync();
+
+            await stepSemaphore.NotifyEndOfLastStepAsync();
+        });
+
+        var task = template.RenderAsync(output);
+
+        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        output.ToString().ShouldEqual("foo");
+        template.Output.ShouldBe<StringWriter>().ToString().ShouldEqual(string.Empty);
+
+        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        output.ToString().ShouldEqual("foo bar");
+        template.Output.ShouldBe<StringWriter>().ToString().ShouldEqual(string.Empty);
+
+        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        output.ToString().ShouldEqual("foo bar baz");
+
+        await task;
+        output.ToString().ShouldEqual("foo bar baz");
+        template.Output.ShouldEqual(TextWriter.Null);
+    }
+
+    [Test]
+    public async Task should_buffer_output_until_flushed()
+    {
+        var output = new StringWriter();
+        var stepSemaphore = new StepSemaphore();
+
+        var template = new Template(async t =>
+        {
+            await stepSemaphore.WaitForNextStepAsync();
+
+            t.WriteLiteral("foo");
+            await t.Output.FlushAsync();
+
+            await stepSemaphore.WaitForNextStepAsync();
+
+            t.WriteLiteral(" bar");
+            await t.Output.FlushAsync();
+
+            await stepSemaphore.WaitForNextStepAsync();
+
+            t.WriteLiteral(" baz");
+            await t.Output.FlushAsync();
+
+            await stepSemaphore.NotifyEndOfLastStepAsync();
+        });
+
+        var task = template.RenderAsync(output);
+
+        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        output.ToString().ShouldEqual(string.Empty);
+        template.Output.ShouldBe<StringWriter>().ToString().ShouldEqual("foo");
+
+        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        output.ToString().ShouldEqual(string.Empty);
+        template.Output.ShouldBe<StringWriter>().ToString().ShouldEqual("foo bar");
+
+        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        await task;
+        output.ToString().ShouldEqual("foo bar baz");
+    }
+
     private class Template(Func<Template, Task> executeAction) : RazorTemplate
     {
         public Template(Action<Template> executeAction)
diff --git a/src/RazorBlade.Tests/Support/AssertExtensions.cs b/src/RazorBlade.Tests/Support/AssertExtensions.cs
index 1d0da29..cba304e 100644
--- a/src/RazorBlade.Tests/Support/AssertExtensions.cs
+++ b/src/RazorBlade.Tests/Support/AssertExtensions.cs
@@ -1,10 +1,14 @@
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using JetBrains.Annotations;
 using NUnit.Framework;
 
 namespace RazorBlade.Tests.Support;
 
+#if NET
+[StackTraceHidden]
+#endif
 internal static class AssertExtensions
 {
     public static void ShouldEqual<T>(this T? actual, T? expected)
@@ -19,6 +23,13 @@ public static void ShouldBeTrue(this bool actual)
     public static void ShouldBeFalse(this bool actual)
         => Assert.That(actual, Is.False);
 
+    public static TExpected ShouldBe<TExpected>(this object? actual)
+        where TExpected : class
+    {
+        Assert.That(actual, Is.InstanceOf<TExpected>());
+        return actual as TExpected ?? throw new AssertionException($"Expected instance of {typeof(TExpected).Name}");
+    }
+
     [ContractAnnotation("notnull => halt")]
     public static void ShouldBeNull(this object? actual)
         => Assert.That(actual, Is.Null);
diff --git a/src/RazorBlade.Tests/Support/StepSemaphore.cs b/src/RazorBlade.Tests/Support/StepSemaphore.cs
new file mode 100644
index 0000000..72d84f5
--- /dev/null
+++ b/src/RazorBlade.Tests/Support/StepSemaphore.cs
@@ -0,0 +1,35 @@
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace RazorBlade.Tests.Support;
+
+public class StepSemaphore
+{
+    private readonly SemaphoreSlim _semaphoreStartOfStep = new(0);
+    private readonly SemaphoreSlim _semaphoreEndOfStep = new(0);
+    private int _isFirstStep;
+
+    public async Task WaitForNextStepAsync()
+    {
+        if (Interlocked.Exchange(ref _isFirstStep, 1) != 0)
+            _semaphoreEndOfStep.Release();
+
+        var result = await _semaphoreStartOfStep.WaitAsync(10_000);
+        result.ShouldBeTrue();
+    }
+
+    public Task NotifyEndOfLastStepAsync()
+    {
+        _semaphoreEndOfStep.Release();
+
+        return Task.CompletedTask; // Just to make the API consistent
+    }
+
+    public async Task StartNextStepAndWaitForResultAsync()
+    {
+        _semaphoreStartOfStep.Release();
+
+        var result = await _semaphoreEndOfStep.WaitAsync(10_000);
+        result.ShouldBeTrue();
+    }
+}

From 05fb6171ae61cd5797a446ad830e363afaf576d5 Mon Sep 17 00:00:00 2001
From: Lucas Trzesniewski <lucas.trzesniewski@gmail.com>
Date: Sat, 6 Apr 2024 15:57:02 +0200
Subject: [PATCH 4/6] Refactor tests

---
 src/RazorBlade.Tests/RazorTemplateTests.cs    | 40 +++++------
 src/RazorBlade.Tests/Support/StepSemaphore.cs | 70 +++++++++++++++----
 2 files changed, 76 insertions(+), 34 deletions(-)

diff --git a/src/RazorBlade.Tests/RazorTemplateTests.cs b/src/RazorBlade.Tests/RazorTemplateTests.cs
index 8ab7c80..21faa10 100644
--- a/src/RazorBlade.Tests/RazorTemplateTests.cs
+++ b/src/RazorBlade.Tests/RazorTemplateTests.cs
@@ -123,15 +123,15 @@ public void should_not_execute_template_when_already_cancelled()
         var cts = new CancellationTokenSource();
         cts.Cancel();
 
-        var semaphore = new SemaphoreSlim(0);
-        var template = new Template(_ => semaphore.Release());
+        var executionCount = 0;
+        var template = new Template(_ => Interlocked.Increment(ref executionCount));
 
         Assert.Catch<OperationCanceledException>(() => template.Render(cts.Token));
         Assert.Catch<OperationCanceledException>(() => template.Render(StreamWriter.Null, cts.Token));
         Assert.CatchAsync<OperationCanceledException>(() => template.RenderAsync(cts.Token));
         Assert.CatchAsync<OperationCanceledException>(() => template.RenderAsync(StreamWriter.Null, cts.Token));
 
-        semaphore.CurrentCount.ShouldEqual(0);
+        executionCount.ShouldEqual(0);
     }
 
     [Test]
@@ -142,37 +142,35 @@ public async Task should_flush_output()
 
         var template = new Template(async t =>
         {
-            await stepSemaphore.WaitForNextStepAsync();
+            using var worker = stepSemaphore.CreateWorker();
+            await worker.WaitForNextStepAsync();
 
             t.WriteLiteral("foo");
             await t.FlushAsync();
 
-            await stepSemaphore.WaitForNextStepAsync();
+            await worker.WaitForNextStepAsync();
 
             t.WriteLiteral(" bar");
             await t.FlushAsync();
 
-            await stepSemaphore.WaitForNextStepAsync();
+            await worker.WaitForNextStepAsync();
 
             t.WriteLiteral(" baz");
             await t.FlushAsync();
-
-            await stepSemaphore.NotifyEndOfLastStepAsync();
         });
 
         var task = template.RenderAsync(output);
+        var controller = stepSemaphore.CreateController();
 
-        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        await controller.StartNextStepAndWaitForResultAsync();
         output.ToString().ShouldEqual("foo");
         template.Output.ShouldBe<StringWriter>().ToString().ShouldEqual(string.Empty);
 
-        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        await controller.StartNextStepAndWaitForResultAsync();
         output.ToString().ShouldEqual("foo bar");
         template.Output.ShouldBe<StringWriter>().ToString().ShouldEqual(string.Empty);
 
-        await stepSemaphore.StartNextStepAndWaitForResultAsync();
-        output.ToString().ShouldEqual("foo bar baz");
-
+        await controller.StartNextStepAndWaitForResultAsync();
         await task;
         output.ToString().ShouldEqual("foo bar baz");
         template.Output.ShouldEqual(TextWriter.Null);
@@ -186,35 +184,35 @@ public async Task should_buffer_output_until_flushed()
 
         var template = new Template(async t =>
         {
-            await stepSemaphore.WaitForNextStepAsync();
+            using var worker = stepSemaphore.CreateWorker();
+            await worker.WaitForNextStepAsync();
 
             t.WriteLiteral("foo");
             await t.Output.FlushAsync();
 
-            await stepSemaphore.WaitForNextStepAsync();
+            await worker.WaitForNextStepAsync();
 
             t.WriteLiteral(" bar");
             await t.Output.FlushAsync();
 
-            await stepSemaphore.WaitForNextStepAsync();
+            await worker.WaitForNextStepAsync();
 
             t.WriteLiteral(" baz");
             await t.Output.FlushAsync();
-
-            await stepSemaphore.NotifyEndOfLastStepAsync();
         });
 
         var task = template.RenderAsync(output);
+        var controller = stepSemaphore.CreateController();
 
-        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        await controller.StartNextStepAndWaitForResultAsync();
         output.ToString().ShouldEqual(string.Empty);
         template.Output.ShouldBe<StringWriter>().ToString().ShouldEqual("foo");
 
-        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        await controller.StartNextStepAndWaitForResultAsync();
         output.ToString().ShouldEqual(string.Empty);
         template.Output.ShouldBe<StringWriter>().ToString().ShouldEqual("foo bar");
 
-        await stepSemaphore.StartNextStepAndWaitForResultAsync();
+        await controller.StartNextStepAndWaitForResultAsync();
         await task;
         output.ToString().ShouldEqual("foo bar baz");
     }
diff --git a/src/RazorBlade.Tests/Support/StepSemaphore.cs b/src/RazorBlade.Tests/Support/StepSemaphore.cs
index 72d84f5..4aacebb 100644
--- a/src/RazorBlade.Tests/Support/StepSemaphore.cs
+++ b/src/RazorBlade.Tests/Support/StepSemaphore.cs
@@ -1,4 +1,5 @@
-using System.Threading;
+using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace RazorBlade.Tests.Support;
@@ -7,29 +8,72 @@ public class StepSemaphore
 {
     private readonly SemaphoreSlim _semaphoreStartOfStep = new(0);
     private readonly SemaphoreSlim _semaphoreEndOfStep = new(0);
-    private int _isFirstStep;
 
-    public async Task WaitForNextStepAsync()
+    private bool _hasWorker;
+    private bool _hasController;
+    private bool _isFirstStep = true;
+
+    public IWorker CreateWorker()
     {
-        if (Interlocked.Exchange(ref _isFirstStep, 1) != 0)
-            _semaphoreEndOfStep.Release();
+        if (_hasWorker)
+            throw new InvalidOperationException("A worker has already been created");
 
-        var result = await _semaphoreStartOfStep.WaitAsync(10_000);
-        result.ShouldBeTrue();
+        _hasWorker = true;
+        return new Worker(this);
     }
 
-    public Task NotifyEndOfLastStepAsync()
+    public IController CreateController()
     {
-        _semaphoreEndOfStep.Release();
+        if (_hasController)
+            throw new InvalidOperationException("A controller has already been created");
 
-        return Task.CompletedTask; // Just to make the API consistent
+        _hasController = true;
+        return new Controller(this);
     }
 
-    public async Task StartNextStepAndWaitForResultAsync()
+    private async Task WaitForNextStepAsync()
+    {
+        if (!_isFirstStep)
+            _semaphoreEndOfStep.Release();
+
+        _isFirstStep = false;
+        await WaitWithTimeout(_semaphoreStartOfStep);
+    }
+
+    private void NotifyEndOfLastStep()
+        => _semaphoreEndOfStep.Release();
+
+    private async Task StartNextStepAndWaitForResultAsync()
     {
         _semaphoreStartOfStep.Release();
+        await WaitWithTimeout(_semaphoreEndOfStep);
+    }
 
-        var result = await _semaphoreEndOfStep.WaitAsync(10_000);
-        result.ShouldBeTrue();
+    private static async Task WaitWithTimeout(SemaphoreSlim semaphore)
+        => (await semaphore.WaitAsync(10_000)).ShouldBeTrue();
+
+    public interface IWorker : IDisposable
+    {
+        Task WaitForNextStepAsync();
+    }
+
+    public interface IController
+    {
+        Task StartNextStepAndWaitForResultAsync();
+    }
+
+    private class Worker(StepSemaphore parent) : IWorker
+    {
+        public Task WaitForNextStepAsync()
+            => parent.WaitForNextStepAsync();
+
+        public void Dispose()
+            => parent.NotifyEndOfLastStep();
+    }
+
+    private class Controller(StepSemaphore parent) : IController
+    {
+        public Task StartNextStepAndWaitForResultAsync()
+            => parent.StartNextStepAndWaitForResultAsync();
     }
 }

From 8ff647d2cbf50b751d396e4b76d32486f73d4109 Mon Sep 17 00:00:00 2001
From: Lucas Trzesniewski <lucas.trzesniewski@gmail.com>
Date: Sat, 6 Apr 2024 19:16:56 +0200
Subject: [PATCH 5/6] Add documentation about FlushAsync

---
 README.md                                     | 19 +++++++++++++++++++
 .../Examples/TemplateWithFlush.cshtml         |  6 ++++++
 .../PageWithFlush.cshtml                      |  7 +++++++
 src/RazorBlade.IntegrationTest/Program.cs     |  1 +
 src/RazorBlade.Library/RazorTemplate.cs       | 13 +++++++++++--
 5 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 src/RazorBlade.IntegrationTest/Examples/TemplateWithFlush.cshtml
 create mode 100644 src/RazorBlade.IntegrationTest/PageWithFlush.cshtml

diff --git a/README.md b/README.md
index 84c6688..a3e835e 100644
--- a/README.md
+++ b/README.md
@@ -196,6 +196,25 @@ The `RazorTemplate` base class provides `Render` and `RenderAsync` methods to ex
 
 Templates are stateful and not thread-safe, so it is advised to always create new instances of the templates to render.
 
+### Flushing partial output
+
+By default, the output of a template is buffered while it is executing, then copied to the provided writer when finished. This is necessary for features such as layouts to be supported, but may not always be desired.
+
+The `RazorTemplate` class provides a `FlushAsync` method which will copy the buffered output to the provided `TextWriter` and then flush the writer:
+
+<!-- snippet: TemplateWithFlush.Usage -->
+<a id='snippet-templatewithflush.usage'></a>
+```cshtml
+<div>Lightweight content goes here.</div>
+@await FlushAsync()
+<div>Slower to render content goes here.</div>
+```
+<sup><a href='/src/RazorBlade.IntegrationTest/Examples/TemplateWithFlush.cshtml#L2-L6' title='Snippet source file'>snippet source</a> | <a href='#snippet-templatewithflush.usage' title='Start of snippet'>anchor</a></sup>
+<!-- endSnippet -->
+
+> [!IMPORTANT]  
+> Flushing is not compatible with layout usage.
+
 ### MSBuild
 
 The source generator will process `RazorBlade` MSBuild items which have the `.cshtml` file extension.
diff --git a/src/RazorBlade.IntegrationTest/Examples/TemplateWithFlush.cshtml b/src/RazorBlade.IntegrationTest/Examples/TemplateWithFlush.cshtml
new file mode 100644
index 0000000..f339dc4
--- /dev/null
+++ b/src/RazorBlade.IntegrationTest/Examples/TemplateWithFlush.cshtml
@@ -0,0 +1,6 @@
+@inherits RazorBlade.HtmlTemplate
+@* begin-snippet: TemplateWithFlush.Usage *@
+<div>Lightweight content goes here.</div>
+@await FlushAsync()
+<div>Slower to render content goes here.</div>
+@* end-snippet *@
diff --git a/src/RazorBlade.IntegrationTest/PageWithFlush.cshtml b/src/RazorBlade.IntegrationTest/PageWithFlush.cshtml
new file mode 100644
index 0000000..a3eb88e
--- /dev/null
+++ b/src/RazorBlade.IntegrationTest/PageWithFlush.cshtml
@@ -0,0 +1,7 @@
+@inherits RazorBlade.HtmlTemplate
+<h2>Hello, world!</h2>
+This is the body contents before flushing.
+@await FlushAsync()
+This is the body contents after flushing.
+@await FlushAsync()
+This is the body contents after a second flushing.
diff --git a/src/RazorBlade.IntegrationTest/Program.cs b/src/RazorBlade.IntegrationTest/Program.cs
index a1819fb..69de8c2 100644
--- a/src/RazorBlade.IntegrationTest/Program.cs
+++ b/src/RazorBlade.IntegrationTest/Program.cs
@@ -10,6 +10,7 @@ public static void Main()
         WriteTemplate(new TestTemplate { Name = "World" });
         WriteTemplate(new TestTemplateWithModel(new FooBarModelClass { Foo = "Foo", Bar = "Bar" }));
         WriteTemplate(new PageWithLayout());
+        WriteTemplate(new PageWithFlush());
     }
 
     private static void WriteTemplate(RazorTemplate template)
diff --git a/src/RazorBlade.Library/RazorTemplate.cs b/src/RazorBlade.Library/RazorTemplate.cs
index a7d213e..8462737 100644
--- a/src/RazorBlade.Library/RazorTemplate.cs
+++ b/src/RazorBlade.Library/RazorTemplate.cs
@@ -164,11 +164,20 @@ private protected virtual async Task<IRazorExecutionResult> ExecuteAsyncCore(Tex
     /// <summary>
     /// Writes the buffered output to the target output then flushes the output stream.
     /// </summary>
+    /// <returns>
+    /// An empty <see cref="IEncodedContent"/>, which allows using a direct expression: <c>@await FlushAsync()</c>
+    /// </returns>
     /// <remarks>
     /// This feature is not compatible with layouts.
     /// </remarks>
-    protected internal Task FlushAsync()
-        => _executionScope?.FlushAsync() ?? Task.CompletedTask;
+    protected internal async Task<IEncodedContent> FlushAsync()
+    {
+        if (_executionScope is not { } executionScope)
+            throw new InvalidOperationException("The template is not executing.");
+
+        await executionScope.FlushAsync().ConfigureAwait(false);
+        return HtmlString.Empty;
+    }
 
     /// <summary>
     /// Executes the template and appends the result to <see cref="Output"/>.

From 692d6dea291005379333d54a041f37b707210b80 Mon Sep 17 00:00:00 2001
From: Lucas Trzesniewski <lucas.trzesniewski@gmail.com>
Date: Sat, 6 Apr 2024 20:05:37 +0200
Subject: [PATCH 6/6] Refactor rendering to avoid unnecessary copying

---
 src/RazorBlade.Library/RazorTemplate.cs | 82 ++++++++++++-------------
 1 file changed, 38 insertions(+), 44 deletions(-)

diff --git a/src/RazorBlade.Library/RazorTemplate.cs b/src/RazorBlade.Library/RazorTemplate.cs
index 8462737..9a78915 100644
--- a/src/RazorBlade.Library/RazorTemplate.cs
+++ b/src/RazorBlade.Library/RazorTemplate.cs
@@ -48,23 +48,12 @@ public string Render(CancellationToken cancellationToken = default)
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        var textWriter = new StringWriter();
-        var renderTask = RenderAsyncCore(textWriter, cancellationToken);
+        var renderTask = RenderAsync(cancellationToken);
 
         if (renderTask.IsCompleted)
-        {
-            renderTask.GetAwaiter().GetResult();
-            return textWriter.ToString();
-        }
+            return renderTask.GetAwaiter().GetResult();
 
-        return Task.Run(
-            async () =>
-            {
-                await renderTask.ConfigureAwait(false);
-                return textWriter.ToString();
-            },
-            CancellationToken.None
-        ).GetAwaiter().GetResult();
+        return Task.Run(async () => await renderTask.ConfigureAwait(false), CancellationToken.None).GetAwaiter().GetResult();
     }
 
     /// <summary>
@@ -80,7 +69,7 @@ public void Render(TextWriter textWriter, CancellationToken cancellationToken =
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        var renderTask = RenderAsyncCore(textWriter, cancellationToken);
+        var renderTask = RenderAsync(textWriter, cancellationToken);
 
         if (renderTask.IsCompleted)
         {
@@ -100,11 +89,18 @@ public void Render(TextWriter textWriter, CancellationToken cancellationToken =
     /// </remarks>
     public async Task<string> RenderAsync(CancellationToken cancellationToken = default)
     {
-        cancellationToken.ThrowIfCancellationRequested();
+        var body = await RenderAsyncCore(null, cancellationToken).ConfigureAwait(false);
+
+        switch (body)
+        {
+            case BufferedContent { Output: var bufferedOutput }:
+                return bufferedOutput.ToString();
 
-        var textWriter = new StringWriter();
-        await RenderAsyncCore(textWriter, cancellationToken).ConfigureAwait(false);
-        return textWriter.ToString();
+            default: // Fallback case, shouldn't happen
+                var writer = new StringWriter();
+                body.WriteTo(writer);
+                return writer.ToString();
+        }
     }
 
     /// <summary>
@@ -117,19 +113,26 @@ public async Task<string> RenderAsync(CancellationToken cancellationToken = defa
     /// </remarks>
     public async Task RenderAsync(TextWriter textWriter, CancellationToken cancellationToken = default)
     {
-        cancellationToken.ThrowIfCancellationRequested();
+        var body = await RenderAsyncCore(textWriter, cancellationToken).ConfigureAwait(false);
+
+        switch (body)
+        {
+            case BufferedContent { Output: var bufferedOutput }:
+                await WriteStringBuilderToOutputAsync(bufferedOutput, textWriter, cancellationToken).ConfigureAwait(false);
+                break;
 
-        await RenderAsyncCore(textWriter, cancellationToken).ConfigureAwait(false);
+            default: // Fallback case, shouldn't happen
+                body.WriteTo(textWriter);
+                break;
+        }
     }
 
     /// <summary>
-    /// Renders the template asynchronously including its layout and returns the result as a <see cref="StringBuilder"/>.
+    /// Renders the template and its layout stack.
     /// </summary>
-    private async Task RenderAsyncCore(TextWriter targetOutput, CancellationToken cancellationToken)
+    private async Task<IEncodedContent> RenderAsyncCore(TextWriter? targetOutput, CancellationToken cancellationToken)
     {
-        cancellationToken.ThrowIfCancellationRequested();
-
-        var executionResult = await ExecuteAsyncCore(targetOutput, cancellationToken);
+        var executionResult = await ExecuteAsyncCore(targetOutput, cancellationToken).ConfigureAwait(false);
 
         while (executionResult.Layout is { } layout)
         {
@@ -137,16 +140,7 @@ private async Task RenderAsyncCore(TextWriter targetOutput, CancellationToken ca
             executionResult = await layout.ExecuteLayoutAsync(executionResult).ConfigureAwait(false);
         }
 
-        switch (executionResult.Body)
-        {
-            case BufferedContent { Output: var bufferedOutput }:
-                await WriteStringBuilderToOutputAsync(bufferedOutput, targetOutput, cancellationToken).ConfigureAwait(false);
-                break;
-
-            case { } body: // Fallback case, shouldn't happen
-                body.WriteTo(targetOutput);
-                break;
-        }
+        return executionResult.Body;
     }
 
     /// <summary>
@@ -337,18 +331,18 @@ public async Task FlushAsync()
             // so disallow setting a layout later on, as that would lead to inconsistent results.
             _layoutFrozen = true;
 
-            if (_targetOutput is not null)
-            {
-                var bufferedOutput = BufferedOutput.GetStringBuilder();
-                await WriteStringBuilderToOutputAsync(bufferedOutput, _targetOutput, CancellationToken).ConfigureAwait(false);
-                bufferedOutput.Clear();
+            if (_targetOutput is null)
+                return;
+
+            var bufferedOutput = BufferedOutput.GetStringBuilder();
+            await WriteStringBuilderToOutputAsync(bufferedOutput, _targetOutput, CancellationToken).ConfigureAwait(false);
+            bufferedOutput.Clear();
 
 #if NET8_0_OR_GREATER
-                await _targetOutput.FlushAsync(CancellationToken).ConfigureAwait(false);
+            await _targetOutput.FlushAsync(CancellationToken).ConfigureAwait(false);
 #else
-                await _targetOutput.FlushAsync().ConfigureAwait(false);
+            await _targetOutput.FlushAsync().ConfigureAwait(false);
 #endif
-            }
         }
 
         public BufferedContent ToBufferedContent()