Use ArrayPool for Discovery Deserialization (#7427)

NethermindEth · Sep 13, 2024 · c5e5d4a · c5e5d4a
1 parent 969c4b9
commit c5e5d4a
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 32 deletions.
diff --git a/src/Nethermind/Nethermind.Core/Buffers/ArrayPoolDisposableReturn.cs b/src/Nethermind/Nethermind.Core/Buffers/ArrayPoolDisposableReturn.cs
@@ -0,0 +1,22 @@
+// SPDX-FileCopyrightText: 2024 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System;
+using System.Buffers;
+
+namespace Nethermind.Core.Buffers;
+
+public readonly struct ArrayPoolDisposableReturn : IDisposable
+{
+    private readonly byte[] _array;
+
+    private ArrayPoolDisposableReturn(byte[] array) => _array = array;
+
+    public static ArrayPoolDisposableReturn Rent(int size, out byte[] array)
+    {
+        array = ArrayPool<byte>.Shared.Rent(size);
+        return new(array);
+    }
+
+    public void Dispose() => ArrayPool<byte>.Shared.Return(_array);
+}
diff --git a/src/Nethermind/Nethermind.Db/SimpleFilePublicKeyDb.cs b/src/Nethermind/Nethermind.Db/SimpleFilePublicKeyDb.cs
@@ -11,6 +11,7 @@
 using Microsoft.Win32.SafeHandles;
 
 using Nethermind.Core;
+using Nethermind.Core.Buffers;
 using Nethermind.Core.Collections;
 using Nethermind.Core.Extensions;
 using Nethermind.Logging;
@@ -208,7 +209,7 @@ private void LoadData()
 
             using SafeFileHandle fileHandle = File.OpenHandle(DbPath, FileMode.OpenOrCreate);
 
-            byte[] rentedBuffer = ArrayPool<byte>.Shared.Rent(maxLineLength);
+            using var handle = ArrayPoolDisposableReturn.Rent(maxLineLength, out byte[] rentedBuffer);
             int read = RandomAccess.Read(fileHandle, rentedBuffer, 0);
 
             long offset = 0L;
@@ -270,7 +271,6 @@ private void LoadData()
                 read = RandomAccess.Read(fileHandle, rentedBuffer.AsSpan(bytes.Length), offset);
             }
 
-            ArrayPool<byte>.Shared.Return(rentedBuffer);
             if (bytes.Length > 0)
             {
                 if (_logger.IsWarn) _logger.Warn($"Malformed {Name}. Ignoring...");

diff --git a/src/Nethermind/Nethermind.Evm/Tracing/GethStyle/Custom/JavaScript/Db.cs b/src/Nethermind/Nethermind.Evm/Tracing/GethStyle/Custom/JavaScript/Db.cs
@@ -10,6 +10,7 @@
 using Microsoft.ClearScript.JavaScript;
 using Microsoft.ClearScript.V8;
 using Nethermind.Core;
+using Nethermind.Core.Buffers;
 using Nethermind.Core.Crypto;
 using Nethermind.Int256;
 using Nethermind.State;
@@ -30,21 +31,15 @@ public class Db
 
     public ITypedArray<byte> getState(object address, object hash)
     {
-        byte[] array = ArrayPool<byte>.Shared.Rent(32);
-        try
-        {
-            ReadOnlySpan<byte> bytes = WorldState.Get(new StorageCell(address.ToAddress(), hash.GetHash()));
-            if (bytes.Length < array.Length)
-            {
-                Array.Clear(array);
-            }
-            bytes.CopyTo(array.AsSpan(array.Length - bytes.Length));
-            return array.ToTypedScriptArray();
-        }
-        finally
+        using var handle = ArrayPoolDisposableReturn.Rent(32, out byte[] array);
+
+        ReadOnlySpan<byte> bytes = WorldState.Get(new StorageCell(address.ToAddress(), hash.GetHash()));
+        if (bytes.Length < array.Length)
         {
-            ArrayPool<byte>.Shared.Return(array);
+            Array.Clear(array);
         }
+        bytes.CopyTo(array.AsSpan(array.Length - bytes.Length));
+        return array.ToTypedScriptArray();
     }
 
     public bool exists(object address) => WorldState.TryGetAccount(address.ToAddress(), out AccountStruct account) && !account.IsTotallyEmpty;

diff --git a/src/Nethermind/Nethermind.Network.Discovery/NettyDiscoveryHandler.cs b/src/Nethermind/Nethermind.Network.Discovery/NettyDiscoveryHandler.cs
@@ -1,13 +1,15 @@
 // SPDX-FileCopyrightText: 2022 Demerzel Solutions Limited
 // SPDX-License-Identifier: LGPL-3.0-only
 
+using System.Buffers;
 using System.Net;
 using System.Net.Sockets;
 using DotNetty.Buffers;
 using DotNetty.Transport.Channels;
 using DotNetty.Transport.Channels.Sockets;
 using FastEnumUtility;
 using Nethermind.Core;
+using Nethermind.Core.Buffers;
 using Nethermind.Core.Extensions;
 using Nethermind.Logging;
 using Nethermind.Network.Discovery.Messages;
@@ -116,21 +118,22 @@ private bool TryParseMessage(DatagramPacket packet, out DiscoveryMsg? msg)
         EndPoint address = packet.Sender;
 
         int size = content.ReadableBytes;
-        byte[] msgBytes = new byte[size];
-        content.ReadBytes(msgBytes);
+        using var handle = ArrayPoolDisposableReturn.Rent(size, out byte[] msgBytes);
 
-        Interlocked.Add(ref Metrics.DiscoveryBytesReceived, msgBytes.Length);
+        content.ReadBytes(msgBytes, 0, size);
+
+        Interlocked.Add(ref Metrics.DiscoveryBytesReceived, size);
 
         if (msgBytes.Length < 98)
         {
-            if (_logger.IsDebug) _logger.Debug($"Incorrect discovery message, length: {msgBytes.Length}, sender: {address}");
+            if (_logger.IsDebug) _logger.Debug($"Incorrect discovery message, length: {size}, sender: {address}");
             return false;
         }
 
         byte typeRaw = msgBytes[97];
         if (!FastEnum.IsDefined<MsgType>((int)typeRaw))
         {
-            if (_logger.IsDebug) _logger.Debug($"Unsupported message type: {typeRaw}, sender: {address}, message {msgBytes.ToHexString()}");
+            if (_logger.IsDebug) _logger.Debug($"Unsupported message type: {typeRaw}, sender: {address}, message {msgBytes.AsSpan(0, size).ToHexString()}");
             return false;
         }
 
@@ -139,12 +142,12 @@ private bool TryParseMessage(DatagramPacket packet, out DiscoveryMsg? msg)
 
         try
         {
-            msg = Deserialize(type, msgBytes);
+            msg = Deserialize(type, new ArraySegment<byte>(msgBytes, 0, size));
             msg.FarAddress = (IPEndPoint)address;
         }
         catch (Exception e)
         {
-            if (_logger.IsDebug) _logger.Debug($"Error during deserialization of the message, type: {type}, sender: {address}, msg: {msgBytes.ToHexString()}, {e.Message}");
+            if (_logger.IsDebug) _logger.Debug($"Error during deserialization of the message, type: {type}, sender: {address}, msg: {msgBytes.AsSpan(0, size).ToHexString()}, {e.Message}");
             return false;
         }
 
@@ -185,7 +188,7 @@ protected override void ChannelRead0(IChannelHandlerContext ctx, DatagramPacket
         }
     }
 
-    private DiscoveryMsg Deserialize(MsgType type, byte[] msg)
+    private DiscoveryMsg Deserialize(MsgType type, ArraySegment<byte> msg)
     {
         return type switch
         {

diff --git a/src/Nethermind/Nethermind.Network/IMessageSerializationService.cs b/src/Nethermind/Nethermind.Network/IMessageSerializationService.cs
@@ -1,6 +1,7 @@
 // SPDX-FileCopyrightText: 2022 Demerzel Solutions Limited
 // SPDX-License-Identifier: LGPL-3.0-only
 
+using System;
 using System.Reflection;
 using DotNetty.Buffers;
 
@@ -9,7 +10,7 @@ namespace Nethermind.Network
     public interface IMessageSerializationService
     {
         IByteBuffer ZeroSerialize<T>(T message, AbstractByteBufferAllocator? allocator = null) where T : MessageBase;
-        T Deserialize<T>(byte[] bytes) where T : MessageBase;
+        T Deserialize<T>(ArraySegment<byte> bytes) where T : MessageBase;
         T Deserialize<T>(IByteBuffer buffer) where T : MessageBase;
         void Register(Assembly assembly);
         void Register<T>(IZeroMessageSerializer<T> messageSerializer) where T : MessageBase;

diff --git a/src/Nethermind/Nethermind.Network/MessageSerializationService.cs b/src/Nethermind/Nethermind.Network/MessageSerializationService.cs
@@ -14,13 +14,13 @@ public class MessageSerializationService : IMessageSerializationService
     {
         private readonly ConcurrentDictionary<RuntimeTypeHandle, object> _zeroSerializers = new ConcurrentDictionary<RuntimeTypeHandle, object>();
 
-        public T Deserialize<T>(byte[] bytes) where T : MessageBase
+        public T Deserialize<T>(ArraySegment<byte> bytes) where T : MessageBase
         {
             if (!TryGetZeroSerializer(out IZeroMessageSerializer<T> zeroMessageSerializer))
                 throw new InvalidOperationException($"No {nameof(IZeroMessageSerializer<T>)} registered for {typeof(T).Name}.");
 
-            IByteBuffer byteBuffer = PooledByteBufferAllocator.Default.Buffer(bytes.Length);
-            byteBuffer.WriteBytes(bytes);
+            IByteBuffer byteBuffer = PooledByteBufferAllocator.Default.Buffer(bytes.Count);
+            byteBuffer.WriteBytes(bytes.Array, bytes.Offset, bytes.Count);
             try
             {
                 return zeroMessageSerializer.Deserialize(byteBuffer);

diff --git a/src/Nethermind/Nethermind.Serialization.Json/JavaScriptObjectConverter.cs b/src/Nethermind/Nethermind.Serialization.Json/JavaScriptObjectConverter.cs
@@ -2,15 +2,14 @@
 // SPDX-License-Identifier: LGPL-3.0-only
 
 using System;
-using System.Buffers;
 using System.Collections.Generic;
 using System.Globalization;
 using System.Numerics;
 using System.Text.Json;
 using System.Text.Json.Serialization;
-
 using Microsoft.ClearScript;
 using Microsoft.ClearScript.JavaScript;
+using Nethermind.Core.Buffers;
 
 #nullable enable
 
@@ -60,12 +59,10 @@ public override void Write(Utf8JsonWriter writer, IJavaScriptObject o, JsonSeria
                 return;
             }
 
-            byte[] array = ArrayPool<byte>.Shared.Rent(size);
+            using var handle = ArrayPoolDisposableReturn.Rent(size, out byte[] array);
 
             buffer.ReadBytes(buffer.Offset, buffer.Size, array, 0);
             ByteArrayConverter.Convert(writer, array.AsSpan(0, size), skipLeadingZeros: false);
-
-            ArrayPool<byte>.Shared.Return(array);
         }
         else
         {