From 66150d4935cfcfe8c8a362950e8c6cda0676bdf3 Mon Sep 17 00:00:00 2001
From: Apolixit <friotromain@gmail.com>
Date: Tue, 16 Jul 2024 23:02:43 +0200
Subject: [PATCH] Bind OnConnectionLost on RPC Disconnected event + Add EraEnd

---
 Substrate.NetApi.Test/Extrinsic/EraTest.cs | 17 +++++++++--
 Substrate.NetApi/Model/Extrinsics/Era.cs   |  7 +++++
 Substrate.NetApi/SubstrateClient.cs        | 33 ++++++++--------------
 3 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/Substrate.NetApi.Test/Extrinsic/EraTest.cs b/Substrate.NetApi.Test/Extrinsic/EraTest.cs
index 9b8f31d..0464b24 100644
--- a/Substrate.NetApi.Test/Extrinsic/EraTest.cs
+++ b/Substrate.NetApi.Test/Extrinsic/EraTest.cs
@@ -39,10 +39,21 @@ public void EraEncodeDecodeTest()
         }
 
         [Test]
-        public void EraBeginTest()
+        [TestCase(64u, 49u, 1587u, 1585u)]
+        [TestCase(64u, 45u, 21604404u, 21604397u)]
+        public void EraBeginTest(ulong period, ulong phase, ulong currentBlock, ulong expectedBlock)
         {
-            var era = new Era(64, 49, false);
-            Assert.AreEqual(1585, era.EraStart(1587));
+            var era = new Era(period, phase, false);
+            Assert.AreEqual(expectedBlock, era.EraStart(currentBlock));
+        }
+
+        [Test]
+        [TestCase(64u, 49u, 1587u, 1649u)]
+        [TestCase(64u, 45u, 21604404u, 21604461u)]
+        public void EraEndTest(ulong period, ulong phase, ulong currentBlock, ulong expectedBlock)
+        {
+            var era = new Era(period, phase, false);
+            Assert.AreEqual(expectedBlock, era.EraEnd(currentBlock));
         }
 
         [Test]
diff --git a/Substrate.NetApi/Model/Extrinsics/Era.cs b/Substrate.NetApi/Model/Extrinsics/Era.cs
index d80ff82..f38e4bf 100644
--- a/Substrate.NetApi/Model/Extrinsics/Era.cs
+++ b/Substrate.NetApi/Model/Extrinsics/Era.cs
@@ -31,6 +31,13 @@ public class Era : IEncodable
         /// <returns></returns>
         public ulong EraStart(ulong currentBlockNumber) => IsImmortal ? 0 : (Math.Max(currentBlockNumber, Phase) - Phase) / Period * Period + Phase;
 
+        /// <summary>
+        /// Era End
+        /// </summary>
+        /// <param name="currentBlockNumber"></param>
+        /// <returns></returns>
+        public ulong EraEnd(ulong currentBlockNumber) => EraStart(currentBlockNumber) + Period;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="Era"/> class.
         /// </summary>
diff --git a/Substrate.NetApi/SubstrateClient.cs b/Substrate.NetApi/SubstrateClient.cs
index 2233910..8cf7951 100644
--- a/Substrate.NetApi/SubstrateClient.cs
+++ b/Substrate.NetApi/SubstrateClient.cs
@@ -233,8 +233,9 @@ public async Task ConnectAsync(bool useMetaData, bool standardSubstrate, Cancell
             OnConnectionSet();
 
             linkedTokenSource.Dispose();
-            //_connectTokenSource.Dispose();
-            //_connectTokenSource = null;
+            _connectTokenSource.Dispose();
+            _connectTokenSource = null;
+
             Logger.Debug("Connected to Websocket.");
 
             var formatter = new JsonMessageFormatter();
@@ -252,6 +253,13 @@ public async Task ConnectAsync(bool useMetaData, bool standardSubstrate, Cancell
             _jsonRpc = new JsonRpc(new WebSocketMessageHandler(_socket, formatter));
             _jsonRpc.TraceSource.Listeners.Add(new SerilogTraceListener.SerilogTraceListener());
             _jsonRpc.TraceSource.Switch.Level = SourceLevels.Warning;
+
+            _jsonRpc.Disconnected += (sender, args) =>
+            {
+                Logger.Debug("Disconnected from websocket.");
+                OnConnectionLost();
+            };
+
             _jsonRpc.AddLocalRpcTarget(Listener, new JsonRpcTargetOptions { AllowNonPublicInvocation = false });
             _jsonRpc.StartListening();
             Logger.Debug("Listening to websocket.");
@@ -288,24 +296,6 @@ public async Task ConnectAsync(bool useMetaData, bool standardSubstrate, Cancell
                     Logger.Warning(ex, "Could not deserialize properties on connect.");
                 }
             }
-
-            // Start a background task to monitor the WebSocket connection
-            _ = Task.Run(async () =>
-            {
-                while (IsConnected)
-                {
-                    await Task.Delay(_connectionCheckDelay);
-
-                    if (!IsConnected)
-                    {
-                        // Raise the ConnectionLost event on a separate thread
-                        ThreadPool.QueueUserWorkItem(state =>
-                        {
-                            OnConnectionLost();
-                        });
-                    }
-                }
-            });
         }
 
         /// <summary>
@@ -496,6 +486,8 @@ public async Task CloseAsync()
         /// <returns> An asynchronous result. </returns>
         public async Task CloseAsync(CancellationToken token)
         {
+            _connectTokenSource?.Cancel();
+
             await Task.Run(async () =>
             {
                 // cancel remaining request tokens
@@ -505,7 +497,6 @@ await Task.Run(async () =>
                 if (_socket != null && _socket.State == WebSocketState.Open)
                 {
                     await _socket.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None);
-                    _connectTokenSource?.Cancel();
                     _jsonRpc?.Dispose();
                     Logger.Debug("Client closed.");
                 }