41 lines
1.9 KiB
C#
41 lines
1.9 KiB
C#
#nullable enable
|
|
using System;
|
|
using System.Buffers;
|
|
|
|
namespace ShrinkNetwork
|
|
{
|
|
/// <summary>Owns rented memory until Dispose. A sender must await completion before disposing.</summary>
|
|
public sealed class ShrinkBufferWriter : IBufferWriter<byte>, IDisposable
|
|
{
|
|
private byte[]? _buffer;
|
|
public ShrinkBufferWriter(int initialCapacity = 256) => _buffer = ArrayPool<byte>.Shared.Rent(Math.Max(1, initialCapacity));
|
|
public int WrittenCount { get; private set; }
|
|
public ReadOnlyMemory<byte> WrittenMemory => Buffer.AsMemory(0, WrittenCount);
|
|
internal Span<byte> WrittenSpan => Buffer.AsSpan(0, WrittenCount);
|
|
private byte[] Buffer => _buffer ?? throw new ObjectDisposedException(nameof(ShrinkBufferWriter));
|
|
public void Advance(int count)
|
|
{
|
|
if (count < 0 || count > Buffer.Length - WrittenCount) throw new ArgumentOutOfRangeException(nameof(count));
|
|
WrittenCount += count;
|
|
}
|
|
public Memory<byte> GetMemory(int sizeHint = 0) { Ensure(sizeHint); return Buffer.AsMemory(WrittenCount); }
|
|
public Span<byte> GetSpan(int sizeHint = 0) { Ensure(sizeHint); return Buffer.AsSpan(WrittenCount); }
|
|
private void Ensure(int sizeHint)
|
|
{
|
|
if (sizeHint < 0) throw new ArgumentOutOfRangeException(nameof(sizeHint));
|
|
sizeHint = Math.Max(1, sizeHint);
|
|
if (sizeHint <= Buffer.Length - WrittenCount) return;
|
|
var next = ArrayPool<byte>.Shared.Rent(checked(Math.Max(Buffer.Length * 2, WrittenCount + sizeHint)));
|
|
Buffer.AsSpan(0, WrittenCount).CopyTo(next);
|
|
ArrayPool<byte>.Shared.Return(Buffer, clearArray: true);
|
|
_buffer = next;
|
|
}
|
|
public void Dispose()
|
|
{
|
|
var buffer = _buffer;
|
|
_buffer = null;
|
|
if (buffer != null) ArrayPool<byte>.Shared.Return(buffer, clearArray: true);
|
|
}
|
|
}
|
|
}
|