Call unsubscribe from unstructured Task to avoid cancellation

Signed-off-by: Adam Fowler <[email protected]>

Author: adam-fowler

Sources/Valkey/Subscriptions/ValkeyConnection+subscribe.swift

@@ -194,7 +194,10 @@ extension ValkeyConnection {
             value = try await process(stream)
             try Task.checkCancellation()
         } catch {
-            _ = try? await unsubscribe(id: id)
+            // call unsubscrobe to avoid it being cancelled
+            _ = await Task {
+                try await unsubscribe(id: id)
+            }.result
             throw error
         }
         _ = try await unsubscribe(id: id)

Tests/IntegrationTests/ValkeyTests.swift

@@ -468,6 +468,30 @@ struct GeneratedCommands {
         }
     }
 
+    @Test
+    @available(valkeySwift 1.0, *)
+    func testCancelSubscription() async throws {
+        let (stream, cont) = AsyncStream.makeStream(of: Void.self)
+        var logger = Logger(label: "Subscriptions")
+        logger.logLevel = .trace
+        try await withValkeyClient(.hostname(valkeyHostname, port: 6379), logger: logger) { client in
+            await withThrowingTaskGroup(of: Void.self) { group in
+                group.addTask {
+                    try await client.withConnection { connection in
+                        try await connection.subscribe(to: "testCancelSubscriptions") { subscription in
+                            cont.finish()
+                            for try await _ in subscription {
+                            }
+                        }
+                        #expect(await connection.isSubscriptionsEmpty())
+                    }
+                }
+                await stream.first { _ in true }
+                group.cancelAll()
+            }
+        }
+    }
+
     @Test
     @available(valkeySwift 1.0, *)
     func testClientSubscriptions() async throws {
@@ -494,6 +518,27 @@ struct GeneratedCommands {
         }
     }
 
+    @Test
+    @available(valkeySwift 1.0, *)
+    func testClientCancelSubscription() async throws {
+        let (stream, cont) = AsyncStream.makeStream(of: Void.self)
+        var logger = Logger(label: "Subscriptions")
+        logger.logLevel = .trace
+        try await withValkeyClient(.hostname(valkeyHostname, port: 6379), logger: logger) { client in
+            await withThrowingTaskGroup(of: Void.self) { group in
+                group.addTask {
+                    try await client.subscribe(to: "testCancelSubscriptions") { subscription in
+                        cont.finish()
+                        for try await _ in subscription {
+                        }
+                    }
+                }
+                await stream.first { _ in true }
+                group.cancelAll()
+            }
+        }
+    }
+
     /// Test two different subscriptions to the same channel both receive messages and that when one ends the other still
     /// receives messages
     @Test

Tests/ValkeyTests/ValkeySubscriptionTests.swift

@@ -729,9 +729,16 @@ struct SubscriptionTests {
                 try await channel.writeInbound(RESPToken(.push([.bulkString("subscribe"), .bulkString("test"), .number(1)])).base)
                 // push message
                 try await channel.writeInbound(RESPToken(.push([.bulkString("message"), .bulkString("test"), .bulkString("Testing!")])).base)
+
             }
             try await group.next()
             group.cancelAll()
+
+            // respond to unsubscribe after cancellation
+            let outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
+            #expect(outbound == RESPToken(.command(["UNSUBSCRIBE", "test"])).base)
+            // push unsubcribe
+            try await channel.writeInbound(RESPToken(.push([.bulkString("unsubscribe"), .bulkString("test"), .number(1)])).base)
         }
         #expect(await connection.isSubscriptionsEmpty())
     }