Semaphore cancellation should increment semaphore value

.swiftpm/xcode/xcshareddata/xcschemes/Semaphore.xcscheme

@@ -40,8 +40,15 @@
       buildConfiguration = "Debug"
       selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
       selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
-      shouldUseLaunchSchemeArgsEnv = "YES"
+      shouldUseLaunchSchemeArgsEnv = "NO"
       codeCoverageEnabled = "YES">
+      <EnvironmentVariables>
+         <EnvironmentVariable
+            key = "LIBDISPATCH_COOPERATIVE_POOL_STRICT"
+            value = "1"
+            isEnabled = "NO">
+         </EnvironmentVariable>
+      </EnvironmentVariables>
       <Testables>
          <TestableReference
             skipped = "NO"

Sources/Semaphore/AsyncSemaphore.swift

@@ -161,9 +161,12 @@ public final class AsyncSemaphore: @unchecked Sendable {
 
         value -= 1
         if value >= 0 {
-            unlock()
+            defer { unlock() }
             // All code paths check for cancellation
-            try Task.checkCancellation()
+            if Task.isCancelled {
+                value += 1
+                throw CancellationError()
+            }
             return
         }
 

Tests/SemaphoreTests/AsyncSemaphoreTests.swift

@@ -3,6 +3,12 @@ import XCTest
 @testable import Semaphore
 
 final class AsyncSemaphoreTests: XCTestCase {
+
+    override func setUp() {
+        super.setUp()
+        // Don't continue after completeWithin(nanoseconds:) causes a XCTFail
+        continueAfterFailure = false
+    }
 
     func testSignalWithoutSuspendedTasks() async {
         // Check DispatchSemaphore behavior
@@ -104,7 +110,7 @@ final class AsyncSemaphoreTests: XCTestCase {
             let ex1 = expectation(description: "wait")
             ex1.isInverted = true
             let ex2 = expectation(description: "woken")
-            Task {
+            let task = Task {
                 await sem.wait()
                 ex1.fulfill()
                 ex2.fulfill()
@@ -115,6 +121,7 @@ final class AsyncSemaphoreTests: XCTestCase {
 
             // When a signal occurs, then the suspended task is resumed.
             sem.signal()
+            await task.value
             wait(for: [ex2], timeout: 0.5)
         }
     }
@@ -134,6 +141,7 @@ final class AsyncSemaphoreTests: XCTestCase {
         }
         try await Task.sleep(nanoseconds: 100_000_000)
         task.cancel()
+        await task.value
         wait(for: [ex], timeout: 1)
     }
 
@@ -157,6 +165,7 @@ final class AsyncSemaphoreTests: XCTestCase {
             ex.fulfill()
         }
         task.cancel()
+        await task.value
         wait(for: [ex], timeout: 5)
     }
 
@@ -173,7 +182,7 @@ final class AsyncSemaphoreTests: XCTestCase {
         let ex1 = expectation(description: "wait")
         ex1.isInverted = true
         let ex2 = expectation(description: "woken")
-        Task {
+        let taskTwo = Task {
             await sem.wait()
             ex1.fulfill()
             ex2.fulfill()
@@ -184,6 +193,7 @@ final class AsyncSemaphoreTests: XCTestCase {
 
         // When a signal occurs, then the suspended task is resumed.
         sem.signal()
+        await taskTwo.value
         wait(for: [ex2], timeout: 0.5)
     }
 
@@ -205,7 +215,7 @@ final class AsyncSemaphoreTests: XCTestCase {
         let ex1 = expectation(description: "wait")
         ex1.isInverted = true
         let ex2 = expectation(description: "woken")
-        Task {
+        let taskTwo = Task {
             await sem.wait()
             ex1.fulfill()
             ex2.fulfill()
@@ -216,9 +226,46 @@ final class AsyncSemaphoreTests: XCTestCase {
 
         // When a signal occurs, then the suspended task is resumed.
         sem.signal()
+        await taskTwo.value
         wait(for: [ex2], timeout: 0.5)
     }
-
+
+    func test_that_cancellation_before_suspension_increments_the_semaphore_two() async {
+        await completeWithin(nanoseconds: NSEC_PER_SEC * 2) {
+            let sem = AsyncSemaphore(value: 1)
+            let task = Task {
+                while !Task.isCancelled {
+                    await Task.yield()
+                }
+                try await sem.waitUnlessCancelled()
+            }
+            task.cancel()
+            try? await task.value
+            await sem.wait()
+        }
+    }
+
+    func test_that_cancellation_while_suspended_increments_the_semaphore_two() async {
+        await completeWithin(nanoseconds: NSEC_PER_SEC * 2) {
+            let sem = AsyncSemaphore(value: 0)
+            let running = Atomic(false)
+            let task = Task {
+                running.mutate { $0 = true }
+                try await sem.waitUnlessCancelled()
+                while !Task.isCancelled {
+                    await Task.yield()
+                }
+            }
+            while !running.value {
+                await Task.yield()
+            }
+            task.cancel()
+            try? await task.value
+            sem.signal()
+            await sem.wait()
+        }
+    }
+
     // Test that semaphore can limit the number of concurrent executions of
     // an actor method.
     func test_semaphore_as_a_resource_limiter_on_actor_method() async {
@@ -395,3 +442,50 @@ final class AsyncSemaphoreTests: XCTestCase {
         }
     }
 }
+
+/// Helper to complete a test within some amount of time or fail.
+/// XCTestExpectation don't work, when using LIBDISPATCH_COOPERATIVE_POOL_STRICT=1 as environment variable
+/// or e.g. running tests on an iOS simulator as the wait(for:timeout:) blocks the pool.
+/// Which means after the await no further async work can execute to fulfill any expectation that wait
+/// is waiting for.
+func completeWithin(nanoseconds nanosecondsDeadline: UInt64,
+                    file: StaticString = #filePath,
+                    line: UInt = #line,
+                    work: () async throws -> Void) async rethrows {
+    let checkDeadlineTask = Task {
+        try await Task.sleep(nanoseconds: nanosecondsDeadline)
+        try Task.checkCancellation()
+        XCTFail("Test timed out.", file: file, line: line)
+    }
+    try await work()
+    checkDeadlineTask.cancel()
+}
+
+final class Atomic<A>: @unchecked Sendable {
+    private var lock = NSRecursiveLock()
+    private var _value: A
+
+    public init(_ value: A) {
+        _value = value
+    }
+
+    public var value: A {
+        synced {
+            _value
+        }
+    }
+
+    public func mutate(_ transform: (inout A) -> Void) {
+        synced {
+            transform(&self._value)
+        }
+    }
+
+    private func synced<Result>(_ action: () throws -> Result) rethrows -> Result {
+        lock.lock()
+        defer {
+            lock.unlock()
+        }
+        return try action()
+    }
+}