Handle concurrent WebSockets messages (#6210)

* Test for #6169 * Use async on Subscription message processing * Fix `ScheduleAction` usages - Use async everywhere * Parametrize test * Reduce number of messages * Add test for multiple concurrent subscriptions * Add locks to avoid concurrency issues - Locks :( * Dispose semaphore * Targeted new * Reduce delays
NethermindEth · Nov 3, 2023 · 1f6baba · 1f6baba
1 parent 122f99c
commit 1f6baba
Show file tree

Hide file tree

Showing 10 changed files with 168 additions and 74 deletions.
diff --git a/src/Nethermind/Nethermind.AccountAbstraction/Subscribe/NewPendingUserOpsSubscription.cs b/src/Nethermind/Nethermind.AccountAbstraction/Subscribe/NewPendingUserOpsSubscription.cs
@@ -60,18 +60,13 @@ public class NewPendingUserOpsSubscription : Subscription
 
         private void OnNewPending(object? sender, UserOperationEventArgs e)
         {
-            ScheduleAction(() =>
+            ScheduleAction(async () =>
             {
                 JsonRpcResult result;
-                if (_includeUserOperations)
-                {
-                    result = CreateSubscriptionMessage(new { UserOperation = new UserOperationRpc(e.UserOperation), e.EntryPoint });
-                }
-                else
-                {
-                    result = CreateSubscriptionMessage(new { UserOperation = e.UserOperation.RequestId, e.EntryPoint });
-                }
-                JsonRpcDuplexClient.SendJsonRpcResult(result);
+                result = _includeUserOperations
+                    ? CreateSubscriptionMessage(new { UserOperation = new UserOperationRpc(e.UserOperation), e.EntryPoint })
+                    : CreateSubscriptionMessage(new { UserOperation = e.UserOperation.RequestId, e.EntryPoint });
+                await JsonRpcDuplexClient.SendJsonRpcResult(result);
                 if (_logger.IsTrace) _logger.Trace($"newPendingUserOperations subscription {Id} printed hash of newPendingUserOperations.");
             });
         }
@@ -89,5 +84,3 @@ public override void Dispose()
         }
     }
 }
-
-
diff --git a/src/Nethermind/Nethermind.AccountAbstraction/Subscribe/NewReceivedUserOpsSubscription.cs b/src/Nethermind/Nethermind.AccountAbstraction/Subscribe/NewReceivedUserOpsSubscription.cs
@@ -60,18 +60,13 @@ public class NewReceivedUserOpsSubscription : Subscription
 
         private void OnNewReceived(object? sender, UserOperationEventArgs e)
         {
-            ScheduleAction(() =>
+            ScheduleAction(async () =>
             {
                 JsonRpcResult result;
-                if (_includeUserOperations)
-                {
-                    result = CreateSubscriptionMessage(new { UserOperation = new UserOperationRpc(e.UserOperation), e.EntryPoint });
-                }
-                else
-                {
-                    result = CreateSubscriptionMessage(new { UserOperation = e.UserOperation.RequestId, e.EntryPoint });
-                }
-                JsonRpcDuplexClient.SendJsonRpcResult(result);
+                result = _includeUserOperations
+                    ? CreateSubscriptionMessage(new { UserOperation = new UserOperationRpc(e.UserOperation), e.EntryPoint })
+                    : CreateSubscriptionMessage(new { UserOperation = e.UserOperation.RequestId, e.EntryPoint });
+                await JsonRpcDuplexClient.SendJsonRpcResult(result);
                 if (_logger.IsTrace) _logger.Trace($"newReceivedUserOperations subscription {Id} printed hash of newReceivedUserOperations.");
             });
         }
@@ -89,5 +84,3 @@ public override void Dispose()
         }
     }
 }
-
-
diff --git a/src/Nethermind/Nethermind.JsonRpc.Test/Modules/SubscribeModuleTests.cs b/src/Nethermind/Nethermind.JsonRpc.Test/Modules/SubscribeModuleTests.cs
@@ -5,6 +5,7 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
+using System.Net.WebSockets;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
@@ -17,17 +18,16 @@
 using Nethermind.Core.Collections;
 using Nethermind.Core.Specs;
 using Nethermind.Core.Test.Builders;
-using Nethermind.Db;
-using Nethermind.Db.Blooms;
 using Nethermind.Facade.Eth;
 using Nethermind.Int256;
 using Nethermind.JsonRpc.Modules;
 using Nethermind.JsonRpc.Modules.Eth;
 using Nethermind.JsonRpc.Modules.Subscribe;
+using Nethermind.JsonRpc.WebSockets;
 using Nethermind.Logging;
 using Nethermind.Serialization.Json;
+using Nethermind.Sockets;
 using Nethermind.Specs;
-using Nethermind.State.Repositories;
 using Nethermind.Synchronization.ParallelSync;
 using Nethermind.TxPool;
 using Newtonsoft.Json;
@@ -819,6 +819,101 @@ public void NewPendingTransactionsSubscription_on_NewPending_with_includeTransac
             expectedResult.Should().Be(serialized);
         }
 
+        [TestCase(2)]
+        [TestCase(5)]
+        [TestCase(10)]
+        [Explicit("Requires a WS server running")]
+        public async Task NewPendingTransactionSubscription_multiple_fast_messages(int messages)
+        {
+            ITxPool txPool = Substitute.For<ITxPool>();
+
+            using ClientWebSocket socket = new();
+            await socket.ConnectAsync(new Uri("ws://localhost:1337/"), CancellationToken.None);
+
+            using ISocketHandler handler = new WebSocketHandler(socket, NullLogManager.Instance);
+            using JsonRpcSocketsClient client = new(
+                clientName: "TestClient",
+                handler: handler,
+                endpointType: RpcEndpoint.Ws,
+                jsonRpcProcessor: null!,
+                jsonRpcService: null!,
+                jsonRpcLocalStats: new NullJsonRpcLocalStats(),
+                jsonSerializer: new EthereumJsonSerializer()
+            );
+
+            using NewPendingTransactionsSubscription subscription = new(
+                jsonRpcDuplexClient: client,
+                txPool: txPool,
+                logManager: LimboLogs.Instance);
+
+            for (int i = 0; i < messages; i++)
+            {
+                Transaction tx = new();
+                txPool.NewPending += Raise.EventWith(new TxEventArgs(tx));
+            }
+
+            // Wait until all messages are sent
+            await Task.Delay(1_000);
+        }
+
+        [TestCase(2)]
+        [TestCase(5)]
+        [TestCase(10)]
+        [Explicit("Requires a WS server running")]
+        public async Task MultipleSubscriptions_concurrent_fast_messages(int messages)
+        {
+            using ClientWebSocket socket = new();
+            await socket.ConnectAsync(new Uri("ws://localhost:1337/"), CancellationToken.None);
+
+            using ISocketHandler handler = new WebSocketHandler(socket, NullLogManager.Instance);
+            using JsonRpcSocketsClient client = new(
+                clientName: "TestClient",
+                handler: handler,
+                endpointType: RpcEndpoint.Ws,
+                jsonRpcProcessor: null!,
+                jsonRpcService: null!,
+                jsonRpcLocalStats: new NullJsonRpcLocalStats(),
+                jsonSerializer: new EthereumJsonSerializer()
+            );
+
+            Task subA = Task.Run(() =>
+            {
+                ITxPool txPool = Substitute.For<ITxPool>();
+                using NewPendingTransactionsSubscription subscription = new(
+                    // ReSharper disable once AccessToDisposedClosure
+                    jsonRpcDuplexClient: client,
+                    txPool: txPool,
+                    logManager: LimboLogs.Instance);
+
+                for (int i = 0; i < messages; i++)
+                {
+                    Transaction tx = new();
+                    txPool.NewPending += Raise.EventWith(new TxEventArgs(tx));
+                }
+            });
+            Task subB = Task.Run(() =>
+            {
+                IBlockTree blockTree = Substitute.For<IBlockTree>();
+                using NewHeadSubscription subscription = new(
+                    // ReSharper disable once AccessToDisposedClosure
+                    jsonRpcDuplexClient: client,
+                    blockTree: blockTree,
+                    specProvider: new TestSpecProvider(new ReleaseSpec()),
+                    logManager: LimboLogs.Instance);
+
+                for (int i = 0; i < messages; i++)
+                {
+                    BlockReplacementEventArgs eventArgs = new(Build.A.Block.TestObject);
+                    blockTree.BlockAddedToMain += Raise.EventWith(eventArgs);
+                }
+            });
+
+            await Task.WhenAll(subA, subB);
+
+            // Wait until all messages are sent
+            await Task.Delay(1_000);
+        }
+
         [Test]
         public void NewHeadSubscription_with_baseFeePerGas_test()
         {

diff --git a/...Nethermind/Nethermind.JsonRpc/Modules/Subscribe/DroppedPendingTransactionsSubscription.cs b/...Nethermind/Nethermind.JsonRpc/Modules/Subscribe/DroppedPendingTransactionsSubscription.cs
@@ -26,10 +26,10 @@ public class DroppedPendingTransactionsSubscription : Subscription
 
         private void OnEvicted(object? sender, TxEventArgs e)
         {
-            ScheduleAction(() =>
+            ScheduleAction(async () =>
             {
                 JsonRpcResult result = CreateSubscriptionMessage(e.Transaction.Hash);
-                JsonRpcDuplexClient.SendJsonRpcResult(result);
+                await JsonRpcDuplexClient.SendJsonRpcResult(result);
                 if (_logger.IsTrace)
                     _logger.Trace(
                         $"DroppedPendingTransactions subscription {Id} printed hash of DroppedPendingTransaction.");

diff --git a/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/LogsSubscription.cs b/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/LogsSubscription.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Threading.Tasks;
 using Nethermind.Blockchain;
 using Nethermind.Blockchain.Filters;
 using Nethermind.Blockchain.Find;
@@ -62,10 +63,10 @@ private void OnReceiptsInserted(object? sender, ReceiptsEventArgs e)
 
         private void TryPublishReceiptsInBackground(BlockHeader blockHeader, Func<TxReceipt[]> getReceipts, string eventName, bool removed)
         {
-            ScheduleAction(() => TryPublishEvent(blockHeader, getReceipts(), eventName, removed));
+            ScheduleAction(async () => await TryPublishEvent(blockHeader, getReceipts(), eventName, removed));
         }
 
-        private void TryPublishEvent(BlockHeader blockHeader, TxReceipt[] receipts, string eventName, bool removed)
+        private async Task TryPublishEvent(BlockHeader blockHeader, TxReceipt[] receipts, string eventName, bool removed)
         {
             BlockHeader fromBlock = _blockTree.FindHeader(_filter.FromBlock);
             BlockHeader toBlock = _blockTree.FindHeader(_filter.ToBlock, true);
@@ -80,7 +81,7 @@ private void TryPublishEvent(BlockHeader blockHeader, TxReceipt[] receipts, stri
                 foreach (var filterLog in filterLogs)
                 {
                     JsonRpcResult result = CreateSubscriptionMessage(filterLog);
-                    JsonRpcDuplexClient.SendJsonRpcResult(result);
+                    await JsonRpcDuplexClient.SendJsonRpcResult(result);
                     if (_logger.IsTrace) _logger.Trace($"Logs subscription {Id} printed new log.");
                 }
             }

diff --git a/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/NewHeadSubscription.cs b/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/NewHeadSubscription.cs
@@ -36,11 +36,11 @@ public class NewHeadSubscription : Subscription
 
         private void OnBlockAddedToMain(object? sender, BlockReplacementEventArgs e)
         {
-            ScheduleAction(() =>
+            ScheduleAction(async () =>
             {
                 JsonRpcResult result = CreateSubscriptionMessage(new BlockForRpc(e.Block, _includeTransactions, _specProvider));
 
-                JsonRpcDuplexClient.SendJsonRpcResult(result);
+                await JsonRpcDuplexClient.SendJsonRpcResult(result);
                 if (_logger.IsTrace) _logger.Trace($"NewHeads subscription {Id} printed new block");
             });
         }

diff --git a/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/NewPendingTransactionsSubscription.cs b/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/NewPendingTransactionsSubscription.cs
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: LGPL-3.0-only
 
 using System;
+using System.Threading.Tasks;
 using Nethermind.JsonRpc.Data;
 using Nethermind.JsonRpc.Modules.Eth;
 using Nethermind.Logging;
@@ -31,10 +32,10 @@ public class NewPendingTransactionsSubscription : Subscription
 
         private void OnNewPending(object? sender, TxEventArgs e)
         {
-            ScheduleAction(() =>
+            ScheduleAction(async () =>
             {
                 JsonRpcResult result = CreateSubscriptionMessage(_includeTransactions ? new TransactionForRpc(e.Transaction) : e.Transaction.Hash);
-                JsonRpcDuplexClient.SendJsonRpcResult(result);
+                await JsonRpcDuplexClient.SendJsonRpcResult(result);
                 if (_logger.IsTrace) _logger.Trace($"NewPendingTransactions subscription {Id} printed hash of NewPendingTransaction.");
             });
         }

diff --git a/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/Subscription.cs b/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/Subscription.cs
@@ -22,7 +22,7 @@ protected Subscription(IJsonRpcDuplexClient jsonRpcDuplexClient)
         public string Id { get; }
         public abstract string Type { get; }
         public IJsonRpcDuplexClient JsonRpcDuplexClient { get; }
-        private Channel<Action> SendChannel { get; } = Channel.CreateUnbounded<Action>(new UnboundedChannelOptions() { SingleReader = true });
+        private Channel<Func<Task>> SendChannel { get; } = Channel.CreateUnbounded<Func<Task>>(new UnboundedChannelOptions { SingleReader = true });
 
         public virtual void Dispose()
         {
@@ -42,7 +42,7 @@ protected JsonRpcResult CreateSubscriptionMessage(object result)
                 }, default);
         }
 
-        protected void ScheduleAction(Action action)
+        protected void ScheduleAction(Func<Task> action)
         {
             SendChannel.Writer.TryWrite(action);
         }
@@ -55,11 +55,11 @@ private void ProcessMessages()
             {
                 while (await SendChannel.Reader.WaitToReadAsync())
                 {
-                    while (SendChannel.Reader.TryRead(out Action action))
+                    while (SendChannel.Reader.TryRead(out Func<Task> action))
                     {
                         try
                         {
-                            action();
+                            await action();
                         }
                         catch (Exception e)
                         {

diff --git a/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/SyncingSubscription.cs b/src/Nethermind/Nethermind.JsonRpc/Modules/Subscribe/SyncingSubscription.cs
@@ -46,7 +46,7 @@ private class SubscriptionSyncingResult
 
         private void OnConditionsChange(object? sender, BlockEventArgs e)
         {
-            ScheduleAction(() =>
+            ScheduleAction(async () =>
             {
                 SyncingResult syncingResult = _ethSyncingInfo.GetFullInfo();
                 bool isSyncing = syncingResult.IsSyncing;
@@ -78,7 +78,7 @@ private void OnConditionsChange(object? sender, BlockEventArgs e)
                 }
 
 
-                JsonRpcDuplexClient.SendJsonRpcResult(result);
+                await JsonRpcDuplexClient.SendJsonRpcResult(result);
                 _logger.Trace($"Syncing subscription {Id} printed SyncingResult object.");
             });
         }