Merge pull request litedb-org#2508 from mbdavid/pj/ArrayPool-4tw

[♻️Housekeeping] Remove BufferPool class and use .net implementation
jiuaizxc · Jun 19, 2024 · b82f30e · b82f30e
2 parents 3b7cc00 + ae62518
commit b82f30e
Show file tree

Hide file tree

Showing 14 changed files with 89 additions and 198 deletions.
diff --git a/LiteDB/Engine/Disk/DiskService.cs b/LiteDB/Engine/Disk/DiskService.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.IO;
 using System.Threading;
@@ -25,6 +26,8 @@ internal class DiskService : IDisposable
         private long _dataLength;
         private long _logLength;
 
+        private static readonly ArrayPool<byte> _bufferPool = ArrayPool<byte>.Shared;
+
         public DiskService(
             EngineSettings settings,
             EngineState state,
@@ -222,11 +225,12 @@ internal void MarkAsInvalidState()
             {
                 using (var stream = _dataFactory.GetStream(true, true))
                 {
-                    var buffer = new byte[PAGE_SIZE];
+                    var buffer = _bufferPool.Rent(PAGE_SIZE);
                     stream.Read(buffer, 0, PAGE_SIZE);
                     buffer[HeaderPage.P_INVALID_DATAFILE_STATE] = 1;
                     stream.Position = 0;
                     stream.Write(buffer, 0, PAGE_SIZE);
+                    _bufferPool.Return(buffer, true);
                 }
             });
         }

diff --git a/LiteDB/Engine/Disk/Serializer/BufferPool.cs b/LiteDB/Engine/Disk/Serializer/BufferPool.cs
diff --git a/LiteDB/Engine/Disk/Serializer/BufferReader.cs b/LiteDB/Engine/Disk/Serializer/BufferReader.cs
@@ -1,7 +1,7 @@
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.IO;
-using System.Text;
 using static LiteDB.Constants;
 
 namespace LiteDB.Engine
@@ -20,6 +20,8 @@ internal class BufferReader : IDisposable
 
         private bool _isEOF = false;
 
+        private static readonly ArrayPool<byte> _bufferPool = ArrayPool<byte>.Shared;
+
         /// <summary>
         /// Current global cursor position
         /// </summary>
@@ -102,10 +104,10 @@ public int Read(byte[] buffer, int offset, int count)
                 // fill buffer
                 if (buffer != null)
                 {
-                    Buffer.BlockCopy(_current.Array, 
-                        _current.Offset + _currentPosition, 
-                        buffer, 
-                        offset + bufferPosition, 
+                    Buffer.BlockCopy(_current.Array,
+                        _current.Offset + _currentPosition,
+                        buffer,
+                        offset + bufferPosition,
                         bytesToCopy);
                 }
 
@@ -161,13 +163,13 @@ public string ReadString(int count)
             else
             {
                 // rent a buffer to be re-usable
-                var buffer = BufferPool.Rent(count);
+                var buffer = _bufferPool.Rent(count);
 
                 this.Read(buffer, 0, count);
 
                 value = StringEncoding.UTF8.GetString(buffer, 0, count);
 
-                BufferPool.Return(buffer);
+                _bufferPool.Return(buffer, true);
             }
 
             return value;
@@ -251,13 +253,13 @@ private T ReadNumber<T>(Func<byte[], int, T> convert, int size)
             }
             else
             {
-                var buffer = BufferPool.Rent(size);
+                var buffer = _bufferPool.Rent(size);
 
                 this.Read(buffer, 0, size);
 
                 value = convert(buffer, 0);
 
-                BufferPool.Return(buffer);
+                _bufferPool.Return(buffer, true);
             }
 
             return value;
@@ -328,13 +330,13 @@ public ObjectId ReadObjectId()
             }
             else
             {
-                var buffer = BufferPool.Rent(12);
+                var buffer = _bufferPool.Rent(12);
 
                 this.Read(buffer, 0, 12);
 
                 value = new ObjectId(buffer, 0);
 
-                BufferPool.Return(buffer);
+                _bufferPool.Return(buffer, true);
             }
 
             return value;
@@ -393,7 +395,7 @@ public BsonValue ReadIndexKey()
                 case BsonType.Int64: return this.ReadInt64();
                 case BsonType.Double: return this.ReadDouble();
                 case BsonType.Decimal: return this.ReadDecimal();
-                
+
                 // Use +1 byte only for length
                 case BsonType.String: return this.ReadString(this.ReadByte());
 

diff --git a/LiteDB/Engine/Disk/Serializer/BufferWriter.cs b/LiteDB/Engine/Disk/Serializer/BufferWriter.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.Text;
 using static LiteDB.Constants;
@@ -18,6 +19,8 @@ internal class BufferWriter : IDisposable
 
         private bool _isEOF = false;
 
+        private static readonly ArrayPool<byte> _bufferPool = ArrayPool<byte>.Shared;
+
         /// <summary>
         /// Current global cursor position
         /// </summary>
@@ -98,10 +101,10 @@ public int Write(byte[] buffer, int offset, int count)
                 // fill buffer
                 if (buffer != null)
                 {
-                    Buffer.BlockCopy(buffer, 
+                    Buffer.BlockCopy(buffer,
                         offset + bufferPosition,
                         _current.Array,
-                        _current.Offset + _currentPosition, 
+                        _current.Offset + _currentPosition,
                         bytesToCopy);
                 }
 
@@ -133,7 +136,7 @@ public int Write(byte[] buffer, int offset, int count)
         /// </summary>
         public void Consume()
         {
-            if(_source != null)
+            if (_source != null)
             {
                 while (_source.MoveNext())
                 {
@@ -166,7 +169,7 @@ public void WriteCString(string value)
             }
             else
             {
-                var buffer = BufferPool.Rent(bytesCount);
+                var buffer = _bufferPool.Rent(bytesCount);
 
                 StringEncoding.UTF8.GetBytes(value, 0, value.Length, buffer, 0);
 
@@ -176,7 +179,7 @@ public void WriteCString(string value)
 
                 this.MoveForward(1);
 
-                BufferPool.Return(buffer);
+                _bufferPool.Return(buffer, true);
             }
         }
 
@@ -202,13 +205,13 @@ public void WriteString(string value, bool specs)
             else
             {
                 // rent a buffer to be re-usable
-                var buffer = BufferPool.Rent(count);
+                var buffer = _bufferPool.Rent(count);
 
                 StringEncoding.UTF8.GetBytes(value, 0, value.Length, buffer, 0);
 
                 this.Write(buffer, 0, count);
 
-                BufferPool.Return(buffer);
+                _bufferPool.Return(buffer, true);
             }
 
             if (specs)
@@ -231,13 +234,13 @@ private void WriteNumber<T>(T value, Action<T, byte[], int> toBytes, int size)
             }
             else
             {
-                var buffer = BufferPool.Rent(size);
+                var buffer = _bufferPool.Rent(size);
 
                 toBytes(value, buffer, 0);
 
                 this.Write(buffer, 0, size);
 
-                BufferPool.Return(buffer);
+                _bufferPool.Return(buffer, true);
             }
         }
 
@@ -293,13 +296,13 @@ public void Write(ObjectId value)
             }
             else
             {
-                var buffer = BufferPool.Rent(12);
+                var buffer = _bufferPool.Rent(12);
 
                 value.ToByteArray(buffer, 0);
 
                 this.Write(buffer, 0, 12);
 
-                BufferPool.Return(buffer);
+                _bufferPool.Return(buffer, true);
             }
         }
 

diff --git a/LiteDB/Engine/Disk/Streams/AesStream.cs b/LiteDB/Engine/Disk/Streams/AesStream.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Buffers;
 using System.IO;
 using System.Linq;
 using System.Security.Cryptography;
@@ -24,6 +25,8 @@ public class AesStream : Stream
 
         private static readonly byte[] _emptyContent = new byte[PAGE_SIZE - 1 - 16]; // 1 for aes indicator + 16 for salt
 
+        private static readonly ArrayPool<byte> _bufferPool = ArrayPool<byte>.Shared;
+
         public byte[] Salt { get; }
 
         public override bool CanRead => _stream.CanRead;
@@ -52,6 +55,11 @@ public AesStream(string password, Stream stream)
             // start stream from zero position
             _stream.Position = 0;
 
+            const int checkBufferSize = 32;
+
+            var checkBuffer = _bufferPool.Rent(checkBufferSize);
+            var msBuffer = _bufferPool.Rent(16);
+
             try
             {
                 // new file? create new salt
@@ -104,33 +112,32 @@ public AesStream(string password, Stream stream)
                 // set stream to password checking
                 _stream.Position = 32;
 
-                var checkBuffer = new byte[32];
 
                 if (!isNew)
                 {
                     // check whether bytes 32 to 64 is empty. This indicates LiteDb was unable to write encrypted 1s during last attempt.
-                    _stream.Read(checkBuffer, 0, checkBuffer.Length);
+                    _stream.Read(checkBuffer, 0, checkBufferSize);
                     isNew = checkBuffer.All(x => x == 0);
 
                     // reset checkBuffer and stream position
-                    Array.Clear(checkBuffer, 0, checkBuffer.Length);
+                    Array.Clear(checkBuffer, 0, checkBufferSize);
                     _stream.Position = 32;
                 }
 
                 // fill checkBuffer with encrypted 1 to check when open
                 if (isNew)
                 {
-                    checkBuffer.Fill(1, 0, checkBuffer.Length);
+                    checkBuffer.Fill(1, 0, checkBufferSize);
 
-                    _writer.Write(checkBuffer, 0, checkBuffer.Length);
+                    _writer.Write(checkBuffer, 0, checkBufferSize);
 
                     //ensure that the "hidden" page in encrypted files is created correctly
                     _stream.Position = PAGE_SIZE - 1;
                     _stream.WriteByte(0);
                 }
                 else
                 {
-                    _reader.Read(checkBuffer, 0, checkBuffer.Length);
+                    _reader.Read(checkBuffer, 0, checkBufferSize);
 
                     if (!checkBuffer.All(x => x == 1))
                     {
@@ -140,8 +147,7 @@ public AesStream(string password, Stream stream)
 
                 _stream.Position = PAGE_SIZE;
                 _stream.FlushToDisk();
-
-                using (var ms = new MemoryStream(new byte[16]))
+                using (var ms = new MemoryStream(msBuffer))
                 using (var tempStream = new CryptoStream(ms, _decryptor, CryptoStreamMode.Read))
                 {
                     tempStream.Read(_decryptedZeroes, 0, _decryptedZeroes.Length);
@@ -153,6 +159,11 @@ public AesStream(string password, Stream stream)
 
                 throw;
             }
+            finally
+            {
+                _bufferPool.Return(msBuffer, true);
+                _bufferPool.Return(checkBuffer, true);
+            }
         }
 
         /// <summary>

diff --git a/LiteDB/Engine/Engine/Upgrade.cs b/LiteDB/Engine/Engine/Upgrade.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
@@ -9,6 +10,9 @@ namespace LiteDB.Engine
 {
     public partial class LiteEngine
     {
+
+        private static readonly ArrayPool<byte> _bufferPool = ArrayPool<byte>.Shared;
+
         /// <summary>
         /// If Upgrade=true, run this before open Disk service
         /// </summary>
@@ -19,20 +23,23 @@ private void TryUpgrade()
             // if file not exists, just exit
             if (!File.Exists(filename)) return;
 
+            const int bufferSize = 1024;
+            var buffer = _bufferPool.Rent(bufferSize);
+
             using (var stream = new FileStream(
                 _settings.Filename,
                 FileMode.Open,
                 FileAccess.Read,
-                FileShare.Read, 1024))
+                FileShare.Read, bufferSize))
             {
-                var buffer = new byte[1024];
+
 
                 stream.Position = 0;
-                stream.Read(buffer, 0, buffer.Length);
+                stream.Read(buffer, 0, bufferSize);
 
                 if (FileReaderV7.IsVersion(buffer) == false) return;
             }
-
+            _bufferPool.Return(buffer, true);
             // run rebuild process
             this.Recovery(_settings.Collation);
         }