That was a nightmare of an bug fixed. The PooledMemoryBufferWriter.cs…

… wasn't properly resizing
Nexus-Mods · Jan 23, 2024 · b30774f · b30774f
1 parent ffcbcf3
commit b30774f
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 36 deletions.
diff --git a/benchmarks/NexusMods.EventSourcing.Benchmarks/Program.cs b/benchmarks/NexusMods.EventSourcing.Benchmarks/Program.cs
@@ -34,7 +34,7 @@
 benchmarks.Setup();
 
 var sw = Stopwatch.StartNew();
-for (var i = 0; i < 100000; i++)
+for (var i = 0; i < 100; i++)
 {
     benchmarks.LoadAllEntities();
 }

diff --git a/src/NexusMods.EventSourcing.Abstractions/IndexedMultiEntityAttributeDefinition.cs b/src/NexusMods.EventSourcing.Abstractions/IndexedMultiEntityAttributeDefinition.cs
@@ -2,6 +2,7 @@
 using System.Buffers;
 using System.Buffers.Binary;
 using System.Collections.Generic;
+using System.Linq;
 using NexusMods.EventSourcing.Abstractions.Serialization;
 
 namespace NexusMods.EventSourcing.Abstractions;
@@ -85,36 +86,24 @@ public class IndexedMultiEntityAccumulator<TKey, TOther> : IAccumulator
     /// <inheritdoc />
     public void WriteTo(IBufferWriter<byte> writer, ISerializationRegistry registry)
     {
-        var getSpan = writer.GetSpan(2);
-        BinaryPrimitives.WriteUInt16BigEndian(getSpan, (ushort) _values.Count);
-        writer.Advance(2);
-
-        foreach (var (key, value) in _values)
-        {
-            registry.Serialize(writer, key);
-            registry.Serialize(writer, value);
-        }
+        registry.Serialize(writer, _keys.Keys.ToArray());
+        registry.Serialize(writer, _keys.Values.ToArray());
     }
 
     /// <inheritdoc />
     public int ReadFrom(ref ReadOnlySpan<byte> input, ISerializationRegistry registry)
     {
-        var originalSize = input.Length;
-        var data = input;
+        var keySize = registry.Deserialize<EntityId<TOther>[]>(input, out var keys);
+        input = input.Slice(keySize);
+        var valueSize = registry.Deserialize<TKey[]>(input, out var values);
 
-        var count = BinaryPrimitives.ReadUInt16BigEndian(data);
-        data = data.Slice(sizeof(ushort));
-
-        for (var idx = 0; idx < count; idx++)
+        for (var i = 0; i < keys.Length; i++)
         {
-            var amountRead = registry.Deserialize(data, out TKey key);
-            data = data.Slice(amountRead);
-            amountRead = registry.Deserialize(data, out EntityId<TOther> value);
-            data = data.Slice(amountRead);
-            _values.Add(key, value);
-            _keys.Add(value, key);
+            _keys.Add(keys[i], values[i]);
+            _values.Add(values[i], keys[i]);
         }
 
-        return originalSize - data.Length;
+
+        return keySize + valueSize;
     }
 }
diff --git a/src/NexusMods.EventSourcing/PooledMemoryBufferWriter.cs b/src/NexusMods.EventSourcing/PooledMemoryBufferWriter.cs
@@ -43,14 +43,15 @@ public void Reset()
     /// <returns></returns>
     public ReadOnlySpan<byte> GetWrittenSpan() => _data.Span.SliceFast(0, _idx);
 
-    private void Expand()
+    private void Expand(int atLeast)
     {
-        var newSize = _data.Length * 2;
-        var newData = MemoryPool<byte>.Shared.Rent(newSize);
-        _data.CopyTo(newData.Memory);
-        _owner.Dispose();
-        _owner = newData;
-        _data = newData.Memory;
+        var newSize = _size;
+        while (newSize < atLeast)
+            newSize *= 2;
+
+        var newData = new Memory<byte>(new byte[newSize]);
+        _data.CopyTo(newData);
+        _data = newData;
         _size = newSize;
     }
 
@@ -59,25 +60,30 @@ private void Expand()
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public void Advance(int count)
     {
-        if (_idx + count > _size)
-            Expand();
+        if (_idx + count >= _size)
+            Expand(_idx + count);
+        Debug.Assert(_idx + count <= _size);
         _idx += count;
     }
 
     /// <inheritdoc />
     public Memory<byte> GetMemory(int sizeHint = 0)
     {
         if (_idx + sizeHint > _size)
-            Expand();
+            Expand(_idx + sizeHint);
+
+        Debug.Assert(_idx + sizeHint <= _size);
         return _data[_idx..];
     }
 
     /// <inheritdoc />
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public Span<byte> GetSpan(int sizeHint = 0)
     {
-        if (_idx + sizeHint > _size)
-            Expand();
-        return _data.Span.SliceFast(_idx);
+        if (_idx + sizeHint >= _size)
+            Expand(_idx + sizeHint);
+
+        Debug.Assert(_idx + sizeHint <= _size);
+        return _data.Span.SliceFast(_idx, sizeHint);
     }
 }