solve 'edge' case and pass cancellation token

This commit is contained in:
Dennis Arfan 2025-08-01 22:13:13 +07:00
parent 0472bfe58e
commit d3dfdd6a74
15 changed files with 208 additions and 551 deletions

View file

@ -24,7 +24,8 @@ public class ArrayOwner<T> : IBuffer<T> where T : unmanaged
public ref T this[int index] => ref _buffer[index];
public Span<T> Span => _buffer;
public Span<T> Span => _buffer.AsSpan(0, Length);
public Memory<T> Memory => _buffer.AsMemory(0, Length);
public T[] Array => _buffer;

View file

@ -5,6 +5,8 @@ public interface IBuffer<T> : IDisposable where T : unmanaged
ref T this[int index] { get; }
Span<T> Span { get; }
Memory<T> Memory { get; }
int Length { get; }
}

View file

@ -4,34 +4,48 @@ using System.Runtime.InteropServices;
namespace StitchATon2.Infra.Buffers;
internal sealed unsafe class ImmovableMemory<T> : MemoryManager<T> where T : unmanaged
internal sealed unsafe class ImmovableMemory<T> : MemoryManager<T>, IBuffer<T> where T : unmanaged
{
private readonly T* _pointer;
private readonly int _length;
internal readonly T* Pointer;
private bool _disposed;
public ImmovableMemory(int count)
public ImmovableMemory(int length)
{
_pointer = (T*)NativeMemory.Alloc((nuint)count, (nuint)Unsafe.SizeOf<T>());
_length = count;
Pointer = (T*)NativeMemory.Alloc((nuint)length, (nuint)Unsafe.SizeOf<T>());
Length = length;
}
protected override void Dispose(bool disposing)
{
if (!_disposed)
{
NativeMemory.Free(_pointer);
_disposed = true;
}
if (_disposed) return;
NativeMemory.Free(Pointer);
_disposed = true;
}
public override Span<T> GetSpan()
=> new(_pointer, _length);
=> _disposed
? throw new ObjectDisposedException(nameof(ImmovableMemory<T>))
: new Span<T>(Pointer, Length);
public override MemoryHandle Pin(int elementIndex = 0)
=> new(_pointer + elementIndex);
=> _disposed
? throw new ObjectDisposedException(nameof(ImmovableMemory<T>))
: new MemoryHandle(Pointer + elementIndex);
public override void Unpin()
{
}
public ref T this[int index]
{
get
{
if (_disposed) throw new ObjectDisposedException(nameof(ImmovableMemory<T>));
return ref Unsafe.AsRef<T>(Pointer + index);
}
}
public Span<T> Span => GetSpan();
public int Length { get; }
}

View file

@ -6,7 +6,7 @@ namespace StitchATon2.Infra.Buffers;
public static class MemoryAllocator
{
public static IBuffer<T> Allocate<T>(int count) where T : unmanaged
=> new UnmanagedMemory<T>(count);
=> new ImmovableMemory<T>(count);
public static IMemoryOwner<T> AllocateManaged<T>(int count)
=> MemoryPool<T>.Shared.Rent(count);
@ -14,31 +14,23 @@ public static class MemoryAllocator
public static ArrayOwner<T> AllocateArray<T>(int count) where T : unmanaged
=> new(ArrayPool<T>.Shared, count);
public static MemoryManager<T> AllocateImmovable<T>(int count) where T : unmanaged
=> new ImmovableMemory<T>(count);
public static unsafe IBuffer<T> Clone<T>(this IBuffer<T> buffer) where T : unmanaged
{
if (buffer is UnmanagedMemory<T> unmanagedMemory)
{
var newBuffer = new UnmanagedMemory<T>(buffer.Length);
var byteCount = (uint)(Unsafe.SizeOf<T>() * buffer.Length);
Unsafe.CopyBlock(newBuffer.Pointer, unmanagedMemory.Pointer, byteCount);
return newBuffer;
}
if (buffer is not ImmovableMemory<T> unmanagedMemory)
throw new NotSupportedException();
throw new NotSupportedException();
var newBuffer = new ImmovableMemory<T>(buffer.Length);
var byteCount = (uint)(Unsafe.SizeOf<T>() * buffer.Length);
Unsafe.CopyBlock(newBuffer.Pointer, unmanagedMemory.Pointer, byteCount);
return newBuffer;
}
public static unsafe void Copy<T>(this IBuffer<T> source, IBuffer<T> destination, int count) where T : unmanaged
{
if (source is UnmanagedMemory<T> sourceBuffer && destination is UnmanagedMemory<T> destinationBuffer)
{
var byteCount = (uint)(Unsafe.SizeOf<T>() * count);
Unsafe.CopyBlock(destinationBuffer.Pointer, sourceBuffer.Pointer, byteCount);
return;
}
if (source is not ImmovableMemory<T> sourceBuffer || destination is not ImmovableMemory<T> destinationBuffer)
throw new NotSupportedException();
throw new NotSupportedException();
var byteCount = (uint)(Unsafe.SizeOf<T>() * count);
Unsafe.CopyBlock(destinationBuffer.Pointer, sourceBuffer.Pointer, byteCount);
}
}

View file

@ -7,11 +7,13 @@ namespace StitchATon2.Infra.Buffers;
/// Provide non-thread safe anti GC contiguous memory.
/// </summary>
/// <typeparam name="T"></typeparam>
[Obsolete("Use immovable memory instead")]
internal sealed unsafe class UnmanagedMemory<T> : IBuffer<T> where T : unmanaged
{
internal readonly T* Pointer;
private bool _disposed;
public Memory<T> Memory => throw new NotImplementedException();
public int Length { get; }
public ref T this[int index] => ref Unsafe.AsRef<T>(Pointer + index);

View file

@ -63,26 +63,32 @@ public class PngPipeEncoder : IDisposable
public async Task WriteDataAsync(IBuffer<byte> buffer, bool disposeBuffer = true, CancellationToken cancellationToken = default)
{
_zlibStream.Write([0]);
var offset = 0;
while (buffer.Length - offset > FlushThreshold)
try
{
_zlibStream.Write(buffer.Span.Slice(offset, FlushThreshold));
await _zlibStream.FlushAsync(cancellationToken);
offset += FlushThreshold;
if(_memoryStream.Length >= BufferSize)
await FlushAsync(cancellationToken);
_zlibStream.Write([0]);
var offset = 0;
while (buffer.Length - offset > FlushThreshold)
{
_zlibStream.Write(buffer.Span.Slice(offset, FlushThreshold));
await _zlibStream.FlushAsync(cancellationToken);
offset += FlushThreshold;
if (_memoryStream.Length >= BufferSize)
await FlushAsync(cancellationToken);
}
if (buffer.Length > offset)
{
_zlibStream.Write(buffer.Span[offset..]);
await _zlibStream.FlushAsync(cancellationToken);
_shouldFlush = true;
}
}
if (buffer.Length > offset)
finally
{
_zlibStream.Write(buffer.Span[offset..]);
await _zlibStream.FlushAsync(cancellationToken);
_shouldFlush = true;
if(disposeBuffer)
buffer.Dispose();
}
if(disposeBuffer) buffer.Dispose();
}
private async Task FlushAsync(CancellationToken cancellationToken)
@ -101,7 +107,7 @@ public class PngPipeEncoder : IDisposable
// write Crc
var crc = Crc32.Compute(buffer.AsSpan(4, dataSize + 4));
BinaryPrimitives.WriteUInt32BigEndian(buffer.AsSpan(dataSize + 8), crc);
_outputPipe.Write(buffer.AsSpan(0, dataSize + 12));
await _outputPipe.FlushAsync(cancellationToken);

View file

@ -1,5 +1,6 @@
using System.Buffers.Binary;
using System.IO.Compression;
using StitchATon2.Infra.Buffers;
namespace StitchATon2.Infra.Encoders;
@ -30,7 +31,7 @@ public class PngStreamEncoder : IDisposable, IAsyncDisposable
~PngStreamEncoder() => Dispose();
public async Task WriteHeader()
public async Task WriteHeader(CancellationToken cancellationToken = default)
{
byte[] headerBytes = [
0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG Signature
@ -54,34 +55,42 @@ public class PngStreamEncoder : IDisposable, IAsyncDisposable
BinaryPrimitives.WriteUInt32BigEndian(headerBytes.AsSpan(29), crc);
await _stream.WriteAsync(headerBytes);
await _stream.WriteAsync(headerBytes, cancellationToken);
}
public async Task WriteData(Memory<byte> data)
public async Task WriteDataAsync(IBuffer<byte> buffer, bool disposeBuffer = true, CancellationToken cancellationToken = default)
{
_zlibStream.Write([0]);
var dataSlice = data;
while (dataSlice.Length > FlushThreshold)
try
{
await _zlibStream.WriteAsync(dataSlice[..FlushThreshold]);
await _zlibStream.FlushAsync();
dataSlice = dataSlice[FlushThreshold..];
if(_memoryStream.Length >= BufferSize)
await Flush();
_zlibStream.Write([0]);
var dataSlice = buffer.Memory;
while (dataSlice.Length > FlushThreshold)
{
await _zlibStream.WriteAsync(dataSlice[..FlushThreshold], cancellationToken);
await _zlibStream.FlushAsync(cancellationToken);
dataSlice = dataSlice[FlushThreshold..];
if (_memoryStream.Length >= BufferSize)
await FlushAsync(cancellationToken);
}
if (dataSlice.Length > 0)
{
await _zlibStream.WriteAsync(dataSlice, cancellationToken);
await _zlibStream.FlushAsync(cancellationToken);
_shouldFlush = true;
}
}
if (dataSlice.Length > 0)
finally
{
await _zlibStream.WriteAsync(dataSlice);
await _zlibStream.FlushAsync();
_shouldFlush = true;
if(disposeBuffer)
buffer.Dispose();
}
}
private async Task Flush()
private async Task FlushAsync(CancellationToken cancellationToken)
{
await _zlibStream.FlushAsync();
await _zlibStream.FlushAsync(cancellationToken);
var dataSize = (int)(_memoryStream.Length - 8);
_memoryStream.Write("\0\0\0\0"u8);
@ -96,16 +105,16 @@ public class PngStreamEncoder : IDisposable, IAsyncDisposable
var crc = Crc32.Compute(buffer.AsSpan(4, dataSize + 4));
BinaryPrimitives.WriteUInt32BigEndian(buffer.AsSpan(dataSize + 8), crc);
await _stream.WriteAsync(buffer.AsMemory(0, dataSize + 12));
await _stream.WriteAsync(buffer.AsMemory(0, dataSize + 12), cancellationToken);
_memoryStream.SetLength(8);
_memoryStream.Position = 8;
_shouldFlush = false;
}
public async ValueTask WriteEndOfFile()
public async ValueTask WriteEndOfFileAsync(CancellationToken cancellationToken = default)
{
if(_shouldFlush)
await Flush();
await FlushAsync(cancellationToken);
var endChunk = new byte[] {
0x00, 0x00, 0x00, 0x00, // Length
@ -113,7 +122,7 @@ public class PngStreamEncoder : IDisposable, IAsyncDisposable
0xAE, 0x42, 0x60, 0x82, // Crc
};
await _stream.WriteAsync(endChunk);
await _stream.WriteAsync(endChunk, cancellationToken);
await DisposeAsync();
}