Don't allow SUBSCRIBE command to be cancelled (#214)

* Don't allow SUBSCRIBE command to be cancelled

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

* Re-instate testCancelSubscribe

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

---------

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

Author: adam-fowler

Sources/Valkey/Subscriptions/ValkeyConnection+subscribe.swift

@@ -204,37 +204,26 @@ extension ValkeyConnection {
     ) async throws -> (Int, ValkeySubscription) {
         let requestID = Self.requestIDGenerator.next()
         let (stream, streamContinuation) = ValkeySubscription.makeStream()
-        return try await withTaskCancellationHandler {
-            if Task.isCancelled {
-                throw ValkeyClientError(.cancelled)
-            }
-            let subscriptionID: Int = try await withCheckedThrowingContinuation { continuation in
-                self.channelHandler.subscribe(
-                    command: command,
-                    streamContinuation: streamContinuation,
-                    filters: filters,
-                    promise: .swift(continuation),
-                    requestID: requestID
-                )
-            }
-            return (subscriptionID, stream)
-        } onCancel: {
-            self.cancel(requestID: requestID)
+        if Task.isCancelled {
+            throw ValkeyClientError(.cancelled)
         }
+        let subscriptionID: Int = try await withCheckedThrowingContinuation { continuation in
+            self.channelHandler.subscribe(
+                command: command,
+                streamContinuation: streamContinuation,
+                filters: filters,
+                promise: .swift(continuation),
+                requestID: requestID
+            )
+        }
+        return (subscriptionID, stream)
     }
 
     @usableFromInline
     func unsubscribe(id: Int) async throws {
         let requestID = Self.requestIDGenerator.next()
-        try await withTaskCancellationHandler {
-            if Task.isCancelled {
-                throw ValkeyClientError(.cancelled)
-            }
-            try await withCheckedThrowingContinuation { continuation in
-                self.channelHandler.unsubscribe(id: id, promise: .swift(continuation), requestID: requestID)
-            }
-        } onCancel: {
-            self.cancel(requestID: requestID)
+        try await withCheckedThrowingContinuation { continuation in
+            self.channelHandler.unsubscribe(id: id, promise: .swift(continuation), requestID: requestID)
         }
     }
 

Tests/ValkeyTests/ValkeySubscriptionTests.swift

@@ -683,17 +683,25 @@ struct SubscriptionTests {
 
         try await withThrowingTaskGroup(of: Void.self) { group in
             group.addTask {
-                await #expect(throws: ValkeyClientError(.cancelled)) {
-                    try await connection.subscribe(to: "test") { _ in }
+                await #expect(throws: CancellationError.self) {
+                    try await connection.subscribe(to: "test") { subscription in
+                        for try await _ in subscription {}
+                    }
                 }
             }
-            group.addTask {
-                let outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
-                // expect SUBSCRIBE command
-                #expect(outbound == RESPToken(.command(["SUBSCRIBE", "test"])).base)
-            }
-            try await group.next()
+            var outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
+            // expect SUBSCRIBE command
+            #expect(outbound == RESPToken(.command(["SUBSCRIBE", "test"])).base)
             group.cancelAll()
+            try await Task.sleep(for: .milliseconds(10))
+
+            // push subscribe
+            try await channel.writeInbound(RESPToken(.push([.bulkString("subscribe"), .bulkString("test"), .number(1)])).base)
+            // expect UNSUBSCRIBE command
+            outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
+            #expect(outbound == RESPToken(.command(["UNSUBSCRIBE", "test"])).base)
+            // push unsubscribe
+            try await channel.writeInbound(RESPToken(.push([.bulkString("unsubscribe"), .bulkString("test"), .number(0)])).base)
         }
         #expect(await connection.isSubscriptionsEmpty())
     }