Skip to content

Implements Stream.ToOwnedMemory for .NET Framework. #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
isaksky merged 4 commits into from
May 19, 2025
Merged

Implements Stream.ToOwnedMemory for .NET Framework. #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f803497
Create as large a buffer as possible when reading from Stream.
oconnor0 May 19, 2025
1dc1411
Implement Stream.ToOwnedMemory for .NET Framework.
oconnor0 May 19, 2025
5be43bc
Test that reading from a stream with an unknown size does so correctly.
oconnor0 May 19, 2025
d6d41f0
Avoid sharing stream probes.
oconnor0 May 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 68 additions & 1 deletion Xledger.Collections.Test/TestMemoryOwner.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,22 @@ public void TestSlice() {
Assert.Equal(array.AsMemory().Slice(3).Slice(1, 4), sliced2.Memory);
}

#if NET
[Fact]
public void TestStream_ToMemoryOwner() {
// This length should be larger than the default GetCopyBufferSize.
byte[] array = new byte[4 * 1024 * 1024 + 3];
new Random().NextBytes(array);
var ms = new MemoryStream(array);
using var memoryOwner = ms.ToOwnedMemory();
#if NET
Assert.Equal(array.AsMemory(), memoryOwner.Memory);
#else
var mem = memoryOwner.Memory;
Assert.Equal(array.Length, mem.Length);
for (int i = 0; i < array.Length; ++i) {
Assert.Equal(array[i], mem.Span[i]);
}
#endif
}

[Fact]
Expand All @@ -41,7 +48,67 @@ public async Task TestStream_ToMemoryOwnerAsync() {
new Random().NextBytes(array);
var ms = new MemoryStream(array);
using var memoryOwner = await ms.ToOwnedMemoryAsync();
#if NET
Assert.Equal(array.AsMemory(), memoryOwner.Memory);
#else
var mem = memoryOwner.Memory;
Assert.Equal(array.Length, mem.Length);
for (int i = 0; i < array.Length; ++i) {
Assert.Equal(array[i], mem.Span[i]);
}
#endif
}

// Tests that reading from a stream with an unknown size does so correctly.
[Fact]
public void TestUnsizedStream_ToMemoryOwner() {
// This length should be larger than the default GetCopyBufferSize.
byte[] array = new byte[2 * 1024 * 1024 + 17];
new Random().NextBytes(array);
var ms = new UnsizedMemoryStream(array);
using var memoryOwner = ms.ToOwnedMemory();
#if NET
Assert.Equal(array.AsMemory(), memoryOwner.Memory);
#else
var mem = memoryOwner.Memory;
Assert.Equal(array.Length, mem.Length);
for (int i = 0; i < array.Length; ++i) {
Assert.Equal(array[i], mem.Span[i]);
}
#endif
}

// Tests that reading from a stream with an unknown size does so correctly.
[Fact]
public async Task TestUnsizedStream_ToMemoryOwnerAsync() {
// This length should be larger than the default GetCopyBufferSize.
byte[] array = new byte[2 * 1024 * 1024 + 17];
new Random().NextBytes(array);
var ms = new UnsizedMemoryStream(array);
using var memoryOwner = await ms.ToOwnedMemoryAsync();
#if NET
Assert.Equal(array.AsMemory(), memoryOwner.Memory);
#else
var mem = memoryOwner.Memory;
Assert.Equal(array.Length, mem.Length);
for (int i = 0; i < array.Length; ++i) {
Assert.Equal(array[i], mem.Span[i]);
}
#endif
}

class UnsizedMemoryStream(byte[] buffer) : MemoryStream(buffer) {
public override bool CanSeek => false;

public override int Capacity {
get => throw new NotImplementedException();
set => throw new NotImplementedException();
}

public override long Length => throw new NotImplementedException();

public override void SetLength(long value) {
throw new NotImplementedException();
}
}
}
99 changes: 63 additions & 36 deletions Xledger.Collections/Memory/Extensions.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,51 +24,69 @@ public static IMemoryOwner<T> Slice<T>(this IMemoryOwner<T> memoryOwner, int sta
return new SizedMemoryOwner<T>(memoryOwner, start, length);
}

static readonly int ArrayMaxLength =
#if NET
Array.MaxLength;
#else
int.MaxValue;
#endif

static readonly ThreadLocal<byte[]> PROBE = new(() => new byte[1]);

public static IMemoryOwner<byte> ToOwnedMemory(this Stream source, bool leaveOpen = false) {
if (source == null) {
throw new ArgumentNullException(nameof(source));
}

int initialBufLen = GetCopyBufferSize(source);
(bool canHoldEntireStream, int initialBufLen) = GetBufferSize(source);

var currentOwner = MemoryPool<byte>.Shared.Rent(initialBufLen);
var currentBuffer = currentOwner.Memory;
var currentBuffer = ArrayPool<byte>.Shared.Rent(initialBufLen);
var currentOwner = currentBuffer.ToOwnedMemory(ArrayPool<byte>.Shared);
int totalBytesRead = 0;

try {
while (true) {
var dest = currentBuffer.Slice(totalBytesRead);

int bytesRead = source.Read(dest.Span);
int bytesRead = source.Read(
currentBuffer,
totalBytesRead,
currentBuffer.Length - totalBytesRead);

if (bytesRead == 0) {
break;
}

totalBytesRead += bytesRead;

if (canHoldEntireStream && totalBytesRead == initialBufLen) {
// We've read the entire stream.
break;
}

if (totalBytesRead != currentBuffer.Length) {
continue;
}

if (currentBuffer.Length == Array.MaxLength) {
if (currentBuffer.Length == ArrayMaxLength) {
#if NET
Span<byte> probe = stackalloc byte[1];
if (source.Read(probe) > 0) {
throw new IOException($"Stream exceeds the maximum bufferable array size of {Array.MaxLength} bytes.");
#else
if (source.Read(PROBE.Value, 0, 1) > 0) {
#endif
throw new IOException($"Stream exceeds the maximum bufferable array size of {ArrayMaxLength} bytes.");
}
break; // we are at the end of the stream
}

var newCapacity = (long)currentBuffer.Length * 2;
if (newCapacity > Array.MaxLength) {
newCapacity = Array.MaxLength;
if (newCapacity > ArrayMaxLength) {
newCapacity = ArrayMaxLength;
}

var newOwner = MemoryPool<byte>.Shared.Rent((int)newCapacity);
var newBuffer = newOwner.Memory;
var newBuffer = ArrayPool<byte>.Shared.Rent((int)newCapacity);
var newOwner = newBuffer.ToOwnedMemory(ArrayPool<byte>.Shared);

currentBuffer.CopyTo(newBuffer);
currentBuffer.CopyTo(newBuffer.AsSpan());
currentOwner.Dispose();
currentOwner = newOwner;
currentBuffer = newBuffer;
Expand All @@ -85,53 +103,58 @@ public static IMemoryOwner<byte> ToOwnedMemory(this Stream source, bool leaveOpe
return currentOwner.Slice(0, totalBytesRead);
}

static readonly byte[] ASYNC_PROBE = new byte[1];

public static async Task<IMemoryOwner<byte>> ToOwnedMemoryAsync(this Stream source, bool leaveOpen = false, CancellationToken tok = default) {
if (source == null) {
throw new ArgumentNullException(nameof(source));
}

int initialBufLen = GetCopyBufferSize(source);
(bool canHoldEntireStream, int initialBufLen) = GetBufferSize(source);

var currentOwner = MemoryPool<byte>.Shared.Rent(initialBufLen);
var currentBuffer = currentOwner.Memory;
var currentBuffer = ArrayPool<byte>.Shared.Rent(initialBufLen);
var currentOwner = currentBuffer.ToOwnedMemory(ArrayPool<byte>.Shared);
int totalBytesRead = 0;

try {
while (true) {
tok.ThrowIfCancellationRequested();

var dest = currentBuffer.Slice(totalBytesRead);

int bytesRead = await source.ReadAsync(dest, tok).ConfigureAwait(false);
int bytesRead = await source.ReadAsync(
currentBuffer,
totalBytesRead,
currentBuffer.Length - totalBytesRead,
tok).ConfigureAwait(false);

if (bytesRead == 0) {
break;
}

totalBytesRead += bytesRead;

if (canHoldEntireStream && totalBytesRead == initialBufLen) {
// We've read the entire stream.
break;
}

if (totalBytesRead != currentBuffer.Length) {
continue;
}

if (currentBuffer.Length == Array.MaxLength) {
if (await source.ReadAsync(ASYNC_PROBE, tok).ConfigureAwait(false) > 0) {
throw new IOException($"Stream exceeds the maximum bufferable array size of {Array.MaxLength} bytes.");
if (currentBuffer.Length == ArrayMaxLength) {
if (await source.ReadAsync(PROBE.Value, 0, 1, tok).ConfigureAwait(false) > 0) {
throw new IOException($"Stream exceeds the maximum bufferable array size of {ArrayMaxLength} bytes.");
}
break; // we are at the end of the stream
}

var newCapacity = (long)currentBuffer.Length * 2;
if (newCapacity > Array.MaxLength) {
newCapacity = Array.MaxLength;
if (newCapacity > ArrayMaxLength) {
newCapacity = ArrayMaxLength;
}

var newOwner = MemoryPool<byte>.Shared.Rent((int)newCapacity);
var newBuffer = newOwner.Memory;
var newBuffer = ArrayPool<byte>.Shared.Rent((int)newCapacity);
var newOwner = newBuffer.ToOwnedMemory(ArrayPool<byte>.Shared);

currentBuffer.CopyTo(newBuffer);
currentBuffer.CopyTo(newBuffer.AsSpan());
currentOwner.Dispose();
currentOwner = newOwner;
currentBuffer = newBuffer;
Expand All @@ -148,8 +171,9 @@ public static async Task<IMemoryOwner<byte>> ToOwnedMemoryAsync(this Stream sour
return currentOwner.Slice(0, totalBytesRead);
}

// Copied from System.IO.Stream, adapted to be static
static int GetCopyBufferSize(Stream stream) {
// Initially copied from System.IO.Stream, adapted to be static and to match
// the use above which is to copy an entire stream into a single array.
static (bool isSufficient, int length) GetBufferSize(Stream stream) {
// This value was originally picked to be the largest multiple of 4096 that is still smaller than the large object heap threshold (85K).
// The CopyTo{Async} buffer is short-lived and is likely to be collected at Gen0, and it offers a significant improvement in Copy
// performance. Since then, the base implementations of CopyTo{Async} have been updated to use ArrayPool, which will end up rounding
Expand All @@ -158,6 +182,7 @@ static int GetCopyBufferSize(Stream stream) {
// benefits to using the larger buffer size. So, for now, this value remains.
const int DefaultCopyBufferSize = 81920;

bool isSufficient = false;
int bufferSize = DefaultCopyBufferSize;

if (stream.CanSeek) {
Expand All @@ -172,16 +197,18 @@ static int GetCopyBufferSize(Stream stream) {
bufferSize = 1;
} else {
long remaining = length - position;
if (remaining > 0) {
// In the case of a positive overflow, stick to the default size
bufferSize = (int)Math.Min(bufferSize, remaining);
if (remaining > ArrayMaxLength) {
throw new IOException($"Stream exceeds the maximum bufferable array size of {ArrayMaxLength} bytes.");
} else if (remaining > 0) {
// If there is some remaining amount in the stream, we copy into a buffer of that size.
isSufficient = true;
bufferSize = (int)remaining;
}
}
}

return bufferSize;
return (isSufficient, bufferSize);
}
#endif

/// <summary>
/// Adapt an array to IMemoryOwner. If you pass in an ArrayPool owner, the Array will be returned to the pool on dispose.
Expand Down
2 changes: 1 addition & 1 deletion Xledger.Collections/Xledger.Collections.csproj
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
<AssemblyName>Xledger.Collections</AssemblyName>

<TargetFrameworks>net48;net8.0</TargetFrameworks>
<LangVersion>12.0</LangVersion>
<LangVersion>13.0</LangVersion>
<ImplicitUsings>disable</ImplicitUsings>
<Nullable>disable</Nullable>

Expand Down