Skip to content

Commit 18ca20f

Browse files
committed
[Proto] Fix ProtoWriter buffer growth for large byte arrays
Adjusts ProtoWriter to account for varint length prefix when growing the buffer before writing large byte arrays. Adds a test to verify correct encoding and decoding of 1MB byte arrays.
1 parent 6212f2f commit 18ca20f

2 files changed

Lines changed: 25 additions & 2 deletions

File tree

Lagrange.Proto.Test/SegmentBufferTest.cs

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using Lagrange.Proto.Primitives;
12
using Lagrange.Proto.Utility;
23

34
namespace Lagrange.Proto.Test;
@@ -29,7 +30,28 @@ public void TestCreateReadeOnlyMemory()
2930
Assert.That(reader.Length, Is.EqualTo(_raw.Length));
3031
Assert.That(reader.Span.ToArray(), Is.EqualTo(_raw));
3132
}
32-
33+
34+
[Test]
35+
public void TestProtoWriteLargeByteArray()
36+
{
37+
var segments = new SegmentBufferWriter();
38+
var writer = new ProtoWriter(segments);
39+
40+
var random = GC.AllocateUninitializedArray<byte>(1024 * 1024 * 1); // 1 MB
41+
Random.Shared.NextBytes(random);
42+
43+
writer.EncodeBytes(random.AsSpan());
44+
writer.Flush();
45+
46+
var writtenData = segments.CreateReadOnlyMemory();
47+
var reader = new ProtoReader(writtenData.Span);
48+
int readData = reader.DecodeVarInt<int>();
49+
var readSpan = reader.CreateSpan(readData);
50+
51+
Assert.That(readData, Is.EqualTo(random.Length));
52+
Assert.That(readSpan.ToArray(), Is.EqualTo(random));
53+
}
54+
3355
[TearDown]
3456
public void TearDown()
3557
{

Lagrange.Proto/Primitives/ProtoWriter.cs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,8 @@ public void EncodeString(ReadOnlySpan<char> str)
6666
public void EncodeBytes(ReadOnlySpan<byte> bytes)
6767
{
6868
int length = bytes.Length;
69-
if (_memory.Length - BytesPending < length) Grow(length);
69+
int count = ProtoHelper.GetVarIntLength(length);
70+
if (_memory.Length - BytesPending < length) Grow(length + count);
7071

7172
EncodeVarInt(length);
7273
bytes.CopyTo(_memory.Span[BytesPending..]);

0 commit comments

Comments
 (0)