Fix Semaphore.withPermits leaking permits when interrupted

.changeset/fix-semaphore-with-permits-interrupt-leak.md

@@ -0,0 +1,5 @@
+---
+"effect": patch
+---
+
+Fix `Semaphore.withPermits` leaking permits when interrupted between acquiring them and installing their release.

packages/effect/src/Semaphore.ts

@@ -204,6 +204,24 @@ export interface Semaphore {
  */
 export const makeUnsafe = (permits: number): Semaphore => new SemaphoreImpl(permits)
 
+const waitForPermits = <A, E, R>(
+  self: SemaphoreImpl,
+  n: number,
+  effect: Effect.Effect<A, E, R>
+): Effect.Effect<A, E, R> =>
+  internal.callback((resume) => {
+    if (self.free >= n) return resume(effect)
+    const observer = () => {
+      if (self.free < n) return
+      self.waiters.delete(observer)
+      resume(effect)
+    }
+    self.waiters.add(observer)
+    return internal.sync(() => {
+      self.waiters.delete(observer)
+    })
+  })
+
 class SemaphoreImpl implements Semaphore {
   public waiters = new Set<() => void>()
   public taken = 0
@@ -220,18 +238,7 @@ class SemaphoreImpl implements Semaphore {
   take(n: number): Effect.Effect<number> {
     const take: Effect.Effect<number> = internal.suspend(() => {
       if (this.free < n) {
-        return internal.callback((resume) => {
-          if (this.free >= n) return resume(take)
-          const observer = () => {
-            if (this.free < n) return
-            this.waiters.delete(observer)
-            resume(take)
-          }
-          this.waiters.add(observer)
-          return internal.sync(() => {
-            this.waiters.delete(observer)
-          })
-        })
+        return waitForPermits(this, n, take)
       }
       this.taken += n
       return internal.succeed(n)
@@ -247,58 +254,56 @@ class SemaphoreImpl implements Semaphore {
     })
   }
 
-  updateTakenUnsafe(fiber: Fiber<any, any>, f: (n: number) => number): number {
-    this.taken = f(this.taken)
+  releaseUnsafe(fiber: Fiber<any, any>, n: number): number {
+    this.taken -= n
     if (this.waiters.size > 0) {
       fiber.currentDispatcher.scheduleTask(() => {
-        const iter = this.waiters.values()
-        let item = iter.next()
-        while (item.done === false && this.free > 0) {
-          item.value()
-          item = iter.next()
+        for (const observer of this.waiters) {
+          if (this.free <= 0) break
+          observer()
         }
       }, 0)
     }
     return this.free
   }
 
-  updateTaken(f: (n: number) => number): Effect.Effect<number> {
-    return core.withFiber((fiber) => internal.succeed(this.updateTakenUnsafe(fiber, f)))
-  }
-
   resize(permits: number) {
     return core.withFiber((fiber) => {
       this.permits = permits
       if (this.free < 0) return internal.void
-      this.updateTakenUnsafe(fiber, (taken) => taken)
+      this.releaseUnsafe(fiber, 0)
       return internal.void
     })
   }
 
   release(n: number): Effect.Effect<number> {
-    return this.updateTaken((taken) => taken - n)
+    return core.withFiber((fiber) => internal.succeed(this.releaseUnsafe(fiber, n)))
   }
 
   get releaseAll(): Effect.Effect<number> {
-    return this.updateTaken((_) => 0)
+    return core.withFiber((fiber) => internal.succeed(this.releaseUnsafe(fiber, this.taken)))
   }
 
   withPermits(n: number) {
     return <A, E, R>(self: Effect.Effect<A, E, R>) =>
-      internal.uninterruptibleMask((restore) =>
-        internal.flatMap(
-          restore(this.take(n)),
-          (permits) =>
-            internal.onExitPrimitive(
-              restore(self),
-              () => {
-                this.updateTakenUnsafe(internal.getCurrentFiber()!, (taken) => taken - permits)
-                return undefined
-              },
-              true
-            )
-        )
-      )
+      internal.uninterruptibleMask((restore) => {
+        const acquire: Effect.Effect<A, E, R> = internal.suspend(() => {
+          if (this.free < n) {
+            const wait = waitForPermits(this, n, internal.void)
+            return internal.flatMap(restore(wait), () => acquire)
+          }
+          this.taken += n
+          return internal.onExitPrimitive(
+            restore(self),
+            () => {
+              this.releaseUnsafe(internal.getCurrentFiber()!, n)
+              return undefined
+            },
+            true
+          )
+        })
+        return acquire
+      })
   }
 
   readonly withPermit = this.withPermits(1)
@@ -309,7 +314,7 @@ class SemaphoreImpl implements Semaphore {
         if (this.free < n) return internal.succeedNone
         this.taken += n
         return internal.onExitPrimitive(restore(internal.asSome(self)), () => {
-          this.updateTakenUnsafe(internal.getCurrentFiber()!, (taken) => taken - n)
+          this.releaseUnsafe(internal.getCurrentFiber()!, n)
           return undefined
         }, true)
       })

packages/effect/test/Semaphore.test.ts

@@ -332,6 +332,129 @@ describe("Semaphore", () => {
       assert.isTrue(acquired)
     }))
 
+  it.effect("withPermits interrupted at any operation does not leak permits", () =>
+    Effect.gen(function*() {
+      let operations = 0
+      const counted = new Scheduler.MixedScheduler()
+      const baseline = yield* (yield* Semaphore.make(1)).withPermits(1)(Effect.void).pipe(
+        Effect.provideService(Scheduler.Scheduler, {
+          executionMode: counted.executionMode,
+          makeDispatcher: () => counted.makeDispatcher(),
+          shouldYield: (fiber) => {
+            operations++
+            return counted.shouldYield(fiber)
+          }
+        }),
+        Effect.forkChild
+      )
+      yield* Fiber.await(baseline)
+
+      const recovered: Array<boolean> = []
+      let ranUnderPermits = false
+
+      for (let at = 1; at <= operations; at++) {
+        const sem = yield* Semaphore.make(1)
+        const base = new Scheduler.MixedScheduler()
+        let seen = 0
+        const scheduler: Scheduler.Scheduler = {
+          executionMode: base.executionMode,
+          makeDispatcher: () => base.makeDispatcher(),
+          shouldYield: (fiber) => {
+            if (++seen === at) fiber.interruptUnsafe()
+            return base.shouldYield(fiber)
+          }
+        }
+
+        const holder = yield* sem.withPermits(1)(
+          Effect.sync(() => {
+            ranUnderPermits = true
+          })
+        ).pipe(
+          Effect.provideService(Scheduler.Scheduler, scheduler),
+          Effect.forkChild
+        )
+        yield* Fiber.await(holder)
+
+        const result = yield* sem.withPermitsIfAvailable(1)(Effect.void)
+        recovered.push(Option.isSome(result))
+      }
+
+      assert.deepStrictEqual(recovered, Array.from({ length: operations }, () => true))
+      assert.isTrue(ranUnderPermits)
+    }))
+
+  it.effect("queued withPermits interrupted at any operation does not leak permits", () =>
+    Effect.gen(function*() {
+      const queuedScheduler = (
+        tasks: Array<() => void>,
+        onOperation: (fiber: Fiber.Fiber<unknown, unknown>) => void
+      ) => {
+        const scheduler: Scheduler.Scheduler = {
+          executionMode: "async",
+          makeDispatcher: () => ({
+            scheduleTask(task) {
+              tasks.push(task)
+            },
+            flush() {}
+          }),
+          shouldYield: (fiber) => {
+            onOperation(fiber)
+            return false
+          }
+        }
+        return scheduler
+      }
+
+      let operations = 0
+      const baselineSem = yield* Semaphore.make(1)
+      const baselineTasks: Array<() => void> = []
+      const baselineHolder = yield* Effect.forkChild(baselineSem.withPermits(1)(Effect.never))
+      yield* Effect.yieldNow
+      yield* Effect.forkChild(
+        baselineSem.withPermits(1)(Effect.void).pipe(
+          Effect.provideService(Scheduler.Scheduler, queuedScheduler(baselineTasks, () => operations++))
+        ),
+        { startImmediately: true }
+      )
+      yield* Fiber.interrupt(baselineHolder)
+      yield* Effect.yieldNow
+      while (baselineTasks.length > 0) baselineTasks.shift()!()
+
+      const recovered: Array<boolean> = []
+      let ranUnderPermits = false
+
+      for (let at = 1; at <= operations; at++) {
+        const sem = yield* Semaphore.make(1)
+        const tasks: Array<() => void> = []
+        let seen = 0
+        const scheduler = queuedScheduler(tasks, (fiber) => {
+          if (++seen === at) fiber.interruptUnsafe()
+        })
+
+        const holder = yield* Effect.forkChild(sem.withPermits(1)(Effect.never))
+        yield* Effect.yieldNow
+
+        yield* Effect.forkChild(
+          sem.withPermits(1)(
+            Effect.sync(() => {
+              ranUnderPermits = true
+            })
+          ).pipe(Effect.provideService(Scheduler.Scheduler, scheduler)),
+          { startImmediately: true }
+        )
+
+        yield* Fiber.interrupt(holder)
+        yield* Effect.yieldNow
+        while (tasks.length > 0) tasks.shift()!()
+
+        const result = yield* sem.withPermitsIfAvailable(1)(Effect.void)
+        recovered.push(Option.isSome(result))
+      }
+
+      assert.deepStrictEqual(recovered, Array.from({ length: operations }, () => true))
+      assert.isTrue(ranUnderPermits)
+    }))
+
   it.effect("takeIfAvailable acquires permits when they are available", () =>
     Effect.gen(function*() {
       const sem = yield* Semaphore.make(2)