RavenDB-23704 - add a test

grisha-kotler · Feb 23, 2025 · 3561aec · 3561aec
1 parent 7a7e391
commit 3561aec
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 8 deletions.
diff --git a/src/Sparrow.Server/ByteString.cs b/src/Sparrow.Server/ByteString.cs
@@ -749,7 +749,7 @@ public override string ToString()
         /// </summary>
         private readonly List<SegmentInformation> _wholeSegments;
         private readonly int _initialAllocationBlockSize;
-        private int _allocationBlockSize;
+        public int AllocationBlockSize { get; private set; }
 
         internal long _totalAllocated, _currentlyAllocated;
 
@@ -780,7 +780,7 @@ public ByteStringContext(SharedMultipleUseFlag lowMemoryFlag, int allocationBloc
 
             _lowMemoryFlag = lowMemoryFlag;
             _initialAllocationBlockSize = allocationBlockSize;
-            _allocationBlockSize = allocationBlockSize;
+            AllocationBlockSize = allocationBlockSize;
 
             _wholeSegments = new List<SegmentInformation>();
             _internalReadyToUseMemorySegments = new List<SegmentInformation>();
@@ -825,7 +825,7 @@ public void Reset()
                 stack?.Clear();
             }
             _internalReadyToUseMemorySegments.Clear();// memory here will be released from _wholeSegments
-            _allocationBlockSize = _initialAllocationBlockSize;
+            AllocationBlockSize = _initialAllocationBlockSize;
 
             _externalStringPool.Clear();
             _externalFastPoolCount = 0;
@@ -1040,9 +1040,9 @@ private ByteString AllocateExternal(byte* valuePtr, int size, ByteStringType typ
             {
                 if (_externalCurrentLeft == 0)
                 {
-                    var tmp = Math.Min(2 * Sparrow.Global.Constants.Size.Megabyte, _allocationBlockSize * 2);
+                    var tmp = Math.Min(2 * Sparrow.Global.Constants.Size.Megabyte, AllocationBlockSize * 2);
                     AllocateExternalSegment(tmp);
-                    _allocationBlockSize = tmp;
+                    AllocationBlockSize = tmp;
                 }
 
                 storagePtr = (ByteStringStorage*)_externalCurrent.Current;
@@ -1074,7 +1074,7 @@ private ByteString AllocateInternal(int length, ByteStringType type)
             // This is even bigger than the configured allocation block size. There is no reason why we shouldn't
             // allocate it directly. When released (if released) this will be reused as a segment, ensuring that the context
             // could handle that.
-            if (allocationSize > _allocationBlockSize)
+            if (allocationSize > AllocationBlockSize)
             {
                 var segment = GetFromReadyToUseMemorySegments(allocationUnit);
                 if (segment != null)
@@ -1184,8 +1184,8 @@ private ByteString AllocateInternalUnlikely(int length, int allocationUnit, Byte
             }
             else
             {
-                _allocationBlockSize = Math.Min(2 * Sparrow.Global.Constants.Size.Megabyte, _allocationBlockSize * 2);
-                var toAllocate = Math.Max(_allocationBlockSize, allocationUnit);
+                AllocationBlockSize = Math.Min(2 * Sparrow.Global.Constants.Size.Megabyte, AllocationBlockSize * 2);
+                var toAllocate = Math.Max(AllocationBlockSize, allocationUnit);
                 _internalCurrent = AllocateSegment(toAllocate);
                 Debug.Assert(_internalCurrent.SizeLeft >= allocationUnit, $"{_internalCurrent.SizeLeft} >= {allocationUnit}");
             }

diff --git a/test/FastTests/Sparrow/ByteStringTests.cs b/test/FastTests/Sparrow/ByteStringTests.cs
@@ -161,6 +161,39 @@ public void AllocateAndReleaseShouldReuseRepeatedly()
             }
         }
 
+        [Fact]
+        public void CanResetAllocationBlockSize()
+        {
+            using (var context = new ByteStringContext<ByteStringDirectAllocator>(SharedMultipleUseFlag.None))
+            {
+                while (context.AllocationBlockSize == ByteStringContext.DefaultAllocationBlockSizeInBytes)
+                {
+                    context.Allocate(ByteStringContext.MinBlockSizeInBytes / 2, out _);
+                }
+
+                Assert.NotEqual(ByteStringContext.DefaultAllocationBlockSizeInBytes, context.AllocationBlockSize);
+
+                context.Reset();
+
+                Assert.Equal(ByteStringContext.DefaultAllocationBlockSizeInBytes, context.AllocationBlockSize);
+            }
+
+            var blockSizeInBytes = ByteStringContext.MinBlockSizeInBytes * 8;
+            using (var context = new ByteStringContext<ByteStringDirectAllocator>(SharedMultipleUseFlag.None, allocationBlockSize: blockSizeInBytes))
+            {
+                while (context.AllocationBlockSize == blockSizeInBytes)
+                {
+                    context.Allocate(blockSizeInBytes / 2, out _);
+                }
+
+                Assert.NotEqual(ByteStringContext.MinBlockSizeInBytes, context.AllocationBlockSize);
+
+                context.Reset();
+
+                Assert.Equal(blockSizeInBytes, context.AllocationBlockSize);
+            }
+        }
+
 #if VALIDATE
         [Fact]
         public void ValidationKeyAfterAllocateAndReleaseReuseShouldBeDifferent()