Search code examples
c#asynchronousmemory-management

Avoiding Span<T>.ToArray() in an async method?


Upgraded some code to be asynchronous and realized that Span<T> cannot be used in an async method.

Invoking Span<T>.ToArray() fixes the issue at the cost of allocating every single time.

Can you suggest a pattern on how to avoid allocating in this case?

using ISO9660.Logical;

namespace ISO9660.Physical;

public static class DiscExtensions
{
    public static async Task ReadFileRawAsync(this Disc disc, IsoFileSystemEntryFile file, Stream stream)
    {
        await ReadFileAsync(disc, file, stream, ReadFileRaw);
    }

    public static async Task ReadFileUserAsync(this Disc disc, IsoFileSystemEntryFile file, Stream stream)
    {
        await ReadFileAsync(disc, file, stream, ReadFileUser);
    }

    private static async Task ReadFileAsync(Disc disc, IsoFileSystemEntryFile file, Stream stream, ReadFileHandler handler)
    {
        int position = (int)file.Position;

        Track track = disc.Tracks.FirstOrDefault(s => position >= s.Position)
                      ?? throw new InvalidOperationException("Failed to determine track for file.");

        int sectors = (int)Math.Ceiling((double)file.Length / track.Sector.GetUserDataLength());

        for (int i = position; i < position + sectors; i++)
        {
            ISector sector = await track.ReadSectorAsync(i);

            byte[] array = handler(file, stream, sector).ToArray(); // sucks

            await stream.WriteAsync(array);
        }
    }

    private static Span<byte> ReadFileRaw(IsoFileSystemEntryFile file, Stream stream, ISector sector)
    {
        Span<byte> data = sector.GetData();

        int size = data.Length;

        Span<byte> span = data[..size];

        return span;
    }

    private static Span<byte> ReadFileUser(IsoFileSystemEntryFile file, Stream stream, ISector sector)
    {
        Span<byte> data = sector.GetUserData();

        int size = (int)Math.Min(Math.Max(file.Length - stream.Length, 0), data.Length);

        Span<byte> span = data[..size];

        return span;
    }

    private delegate Span<byte> ReadFileHandler(IsoFileSystemEntryFile file, Stream stream, ISector sector);
}

ISector interface:

using System;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Runtime.InteropServices.Marshalling;

namespace ISO9660.Physical;

public interface ISector
{
    int Length { get; }

    Span<byte> GetData();

    Span<byte> GetUserData();

    int GetUserDataLength();

    public static Span<byte> GetSpan<T>(scoped ref T sector, int start, int length)
        where T : struct, ISector
    {
        var span = MemoryMarshal.CreateSpan(ref sector, 1);

        var bytes = MemoryMarshal.AsBytes(span);

        var slice = bytes.Slice(start, length);

        return slice;
    }
    
    public static ISector Read<T>(Stream stream)
        where T : struct, ISector
    {
        var size = Unsafe.SizeOf<T>();

        Span<byte> span = stackalloc byte[size];

        stream.ReadExactly(span);

        var read = MemoryMarshal.Read<T>(span);

        return read;
    }

    public static async Task<ISector> ReadAsync<T>(Stream stream)
        where T : struct, ISector
    {
        var buffer = new byte[Unsafe.SizeOf<T>()];

        await stream.ReadExactlyAsync(buffer);

        var sector = MemoryMarshal.Read<T>(buffer);

        return sector;
    }
}

ISector implementation example:

using JetBrains.Annotations;

namespace ISO9660.Physical;

public unsafe struct SectorRawMode1 : ISector
{
    private const int UserDataLength = 2048;

    private const int UserDataPosition = 16;

    [UsedImplicitly]
    public fixed byte Sync[12];

    [UsedImplicitly]
    public fixed byte Header[4];

    [UsedImplicitly]
    public fixed byte UserData[UserDataLength];

    [UsedImplicitly]
    public fixed byte Edc[4];

    [UsedImplicitly]
    public fixed byte Intermediate[8];

    [UsedImplicitly]
    public fixed byte PParity[172];

    [UsedImplicitly]
    public fixed byte QParity[104];

    public readonly int Length => 2352;

    public Span<byte> GetData()
    {
        return ISector.GetSpan(ref this, 0, Length);
    }

    public Span<byte> GetUserData()
    {
        return ISector.GetSpan(ref this, UserDataPosition, UserDataLength);
    }

    public readonly int GetUserDataLength()
    {
        return UserDataLength;
    }
}

Solution

  • Change Span<T> to Memory<T> and push the change upwards. Memory<T> is analogous to Span<T> for async (in practice). The idea is unlike a Span, a Memory can live on the heap (and thus can't refer to a stackalloc).

    There is no MemoryMarshal.CreateMemory so you have to do it the old fashioned way: pin, take address, and construct a Memory form a pointer. Example how to do this: https://stackoverflow.com/a/52191681/14768 Note that the MemoryManager allocation is far smaller than the buffer allocation and you don't need to copy.

    Be very careful how you write your code; this MemoryManager's constructor does indeed take a Span that is the entire ISector, and it will live across async calls; but if you fail to dispose it, memory leak. Also, if the ISector came from the stack rather than the heap, that's on you and you will trash stack.