Created
January 29, 2019 01:11
-
-
Save ahsonkhan/c76a1cc4dc7107537c3fdc0079a68b35 to your computer and use it in GitHub Desktop.
Sample implementation of ArrayBufferWriter, a class that implements IBufferWriter<byte> backed by ArrayPool<byte>.Shared.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
public class ArrayBufferWriter : IBufferWriter<byte>, IDisposable | |
{ | |
private byte[] _rentedBuffer; | |
private int _written; | |
private long _committed; | |
private const int MinimumBufferSize = 256; | |
public ArrayBufferWriter(int initialCapacity = MinimumBufferSize) | |
{ | |
if (initialCapacity <= 0) | |
throw new ArgumentException(nameof(initialCapacity)); | |
_rentedBuffer = ArrayPool<byte>.Shared.Rent(initialCapacity); | |
_written = 0; | |
_committed = 0; | |
} | |
public Memory<byte> OutputAsMemory | |
{ | |
get | |
{ | |
CheckIfDisposed(); | |
return _rentedBuffer.AsMemory(0, _written); | |
} | |
} | |
public Span<byte> OutputAsSpan | |
{ | |
get | |
{ | |
CheckIfDisposed(); | |
return _rentedBuffer.AsSpan(0, _written); | |
} | |
} | |
public int BytesWritten | |
{ | |
get | |
{ | |
CheckIfDisposed(); | |
return _written; | |
} | |
} | |
public long BytesCommitted | |
{ | |
get | |
{ | |
CheckIfDisposed(); | |
return _committed; | |
} | |
} | |
public void Clear() | |
{ | |
CheckIfDisposed(); | |
ClearHelper(); | |
} | |
private void ClearHelper() | |
{ | |
_rentedBuffer.AsSpan(0, _written).Clear(); | |
_written = 0; | |
} | |
public async Task CopyToAsync(Stream stream, CancellationToken cancellationToken = default) | |
{ | |
CheckIfDisposed(); | |
if (stream == null) | |
throw new ArgumentNullException(nameof(stream)); | |
await stream.WriteAsync(_rentedBuffer, 0, _written, cancellationToken).ConfigureAwait(false); | |
_committed += _written; | |
ClearHelper(); | |
} | |
public void CopyTo(Stream stream) | |
{ | |
CheckIfDisposed(); | |
if (stream == null) | |
throw new ArgumentNullException(nameof(stream)); | |
stream.Write(_rentedBuffer, 0, _written); | |
_committed += _written; | |
ClearHelper(); | |
} | |
public void Advance(int count) | |
{ | |
CheckIfDisposed(); | |
if (count < 0) | |
throw new ArgumentException(nameof(count)); | |
if (_written > _rentedBuffer.Length - count) | |
throw new InvalidOperationException("Cannot advance past the end of the buffer."); | |
_written += count; | |
} | |
// Returns the rented buffer back to the pool | |
public void Dispose() | |
{ | |
if (_rentedBuffer == null) | |
{ | |
return; | |
} | |
ArrayPool<byte>.Shared.Return(_rentedBuffer, clearArray: true); | |
_rentedBuffer = null; | |
_written = 0; | |
} | |
private void CheckIfDisposed() | |
{ | |
if (_rentedBuffer == null) | |
throw new ObjectDisposedException(nameof(ArrayBufferWriter)); | |
} | |
public Memory<byte> GetMemory(int sizeHint = 0) | |
{ | |
CheckIfDisposed(); | |
if (sizeHint < 0) | |
throw new ArgumentException(nameof(sizeHint)); | |
CheckAndResizeBuffer(sizeHint); | |
return _rentedBuffer.AsMemory(_written); | |
} | |
public Span<byte> GetSpan(int sizeHint = 0) | |
{ | |
CheckIfDisposed(); | |
if (sizeHint < 0) | |
throw new ArgumentException(nameof(sizeHint)); | |
CheckAndResizeBuffer(sizeHint); | |
return _rentedBuffer.AsSpan(_written); | |
} | |
private void CheckAndResizeBuffer(int sizeHint) | |
{ | |
Debug.Assert(sizeHint >= 0); | |
if (sizeHint == 0) | |
{ | |
sizeHint = MinimumBufferSize; | |
} | |
int availableSpace = _rentedBuffer.Length - _written; | |
if (sizeHint > availableSpace) | |
{ | |
int growBy = sizeHint > _rentedBuffer.Length ? sizeHint : _rentedBuffer.Length; | |
int newSize = checked(_rentedBuffer.Length + growBy); | |
byte[] oldBuffer = _rentedBuffer; | |
_rentedBuffer = ArrayPool<byte>.Shared.Rent(newSize); | |
Debug.Assert(oldBuffer.Length >= _written); | |
Debug.Assert(_rentedBuffer.Length >= _written); | |
oldBuffer.AsSpan(0, _written).CopyTo(_rentedBuffer); | |
ArrayPool<byte>.Shared.Return(oldBuffer, clearArray: true); | |
} | |
Debug.Assert(_rentedBuffer.Length - _written > 0); | |
Debug.Assert(_rentedBuffer.Length - _written >= sizeHint); | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment