Add subscription removal on cancellation

Author: adam-fowler

Sources/Valkey/Connection/ValkeyChannelHandler.swift

@@ -100,6 +100,30 @@ final class ValkeyChannelHandler: ChannelInboundHandler {
         }
     }
 
+    func addSubscription(
+        continuation: ValkeySubscriptionAsyncStream.Continuation,
+        filter: ValkeySubscriptionFilter
+    ) -> Int {
+        self.eventLoop.assertInEventLoop()
+        let id = ValkeySubscriptions.getSubscriptionID()
+        let loopBoundHandler = NIOLoopBound(self, eventLoop: self.eventLoop)
+        continuation.onTermination = { [eventLoop] termination in
+            switch termination {
+            case .cancelled:
+                eventLoop.execute {
+                    loopBoundHandler.value.subscriptions.removeSubscription(id: id)
+                }
+            case .finished:
+                break
+
+            @unknown default:
+                break
+            }
+        }
+        self.subscriptions.addSubscription(id: id, continuation: continuation, filter: filter)
+        return id
+    }
+
     @usableFromInline
     func handlerAdded(context: ChannelHandlerContext) {
         self.context = context

Sources/Valkey/Subscriptions/ValkeyConnection+subscribe.swift

@@ -48,7 +48,7 @@ extension ValkeyConnection {
 
         let channels = Set(channels)
         if self.channel.eventLoop.inEventLoop {
-            let subscriptionID = self.channelHandler.value.subscriptions.addSubscription(
+            let subscriptionID = self.channelHandler.value.addSubscription(
                 continuation: streamContinuation,
                 filter: .channels(channels)
             )
@@ -60,7 +60,7 @@ extension ValkeyConnection {
                 .get()
         } else {
             _ = try await self.channel.eventLoop.flatSubmit {
-                let subscriptionID = self.channelHandler.value.subscriptions.addSubscription(
+                let subscriptionID = self.channelHandler.value.addSubscription(
                     continuation: streamContinuation,
                     filter: .channels(channels)
                 )

Sources/Valkey/Subscriptions/ValkeySubscriptions.swift

@@ -66,16 +66,18 @@ struct ValkeySubscriptions {
         self.subscriptions = []
     }
 
-    mutating func addSubscription(continuation: ValkeySubscriptionAsyncStream.Continuation, filter: ValkeySubscriptionFilter) -> Int {
-        let id = Self.globalSubscriptionId.wrappingAdd(1, ordering: .relaxed)
+    static func getSubscriptionID() -> Int {
+        Self.globalSubscriptionId.wrappingAdd(1, ordering: .relaxed).newValue
+    }
+
+    mutating func addSubscription(id: Int, continuation: ValkeySubscriptionAsyncStream.Continuation, filter: ValkeySubscriptionFilter) {
         subscriptions.append(
             .init(
-                id: id.newValue,
+                id: id,
                 subscription: .init(continuation: continuation, filter: filter, logger: self.logger),
                 stateMachine: .init(filter: filter)
             )
         )
-        return id.newValue
     }
 
     mutating func removeSubscription(id: Int) {

Tests/ValkeyTests/ValkeySubscriptionTests.swift

@@ -16,7 +16,8 @@ import Logging
 import NIOCore
 import NIOEmbedded
 import Testing
-import Valkey
+
+@testable import Valkey
 
 @Suite
 struct SubscriptionTests {
@@ -225,4 +226,34 @@ struct SubscriptionTests {
             try await group.waitForAll()
         }
     }
+
+    @Test
+    func testRemoveSubscriptionOnCancellation() async throws {
+        let channel = NIOAsyncTestingChannel()
+        let logger = Logger(label: "test")
+        let connection = try await ValkeyConnection.setupChannel(channel, configuration: .init(), logger: logger)
+
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask {
+                let subscription = try await connection.subscribe(to: "test")
+                for try await message in subscription {
+                    #expect(message == .init(channel: "test", message: "Testing!"))
+                }
+            }
+            group.addTask {
+                let outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
+                // expect SUBSCRIBE command
+                #expect(String(buffer: outbound) == "*2\r\n$9\r\nSUBSCRIBE\r\n$4\r\ntest\r\n")
+                // push SUBSCRIBE channel
+                try await channel.writeInbound(ByteBuffer(string: ">3\r\n$9\r\nsubscribe\r\n$4\r\ntest\r\n:1\r\n"))
+                // push SUBSCRIBE message
+                try await channel.writeInbound(ByteBuffer(string: ">3\r\n$7\r\nmessage\r\n$4\r\ntest\r\n$8\r\nTesting!\r\n"))
+            }
+            try await group.next()
+            group.cancelAll()
+        }
+        try await connection.channel.eventLoop.submit {
+            #expect(connection.channelHandler.value.subscriptions.subscriptions.count == 0)
+        }.get()
+    }
 }