Fix behavior when connection drops (#56)

* add a method to check if the TCP connection is open;
dispose of the TransportService correctly in case of disconnections;

Signed-off-by: Níckolas Goline <nickolas.goline+github@gmail.com>

* fix tests

Signed-off-by: Níckolas Goline <nickolas.goline+github@gmail.com>

---------

Signed-off-by: Níckolas Goline <nickolas.goline+github@gmail.com>

Author: nGoline

src/NLightning.Application/Node/Managers/PeerManager.cs

@@ -89,7 +89,14 @@ public async Task StopAsync()
             throw new InvalidOperationException($"{nameof(PeerManager)} is not running");
 
         foreach (var peerKey in _peers.Keys)
-            DisconnectPeer(peerKey);
+            try
+            {
+                DisconnectPeer(peerKey);
+            }
+            catch (Exception e)
+            {
+                _logger.LogWarning(e, "Error disconnecting peer {Peer}", peerKey);
+            }
 
         try
         {
@@ -174,10 +181,11 @@ private void HandleNewPeerConnected(object? _, NewPeerConnectedEventArgs args)
         {
             // Create the peer
             var peerService = _peerServiceFactory.CreateConnectingPeerAsync(args.TcpClient).GetAwaiter().GetResult();
-
             peerService.OnDisconnect += HandlePeerDisconnection;
             peerService.OnChannelMessageReceived += HandlePeerChannelMessage;
 
+            _logger.LogTrace("PeerService created for peer {PeerPubKey}", peerService.PeerPubKey);
+
             var peer = new PeerModel(peerService.PeerPubKey, args.Host, args.Port)
             {
                 LastSeenAt = DateTime.UtcNow
@@ -192,12 +200,23 @@ private void HandleNewPeerConnected(object? _, NewPeerConnectedEventArgs args)
         }
     }
 
-    private void HandlePeerDisconnection(object? _, PeerDisconnectedEventArgs args)
+    private void HandlePeerDisconnection(object? sender, PeerDisconnectedEventArgs args)
     {
         ArgumentNullException.ThrowIfNull(args);
 
         _peers.Remove(args.PeerPubKey);
         _logger.LogInformation("Peer {Peer} disconnected", args.PeerPubKey);
+
+        if (sender is IPeerService peerService)
+        {
+            peerService.OnDisconnect -= HandlePeerDisconnection;
+            peerService.OnChannelMessageReceived -= HandlePeerChannelMessage;
+        }
+        else
+        {
+            _logger.LogWarning("Peer {Peer} disconnected, but we were unable to detach event handlers",
+                               args.PeerPubKey);
+        }
     }
 
     private void HandlePeerChannelMessage(object? _, ChannelMessageEventArgs args)

src/NLightning.Domain/Protocol/Interfaces/IMessageService.cs

@@ -12,9 +12,11 @@ public interface IMessageService : IDisposable
     /// Sends a message.
     /// </summary>
     /// <param name="message">The message to send.</param>
+    /// <param name="throwOnException">If true, exceptions will be thrown instead of being raised via the OnExceptionRaised event.</param>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>A task that represents the asynchronous operation.</returns>
-    Task SendMessageAsync(IMessage message, CancellationToken cancellationToken = default);
+    Task SendMessageAsync(IMessage message, bool throwOnException = false,
+                          CancellationToken cancellationToken = default);
 
     /// <summary>
     /// Event that is raised when a message is received.

src/NLightning.Infrastructure/Node/Factories/PeerServiceFactory.cs

@@ -59,6 +59,7 @@ public async Task<IPeerService> CreateConnectedPeerAsync(CompactPubKey peerPubKe
         }
         catch (Exception ex)
         {
+            transportService.Dispose();
             throw new ConnectionException($"Error connecting to peer {peerPubKey}", ex);
         }
 
@@ -100,6 +101,7 @@ public async Task<IPeerService> CreateConnectingPeerAsync(TcpClient tcpClient)
         }
         catch (Exception ex)
         {
+            transportService.Dispose();
             throw new ConnectionException($"Error establishing connection to peer {ipAddress}:{port}", ex);
         }
 

src/NLightning.Infrastructure/Node/Services/PeerCommunicationService.cs

@@ -18,13 +18,16 @@ namespace NLightning.Infrastructure.Node.Services;
 /// </summary>
 public class PeerCommunicationService : IPeerCommunicationService
 {
-    private readonly CancellationTokenSource _cancellationTokenSource = new();
+    private readonly CancellationTokenSource _cts = new();
     private readonly ILogger<PeerCommunicationService> _logger;
     private readonly IMessageService _messageService;
     private readonly IPingPongService _pingPongService;
     private readonly IServiceProvider _serviceProvider;
     private readonly IMessageFactory _messageFactory;
+    private readonly TaskCompletionSource<bool> _pingPongTcs = new();
+
     private bool _isInitialized;
+    private CancellationTokenSource? _initWaitCancellationTokenSource;
 
     /// <inheritdoc />
     public event EventHandler<IMessage?>? MessageReceived;
@@ -69,16 +72,11 @@ public PeerCommunicationService(ILogger<PeerCommunicationService> logger, IMessa
     /// <inheritdoc />
     public async Task InitializeAsync(TimeSpan networkTimeout)
     {
-        // Always send an init message upon connection
-        _logger.LogTrace("Sending init message to peer {peer}", PeerCompactPubKey);
-        var initMessage = _messageFactory.CreateInitMessage();
-        await _messageService.SendMessageAsync(initMessage, _cancellationTokenSource.Token);
-
-        // Wait for an init message
         _logger.LogTrace("Waiting for init message from peer {peer}", PeerCompactPubKey);
 
         // Set timeout to close connection if the other peer doesn't send an init message
-        _ = Task.Delay(networkTimeout, _cancellationTokenSource.Token).ContinueWith(task =>
+        _initWaitCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(_cts.Token);
+        _ = Task.Delay(networkTimeout, _initWaitCancellationTokenSource.Token).ContinueWith(task =>
         {
             if (!task.IsCanceled && !_isInitialized)
             {
@@ -87,21 +85,35 @@ public async Task InitializeAsync(TimeSpan networkTimeout)
             }
         });
 
-        if (!_messageService.IsConnected)
+        // Always send an init message upon connection
+        _logger.LogTrace("Sending init message to peer {peer}", PeerCompactPubKey);
+        var initMessage = _messageFactory.CreateInitMessage();
+        try
         {
-            throw new ConnectionException($"Failed to connect to peer {PeerCompactPubKey}");
+            await _messageService.SendMessageAsync(initMessage, true, _cts.Token);
+        }
+        catch (Exception e)
+        {
+            _pingPongTcs.TrySetResult(true);
+            throw new ConnectionException($"Failed to send init message to peer {PeerCompactPubKey}", e);
         }
 
         // Set up ping service to keep connection alive
-        SetupPingPongService();
+        if (!_cts.IsCancellationRequested)
+        {
+            if (!_messageService.IsConnected)
+                throw new ConnectionException($"Failed to connect to peer {PeerCompactPubKey}");
+
+            SetupPingPongService();
+        }
     }
 
     /// <inheritdoc />
     public async Task SendMessageAsync(IMessage message, CancellationToken cancellationToken = default)
     {
         try
         {
-            await _messageService.SendMessageAsync(message, cancellationToken);
+            await _messageService.SendMessageAsync(message, cancellationToken: cancellationToken);
         }
         catch (Exception ex)
         {
@@ -112,11 +124,18 @@ public async Task SendMessageAsync(IMessage message, CancellationToken cancellat
     /// <inheritdoc />
     public void Disconnect()
     {
-        _logger.LogInformation("Disconnecting peer {peer}", PeerCompactPubKey);
-        _cancellationTokenSource.Cancel();
-        _messageService.Dispose();
-
-        DisconnectEvent?.Invoke(this, EventArgs.Empty);
+        try
+        {
+            _ = _cts.CancelAsync();
+            _logger.LogTrace("Waiting for ping service to stop for peer {peer}", PeerCompactPubKey);
+            _pingPongTcs.Task.Wait(TimeSpan.FromSeconds(5));
+            _logger.LogTrace("Ping service stopped for peer {peer}", PeerCompactPubKey);
+        }
+        finally
+        {
+            _messageService.Dispose();
+            DisconnectEvent?.Invoke(this, EventArgs.Empty);
+        }
     }
 
     private void SetupPingPongService()
@@ -125,13 +144,10 @@ private void SetupPingPongService()
         _pingPongService.OnPongReceived += HandlePongReceived;
 
         // Setup Ping to keep connection alive
-        _ = _pingPongService.StartPingAsync(_cancellationTokenSource.Token).ContinueWith(task =>
+        _ = _pingPongService.StartPingAsync(_cts.Token).ContinueWith(_ =>
         {
-            if (task.IsFaulted)
-            {
-                RaiseException(new ConnectionException($"Failed to start ping service for peer {PeerCompactPubKey}",
-                                                       task.Exception));
-            }
+            _logger.LogTrace("Ping service stopped for peer {peer}, setting result", PeerCompactPubKey);
+            _pingPongTcs.TrySetResult(true);
         });
 
         _logger.LogInformation("Ping service started for peer {peer}", PeerCompactPubKey);
@@ -154,7 +170,7 @@ private void HandlePingMessageReady(object? sender, IMessage pingMessage)
 
         try
         {
-            _messageService.SendMessageAsync(pingMessage, _cancellationTokenSource.Token).GetAwaiter().GetResult();
+            _messageService.SendMessageAsync(pingMessage, cancellationToken: _cts.Token).GetAwaiter().GetResult();
         }
         catch (Exception ex)
         {
@@ -172,6 +188,7 @@ private void HandleMessageReceived(object? sender, IMessage? message)
         if (!_isInitialized && message.Type == MessageTypes.Init)
         {
             _isInitialized = true;
+            _initWaitCancellationTokenSource?.Cancel();
         }
 
         // Forward the message to subscribers
@@ -206,6 +223,7 @@ private void RaiseException(Exception exception)
         {
             ChannelId? channelId = null;
             var message = errorException.Message;
+            mustDisconnect = true;
 
             if (errorException is ChannelErrorException channelErrorException)
             {
@@ -214,8 +232,16 @@ private void RaiseException(Exception exception)
                     message = channelErrorException.PeerMessage;
             }
 
-            _messageService.SendMessageAsync(new ErrorMessage(new ErrorPayload(channelId, message)));
-            mustDisconnect = true;
+            if (errorException is not ConnectionException)
+            {
+                _logger.LogTrace("Sending error message to peer {peer}. ChannelId: {channelId}, Message: {message}",
+                                 PeerCompactPubKey, channelId, message);
+
+                _ = Task.Run(() => _messageService.SendMessageAsync(
+                                 new ErrorMessage(new ErrorPayload(channelId, message))));
+
+                return;
+            }
         }
         else if (exception is WarningException warningException)
         {
@@ -229,17 +255,25 @@ private void RaiseException(Exception exception)
                     message = channelWarningException.PeerMessage;
             }
 
-            _messageService.SendMessageAsync(new WarningMessage(new ErrorPayload(channelId, message)));
-        }
+            _logger.LogTrace("Sending warning message to peer {peer}. ChannelId: {channelId}, Message: {message}",
+                             PeerCompactPubKey, channelId, message);
 
-        _logger.LogError(exception, "Exception occurred with peer {peer}. {exceptionMessage}", PeerCompactPubKey,
-                         exception.Message);
+            _ = Task.Run(() => _messageService.SendMessageAsync(
+                             new WarningMessage(new ErrorPayload(channelId, message))));
+        }
 
         // Forward the exception to subscribers
         ExceptionRaised?.Invoke(this, exception);
 
         // Disconnect if not already disconnecting
-        if (mustDisconnect && !_cancellationTokenSource.IsCancellationRequested)
+        if (mustDisconnect && !_cts.IsCancellationRequested)
+        {
+            _messageService.OnMessageReceived -= HandleMessageReceived;
+            _messageService.OnExceptionRaised -= HandleExceptionRaised;
+            _pingPongService.DisconnectEvent -= HandleExceptionRaised;
+
+            _logger.LogWarning(exception, "We're disconnecting peer {peer} because of an exception", PeerCompactPubKey);
             Disconnect();
+        }
     }
 }

src/NLightning.Infrastructure/Node/Services/PeerService.cs

@@ -4,6 +4,7 @@
 namespace NLightning.Infrastructure.Node.Services;
 
 using Domain.Crypto.ValueObjects;
+using Domain.Exceptions;
 using Domain.Node.Events;
 using Domain.Node.Interfaces;
 using Domain.Node.Options;
@@ -56,7 +57,18 @@ public PeerService(IPeerCommunicationService peerCommunicationService, FeatureOp
         _peerCommunicationService.DisconnectEvent += HandleDisconnection;
 
         // Initialize communication
-        _peerCommunicationService.InitializeAsync(networkTimeout).GetAwaiter().GetResult();
+        try
+        {
+            _peerCommunicationService.InitializeAsync(networkTimeout).GetAwaiter().GetResult();
+        }
+        catch (Exception e)
+        {
+            _peerCommunicationService.MessageReceived -= HandleMessage;
+            _peerCommunicationService.ExceptionRaised -= HandleException;
+            _peerCommunicationService.DisconnectEvent -= HandleDisconnection;
+
+            throw new ErrorException("Error initializing peer communication", e);
+        }
     }
 
     /// <summary>
@@ -83,7 +95,6 @@ private void HandleMessage(object? sender, IMessage? message)
 
         if (!_isInitialized)
         {
-            _logger.LogTrace("Received message from peer {peer} but was not initialized", PeerPubKey);
             HandleInitialization(message);
         }
         else if (message is IChannelMessage channelMessage)
@@ -103,6 +114,7 @@ private void HandleException(object? sender, Exception e)
 
     private void HandleDisconnection(object? sender, EventArgs e)
     {
+        _logger.LogTrace("Handling disconnection for peer {Peer}", PeerPubKey);
         OnDisconnect?.Invoke(this, new PeerDisconnectedEventArgs(PeerPubKey));
     }