Fixed websocket subscription request revitalization

JKorf · Oct 22, 2024 · 22c571a · 22c571a
1 parent c669e32
commit 22c571a
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 13 deletions.
diff --git a/Kraken.Net/Clients/SpotApi/KrakenSocketClientSpotApi.cs b/Kraken.Net/Clients/SpotApi/KrakenSocketClientSpotApi.cs
@@ -609,7 +609,12 @@ public async Task<CallResult<KrakenCancelAfterResult>> CancelAllOrdersAfterAsync
         /// <inheritdoc />
         protected override async Task<CallResult> RevitalizeRequestAsync(Subscription subscription)
         {
-            var krakenSubscription = (KrakenSubscription)subscription;
+            var krakenSubscription = subscription as KrakenSubscription;
+            if (krakenSubscription == null)
+            {
+                // Heartbeat subscription
+                return new CallResult(null);
+            }
 
             if (!krakenSubscription.TokenRequired)
                 return new CallResult(null);

diff --git a/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenBalanceSubscription.cs b/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenBalanceSubscription.cs
@@ -8,22 +8,21 @@
 
 namespace Kraken.Net.Objects.Sockets.Subscriptions.Spot
 {
-    internal class KrakenBalanceSubscription : Subscription<KrakenSocketResponseV2<KrakenSocketSubResponse>, KrakenSocketResponseV2<KrakenSocketSubResponse>>
+    internal class KrakenBalanceSubscription : KrakenSubscription
     {
         private static readonly MessagePath _typePath = MessagePath.Get().Property("type");
 
         private readonly Action<DataEvent<IEnumerable<KrakenBalanceSnapshot>>>? _snapshotHandler;
         private readonly Action<DataEvent<IEnumerable<KrakenBalanceUpdate>>> _updateHandler;
 
         private bool? _snapshot;
-        private string _token;
 
         public override HashSet<string> ListenerIdentifiers { get; set; }
 
-        public KrakenBalanceSubscription(ILogger logger, bool? snapshot, string token, Action<DataEvent<IEnumerable<KrakenBalanceSnapshot>>>? snapshotHandler, Action<DataEvent<IEnumerable<KrakenBalanceUpdate>>> updateHandler) : base(logger, false)
+        public KrakenBalanceSubscription(ILogger logger, bool? snapshot, string token, Action<DataEvent<IEnumerable<KrakenBalanceSnapshot>>>? snapshotHandler, Action<DataEvent<IEnumerable<KrakenBalanceUpdate>>> updateHandler) : base(logger, true)
         {
             _snapshot = snapshot;
-            _token = token;
+            Token = token;
 
             _snapshotHandler = snapshotHandler;
             _updateHandler = updateHandler;
@@ -50,7 +49,7 @@ public KrakenBalanceSubscription(ILogger logger, bool? snapshot, string token, A
                     {
                         Channel = "balances",
                         Snapshot = _snapshot,
-                        Token = _token
+                        Token = Token
                     }
                 }, false);
         }
@@ -66,7 +65,7 @@ public KrakenBalanceSubscription(ILogger logger, bool? snapshot, string token, A
                     {
                         Channel = "balances",
                         Snapshot = _snapshot,
-                        Token = _token
+                        Token = Token
                     }
                 }, false);
         }

diff --git a/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenOrderSubscription.cs b/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenOrderSubscription.cs
@@ -7,23 +7,22 @@
 
 namespace Kraken.Net.Objects.Sockets.Subscriptions.Spot
 {
-    internal class KrakenOrderSubscription : Subscription<KrakenSocketResponseV2<KrakenSocketSubResponse>, KrakenSocketResponseV2<KrakenSocketSubResponse>>
+    internal class KrakenOrderSubscription : KrakenSubscription
     {
         private static readonly MessagePath _typePath = MessagePath.Get().Property("type");
 
         private readonly Action<DataEvent<IEnumerable<KrakenOrderUpdate>>> _updateHandler;
 
         private bool? _snapshotOrder;
         private bool? _snapshotTrades;
-        private string _token;
 
         public override HashSet<string> ListenerIdentifiers { get; set; }
 
-        public KrakenOrderSubscription(ILogger logger, bool? snapshotOrder, bool? snapshotTrades, string token, Action<DataEvent<IEnumerable<KrakenOrderUpdate>>> updateHandler) : base(logger, false)
+        public KrakenOrderSubscription(ILogger logger, bool? snapshotOrder, bool? snapshotTrades, string token, Action<DataEvent<IEnumerable<KrakenOrderUpdate>>> updateHandler) : base(logger, true)
         {
             _snapshotOrder = snapshotOrder;
             _snapshotTrades = snapshotTrades;
-            _token = token;
+            Token = token;
 
             _updateHandler = updateHandler;
 
@@ -47,7 +46,7 @@ public KrakenOrderSubscription(ILogger logger, bool? snapshotOrder, bool? snapsh
                         Channel = "executions",
                         SnapshotOrders = _snapshotOrder,
                         SnapshotTrades = _snapshotTrades,
-                        Token = _token
+                        Token = Token
                     }
                 }, false);
         }
@@ -62,7 +61,7 @@ public KrakenOrderSubscription(ILogger logger, bool? snapshotOrder, bool? snapsh
                     Parameters = new KrakenSocketSubRequest
                     {
                         Channel = "executions",
-                        Token = _token
+                        Token = Token
                     }
                 }, false);
         }