From 5d0da861992f9fb8921eac626a702b4ecc164f18 Mon Sep 17 00:00:00 2001
From: Benedikt Schenkel <Benedikt.Schenkel@systemslab21.com>
Date: Mon, 18 Dec 2023 10:46:04 +0100
Subject: [PATCH] fix receive timeout

---
 .../PointToInterlockingConnection.cs              | 15 +++++++++++----
 .../PointToInterlockingConnection.cs              | 14 +++++++++++---
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/FieldElementSubsystems/Connections/EulynxBaseline4R1/PointToInterlockingConnection.cs b/src/FieldElementSubsystems/Connections/EulynxBaseline4R1/PointToInterlockingConnection.cs
index 24f776c..1ad4b22 100644
--- a/src/FieldElementSubsystems/Connections/EulynxBaseline4R1/PointToInterlockingConnection.cs
+++ b/src/FieldElementSubsystems/Connections/EulynxBaseline4R1/PointToInterlockingConnection.cs
@@ -52,7 +52,7 @@ public void Connect(IConnection connection)
 
     public async Task<bool> InitializeConnection(GenericPointState state, bool observeAbilityToMove, CancellationToken cancellationToken)
     {
-        if (await ReceiveMessage<PointPdiVersionCheckCommand>(cancellationToken) == null)
+        if (await ReceiveMessageWithTimeout<PointPdiVersionCheckCommand>(cancellationToken) == null)
         {
             _logger.LogError("Unexpected message.");
             return false;
@@ -61,7 +61,7 @@ public async Task<bool> InitializeConnection(GenericPointState state, bool obser
         var versionCheckResponse = new PointPdiVersionCheckMessage(_localId, _remoteId, PointPdiVersionCheckMessageResultPdiVersionCheck.PDIVersionsFromReceiverAndSenderDoMatch, /* TODO */ 0, 0, Array.Empty<byte>());
         await SendMessage(versionCheckResponse);
 
-        if (await ReceiveMessage<PointInitialisationRequestCommand>(cancellationToken) == null)
+        if (await ReceiveMessageWithTimeout<PointInitialisationRequestCommand>(cancellationToken) == null)
         {
             _logger.LogError("Unexpected message.");
             return false;
@@ -123,12 +123,19 @@ private async Task SendMessage(byte[] message)
         await CurrentConnection.SendAsync(message);
     }
 
-    private async Task<T> ReceiveMessage<T>(CancellationToken cancellationToken) where T : Message
+    private async Task<T> ReceiveMessageWithTimeout<T>(CancellationToken cancellationToken) where T : Message
     {
         if (CurrentConnection == null) throw new InvalidOperationException("Connection is null. Did you call Connect()?");
         ResetTimeout();
+        var token = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _timeout.Token).Token;
+        return await ReceiveMessage<T>(token);
+    }
+
+    private async Task<T> ReceiveMessage<T>(CancellationToken cancellationToken) where T : Message
+    {
+        if (CurrentConnection == null) throw new InvalidOperationException("Connection is null. Did you call Connect()?");
 
-        var message = Message.FromBytes(await CurrentConnection.ReceiveAsync(CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _timeout.Token).Token));
+        var message = Message.FromBytes(await CurrentConnection.ReceiveAsync(cancellationToken));
         if (message is T tMessage) return tMessage;
         _logger.LogError("Unexpected message: {}", message);
         throw new InvalidOperationException("Unexpected message.");
diff --git a/src/FieldElementSubsystems/Connections/EulynxBaseline4R2/PointToInterlockingConnection.cs b/src/FieldElementSubsystems/Connections/EulynxBaseline4R2/PointToInterlockingConnection.cs
index 1ee35c9..d9da7ea 100644
--- a/src/FieldElementSubsystems/Connections/EulynxBaseline4R2/PointToInterlockingConnection.cs
+++ b/src/FieldElementSubsystems/Connections/EulynxBaseline4R2/PointToInterlockingConnection.cs
@@ -51,7 +51,7 @@ public void Connect(IConnection connection)
 
     public async Task<bool> InitializeConnection(GenericPointState state, bool observeAbilityToMove, CancellationToken cancellationToken)
     {
-        if (await ReceiveMessage<PointPdiVersionCheckCommand>(cancellationToken) == null)
+        if (await ReceiveMessageWithTimeout<PointPdiVersionCheckCommand>(cancellationToken) == null)
         {
             _logger.LogError("Unexpected message.");
             return false;
@@ -60,7 +60,7 @@ public async Task<bool> InitializeConnection(GenericPointState state, bool obser
         var versionCheckResponse = new PointPdiVersionCheckMessage(_localId, _remoteId, PointPdiVersionCheckMessageResultPdiVersionCheck.PDIVersionsFromReceiverAndSenderDoMatch, /* TODO */ 0, 0, new byte[] { });
         await SendMessage(versionCheckResponse);
 
-        if (await ReceiveMessage<PointInitialisationRequestCommand>(cancellationToken) == null)
+        if (await ReceiveMessageWithTimeout<PointInitialisationRequestCommand>(cancellationToken) == null)
         {
             _logger.LogError("Unexpected message.");
             return false;
@@ -122,12 +122,20 @@ private async Task SendMessage(byte[] message)
         await CurrentConnection.SendAsync(message);
     }
 
+    private async Task<T> ReceiveMessageWithTimeout<T>(CancellationToken cancellationToken) where T : Message
+    {
+        if (CurrentConnection == null) throw new InvalidOperationException("Connection is null. Did you call Connect()?");
+        ResetTimeout();
+        var token = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _timeout.Token).Token;
+        return await ReceiveMessage<T>(token);
+    }
+
     private async Task<T> ReceiveMessage<T>(CancellationToken cancellationToken) where T : Message
     {
         if (CurrentConnection == null) throw new InvalidOperationException("Connection is null. Did you call Connect()?");
         ResetTimeout();
 
-        var message = Message.FromBytes(await CurrentConnection.ReceiveAsync(CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _timeout.Token).Token));
+        var message = Message.FromBytes(await CurrentConnection.ReceiveAsync(cancellationToken));
         if (message is T tMessage) return tMessage;
         _logger.LogError("Unexpected message: {}", message);
         throw new InvalidOperationException("Unexpected message.");