PNG decoder: buffer allocated in ReadChunkData is not disposed when the read is cancelled #3202

// Minimal reproduction for SixLabors.ImageSharp 4.1.2: the buffer rented in
// PngDecoderCore.ReadChunkData is never disposed when the stream read that
// follows the allocation observes cancellation.
//
// The cancellation is made deterministic by cancelling from inside that
// allocation. The internal MemoryDiagnostics.MemoryAllocated event is observed
// through reflection (diagnostics only; nothing in ImageSharp is patched). On
// each allocation the stack is checked for PngDecoderCore.ReadChunkData inside
// the intended phase (Identify or Decode), and the token is cancelled once. The
// next statement in ReadChunkData is the BufferedReadStream read, which throws.
//
// Usage, one case per process:
//   dotnet run -c Release -- A    no cancellation
//   dotnet run -c Release -- B    token cancelled before the operation starts
//   dotnet run -c Release -- C    cancelled at the ReadChunkData allocation in Identify
//   dotnet run -c Release -- D    cancelled at the ReadChunkData allocation in Decode

using System.Collections.Concurrent;
using System.Diagnostics;
using System.Reflection;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Text.RegularExpressions;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Diagnostics;
using SixLabors.ImageSharp.Formats;
using SixLabors.ImageSharp.Formats.Png;
using SixLabors.ImageSharp.Memory;
using SixLabors.ImageSharp.PixelFormats;

var caseId = args.Length == 1 ? args[0].ToUpperInvariant() : string.Empty;
if (caseId is not ("A" or "B" or "C" or "D"))
{
    Console.Error.WriteLine("usage: dotnet run -c Release -- A|B|C|D");
    return 2;
}

Out("case", caseId switch
{
    "A" => "A - no cancellation",
    "B" => "B - token cancelled before the operation starts",
    "C" => "C - cancelled at the ReadChunkData allocation in PNG Identify",
    _ => "D - cancelled at the ReadChunkData allocation in PNG Decode"
});
Out("runtime", RuntimeInformation.FrameworkDescription);
Out("os", RuntimeInformation.OSDescription);
Out("arch", RuntimeInformation.ProcessArchitecture);
Out("imagesharp", typeof(Image).Assembly.GetCustomAttribute<AssemblyInformationalVersionAttribute>()?.InformationalVersion);

var memoryAllocated = typeof(MemoryDiagnostics).GetEvent(
    "MemoryAllocated",
    BindingFlags.NonPublic | BindingFlags.Static);
if (memoryAllocated is null)
{
    Out("result", "cannot run: internal MemoryDiagnostics.MemoryAllocated not found");
    return 3;
}

// Input: a plain 2048x1536 RGBA PNG generated here; no external file.
var png = CreatePng(2048, 1536);
Out("input", $"generated PNG 2048x1536, {png.Length} bytes");

// Settle before taking the baseline.
GC.Collect();
GC.WaitForPendingFinalizers();
GC.Collect();

// The library's own finalizer-based report of undisposed allocations. Subscribing
// also makes it record an allocation stack for each allocation from here on.
var undisposedReports = new ConcurrentQueue<string>();
MemoryDiagnostics.UndisposedAllocation += stack => undisposedReports.Enqueue(stack);

var baseline = MemoryDiagnostics.TotalUndisposedAllocationCount;
Out("baseline.TotalUndisposedAllocationCount", baseline);

using var cancellation = new CancellationTokenSource();
var targetPhase = caseId switch { "C" => "Identify", "D" => "Decode", _ => null };
var otherPhase = caseId switch { "C" => "Decode", "D" => "Identify", _ => null };
var fired = 0;
string? hookStack = null;

var onAllocated = new Action(() =>
{
    if (targetPhase is null || Volatile.Read(ref fired) != 0)
    {
        return;
    }

    var frames = new StackTrace(fNeedFileInfo: false).GetFrames();
    bool inReadChunkData = false, inTarget = false, inOther = false;

    foreach (var frame in frames)
    {
        var method = frame.GetMethod();
        if (method?.DeclaringType?.Name != "PngDecoderCore")
        {
            continue;
        }

        if (method.Name == "ReadChunkData")
        {
            inReadChunkData = true;
        }
        else if (method.Name == targetPhase)
        {
            inTarget = true;
        }
        else if (method.Name == otherPhase)
        {
            inOther = true;
        }
    }

    if (!inReadChunkData || !inTarget || inOther || Interlocked.CompareExchange(ref fired, 1, 0) != 0)
    {
        return;
    }

    hookStack = Frames(frames.Skip(1), 10);

    // The allocation has happened; the read that follows it has not.
    cancellation.Cancel();
});

// The stack at the point where the cancellation is thrown (an awaited
// exception no longer carries it).
string? throwSite = null;
AppDomain.CurrentDomain.FirstChanceException += (_, e) =>
{
    if (e.Exception is OperationCanceledException && throwSite is null)
    {
        // From the token's throw helper when there is one; otherwise from the
        // first frame after this handler and the runtime's dispatch.
        var frames = new StackTrace(fNeedFileInfo: false).GetFrames().Skip(1).ToArray();
        var fromToken = frames.SkipWhile(f =>
            f.GetMethod()?.DeclaringType?.FullName?.StartsWith("System.Threading.CancellationToken", StringComparison.Ordinal) != true).ToArray();
        throwSite = Frames(
            fromToken.Length > 0
                ? fromToken
                : frames.SkipWhile(f => f.GetMethod()?.DeclaringType?.FullName?.StartsWith("System.Runtime.EH", StringComparison.Ordinal) == true),
            8);
    }
};

if (targetPhase is not null)
{
    memoryAllocated.GetAddMethod(nonPublic: true)!.Invoke(null, [onAllocated]);
}

if (caseId == "B")
{
    cancellation.Cancel();
}

var (outcome, awaitStack, immediate) = await RunAsync(png, cancellation.Token);

if (targetPhase is not null)
{
    memoryAllocated.GetRemoveMethod(nonPublic: true)!.Invoke(null, [onAllocated]);
}

Out("trigger", targetPhase is null
    ? "n/a"
    : fired == 1
        ? $"fired once, at the ReadChunkData allocation inside PngDecoderCore.{targetPhase}"
        : "DID NOT FIRE - the intended allocation was not identified, so this run shows nothing");
if (hookStack is not null)
{
    Console.WriteLine("trigger.allocationStack=" + Environment.NewLine + hookStack);
}

Out("outcome", outcome);
if (awaitStack is not null)
{
    Console.WriteLine("outcome.awaitStack=" + Environment.NewLine + awaitStack);
}

Console.WriteLine("cancellation.throwSite=" + (throwSite is null ? "none" : Environment.NewLine + throwSite));
Out("immediate.TotalUndisposedAllocationCount", immediate);
Out("immediate.delta", immediate - baseline);

// Diagnostic only, after the measurement above: once the configuration that
// owns the allocator is unreachable, lets the library's finalizer report any
// allocation that was never disposed.
GC.Collect();
GC.WaitForPendingFinalizers();
GC.Collect();
GC.WaitForPendingFinalizers();
Out("afterGc.TotalUndisposedAllocationCount.delta", MemoryDiagnostics.TotalUndisposedAllocationCount - baseline);
Out("undisposedAllocationReports", undisposedReports.Count);

var index = 0;
foreach (var report in undisposedReports)
{
    index++;
    var lines = report
        .Split('\n')
        .Select(line => Regex.Replace(line.TrimEnd('\r'), @" in .+?:line \d+$", string.Empty))
        .Where(line => line.Contains(" at ", StringComparison.Ordinal))
        .Take(9);
    Console.WriteLine($"undisposedAllocationReport[{index}].allocationStack=" + Environment.NewLine + string.Join(Environment.NewLine, lines));
}

return 0;

// Identify, then decode, with a dedicated configuration that is unreachable once
// this returns. The undisposed count is read here, before any GC.
[MethodImpl(MethodImplOptions.NoInlining)]
static async Task<(string Outcome, string? AwaitStack, int Immediate)> RunAsync(byte[] png, CancellationToken token)
{
    var configuration = new Configuration(new PngConfigurationModule())
    {
        MemoryAllocator = MemoryAllocator.Create()
    };
    var options = new DecoderOptions { Configuration = configuration };

    string outcome;
    string? awaitStack = null;

    try
    {
        using var stream = new MemoryStream(png, writable: false);
        var info = await Image.IdentifyAsync(options, stream, token);

        stream.Position = 0;
        using var image = await Image.LoadAsync<Rgba32>(options, stream, token);

        outcome = $"completed: identified {info.Width}x{info.Height}, decoded {image.Width}x{image.Height}";
    }
    catch (Exception e)
    {
        outcome = $"{e.GetType().FullName}: {e.Message}";
        awaitStack = Frames(new StackTrace(e, fNeedFileInfo: false).GetFrames(), 6);
    }

    var immediate = MemoryDiagnostics.TotalUndisposedAllocationCount;
    GC.KeepAlive(configuration);

    return (outcome, awaitStack, immediate);
}

static byte[] CreatePng(int width, int height)
{
    using var image = new Image<Rgba32>(width, height);
    image.ProcessPixelRows(accessor =>
    {
        for (var y = 0; y < accessor.Height; y++)
        {
            var row = accessor.GetRowSpan(y);
            for (var x = 0; x < row.Length; x++)
            {
                row[x] = new Rgba32((byte)(x * 255 / width), (byte)(y * 255 / height), (byte)((x + y) & 0xFF), 255);
            }
        }
    });

    using var output = new MemoryStream();
    image.SaveAsPng(output);
    return output.ToArray();
}

static string Frames(IEnumerable<StackFrame> frames, int take)
{
    return string.Join(Environment.NewLine, frames.Take(take).Select(frame =>
    {
        var method = frame.GetMethod();
        return $"    at {method?.DeclaringType?.FullName}.{method?.Name}";
    }));
}

static void Out(string key, object? value) => Console.WriteLine($"{key}={value}");