1/**
2 * The `PartitionedSemaphore` module provides a semaphore for limiting
3 * concurrency across a shared permit pool while keeping waiters grouped by
4 * partition key. A `PartitionedSemaphore<K>` is useful when many independent
5 * groups of work compete for the same bounded resource and each group should
6 * make progress without one busy group monopolizing released permits.
7 *
8 * @since 4.0.0
9 */
10import * as Effect from "./Effect.ts"
11import * as Exit from "./Exit.ts"
12import { dual } from "./Function.ts"
13import * as internalEffect from "./internal/effect.ts"
14import * as MutableHashMap from "./MutableHashMap.ts"
15import * as Option from "./Option.ts"
16
17/**
18 * Runtime type identifier used to mark values that implement
19 * `PartitionedSemaphore`.
20 *
21 * **Details**
22 *
23 * This marker is part of the runtime representation of partitioned semaphore
24 * values.
25 *
26 * @category type IDs
27 * @since 4.0.0
28 */
29export const PartitionedTypeId: PartitionedTypeId = "~effect/PartitionedSemaphore"
30
31/**
32 * Literal type of the `PartitionedSemaphore` runtime type identifier.
33 *
34 * **When to use**
35 *
36 * Use to type fields that store the exact `PartitionedSemaphore` runtime marker.
37 *
38 * **Details**
39 *
40 * Use this type when declaring fields that must contain the exact
41 * `PartitionedTypeId` marker value.
42 *
43 * @category type IDs
44 * @since 4.0.0
45 */
46export type PartitionedTypeId = "~effect/PartitionedSemaphore"
47
48/**
49 * A `PartitionedSemaphore` controls access to a shared permit pool while
50 * tracking waiters by partition key.
51 *
52 * **When to use**
53 *
54 * Use to coordinate shared permits across partition keys so waiting groups make
55 * progress without one group monopolizing the pool.
56 *
57 * **Details**
58 *
59 * Waiting permits are distributed across partitions in round-robin order.
60 *
61 * @category models
62 * @since 3.19.4
63 */
64export interface PartitionedSemaphore<in K> {
65 readonly [PartitionedTypeId]: PartitionedTypeId
66 readonly capacity: number
67 readonly available: Effect.Effect<number>
68 readonly take: (key: K, permits: number) => Effect.Effect<void>
69 readonly release: (permits: number) => Effect.Effect<number>
70 readonly withPermits: (
71 key: K,
72 permits: number
73 ) => <A, E, R>(effect: Effect.Effect<A, E, R>) => Effect.Effect<A, E, R>
74 readonly withPermit: (key: K) => <A, E, R>(effect: Effect.Effect<A, E, R>) => Effect.Effect<A, E, R>
75 readonly withPermitsIfAvailable: (
76 permits: number
77 ) => <A, E, R>(effect: Effect.Effect<A, E, R>) => Effect.Effect<Option.Option<A>, E, R>
78}
79
80/**
81 * Alias interface for a `PartitionedSemaphore` keyed by values of type `K`.
82 *
83 * **When to use**
84 *
85 * Use as an alternate exported name for a partitioned permit pool keyed by `K`.
86 *
87 * **Details**
88 *
89 * This interface does not add members beyond `PartitionedSemaphore`; it
90 * provides an alternate exported name for APIs that refer to a partitioned
91 * permit pool.
92 *
93 * @category models
94 * @since 4.0.0
95 */
96export interface Partitioned<in K> extends PartitionedSemaphore<K> {}
97
98/**
99 * Constructs a `PartitionedSemaphore` synchronously, outside of `Effect`.
100 *
101 * **When to use**
102 *
103 * Use when you need to construct a partitioned semaphore synchronously outside
104 * an `Effect` workflow.
105 *
106 * **Details**
107 *
108 * Negative permit counts are clamped to `0`. Non-finite permit counts create
109 * an unbounded semaphore whose acquire and release operations complete
110 * immediately.
111 *
112 * @see {@link make} for creating a partitioned semaphore inside `Effect`
113 *
114 * @category constructors
115 * @since 3.19.4
116 */
117export const makeUnsafe = <K = unknown>(options: {
118 readonly permits: number
119}): PartitionedSemaphore<K> => {
120 const maxPermits = Math.max(0, options.permits)
121
122 if (!Number.isFinite(maxPermits)) {
123 return {
124 [PartitionedTypeId]: PartitionedTypeId,
125 capacity: maxPermits,
126 available: Effect.succeed(maxPermits),
127 take: () => Effect.void,
128 release: () => Effect.succeed(maxPermits),
129 withPermits: () => (effect) => effect,
130 withPermit: () => (effect) => effect,
131 withPermitsIfAvailable: () => (effect) => Effect.asSome(effect)
132 }
133 }
134
135 let totalPermits = maxPermits
136 let waitingPermits = 0
137
138 type Waiter = {
139 permits: number
140 readonly resume: () => void
141 }
142
143 const partitions = MutableHashMap.empty<K, Set<Waiter>>()
144 let iterator = partitions[Symbol.iterator]()
145
146 const releaseUnsafe = (permits: number): number => {
147 while (permits > 0) {
148 if (waitingPermits === 0) {
149 totalPermits = Math.min(maxPermits, totalPermits + permits)
150 return totalPermits
151 }
152
153 let state = iterator.next()
154 if (state.done) {
155 iterator = partitions[Symbol.iterator]()
156 state = iterator.next()
157 if (state.done) {
158 return totalPermits
159 }
160 }
161
162 const waiter = state.value[1].values().next().value
163 if (waiter === undefined) {
164 continue
165 }
166
167 waiter.permits -= 1
168 waitingPermits -= 1
169
170 if (waiter.permits === 0) {
171 waiter.resume()
172 }
173
174 permits -= 1
175 }
176
177 return totalPermits
178 }
179
180 const take = (key: K, permits: number): Effect.Effect<void> => {
181 if (permits <= 0) {
182 return Effect.void
183 }
184
185 if (maxPermits < permits) {
186 return Effect.never
187 }
188
189 return Effect.withFiber((fiber) => {
190 if (totalPermits >= permits) {
191 // Keep the capacity check, deduction, and interruption cleanup in one evaluation step.
192 totalPermits -= permits
193 internalEffect.onExitUnsafe(fiber, (exit) => {
194 if (Exit.isFailure(exit)) {
195 releaseUnsafe(permits)
196 }
197 return undefined
198 })
199 return Effect.void
200 }
201
202 return Effect.callback<void>((resume) => {
203 if (totalPermits >= permits) {
204 resume(take(key, permits))
205 return
206 }
207 const needed = permits - totalPermits
208 if (totalPermits > 0) {
209 totalPermits = 0
210 }
211 waitingPermits += needed
212
213 const waiters = Option.getOrElse(
214 MutableHashMap.get(partitions, key),
215 () => {
216 const set = new Set<Waiter>()
217 MutableHashMap.set(partitions, key, set)
218 return set
219 }
220 )
221
222 const entry: Waiter = {
223 permits: needed,
224 resume: () => {
225 cleanup()
226 resume(Effect.void)
227 }
228 }
229
230 const cleanup = () => {
231 if (waiters.delete(entry) && waiters.size === 0) {
232 MutableHashMap.remove(partitions, key)
233 }
234 }
235
236 waiters.add(entry)
237
238 return Effect.sync(() => {
239 cleanup()
240 waitingPermits -= entry.permits
241 releaseUnsafe(permits - entry.permits)
242 })
243 })
244 })
245 }
246
247 const withPermits =
248 (key: K, permits: number) => <A, E, R>(effect: Effect.Effect<A, E, R>): Effect.Effect<A, E, R> => {
249 if (permits <= 0) {
250 return effect
251 }
252
253 const takePermits = take(key, permits)
254 return Effect.uninterruptibleMask((restore) =>
255 Effect.flatMap(
256 restore(takePermits),
257 () =>
258 Effect.ensuring(
259 restore(effect),
260 Effect.sync(() => {
261 releaseUnsafe(permits)
262 })
263 )
264 )
265 )
266 }
267
268 const tryTake = (permits: number): boolean => {
269 if (permits <= 0) {
270 return true
271 }
272
273 if (maxPermits < permits || totalPermits < permits) {
274 return false
275 }
276
277 totalPermits -= permits
278 return true
279 }
280
281 return {
282 [PartitionedTypeId]: PartitionedTypeId,
283 capacity: maxPermits,
284 available: Effect.sync(() => totalPermits),
285 take,
286 release: (permits) => Effect.sync(() => releaseUnsafe(permits)),
287 withPermits,
288 withPermit: (key) => withPermits(key, 1),
289 withPermitsIfAvailable:
290 (permits) => <A, E, R>(effect: Effect.Effect<A, E, R>): Effect.Effect<Option.Option<A>, E, R> => {
291 if (permits <= 0) {
292 return Effect.asSome(effect)
293 }
294
295 return Effect.withFiber((fiber) => {
296 if (!tryTake(permits)) {
297 return Effect.succeed(Option.none())
298 }
299
300 // Register cleanup before the fiber can yield to the user effect.
301 internalEffect.onExitUnsafe(fiber, () => {
302 releaseUnsafe(permits)
303 return undefined
304 })
305 return Effect.asSome(effect)
306 })
307 }
308 }
309}
310
311/**
312 * Creates a `PartitionedSemaphore` inside an `Effect`.
313 *
314 * **When to use**
315 *
316 * Use when semaphore construction should stay inside an `Effect` workflow.
317 *
318 * **Details**
319 *
320 * The `permits` option sets the shared permit capacity. The resulting
321 * semaphore tracks waiters by partition key and distributes released permits
322 * across waiting partitions in round-robin order.
323 *
324 * **Gotchas**
325 *
326 * Negative permit counts are clamped to `0`. Non-finite permit counts create
327 * an unbounded semaphore.
328 *
329 * @see {@link makeUnsafe} for synchronous construction
330 *
331 * @category constructors
332 * @since 3.19.4
333 */
334export const make = <K = unknown>(options: {
335 readonly permits: number
336}): Effect.Effect<PartitionedSemaphore<K>> => Effect.sync(() => makeUnsafe<K>(options))
337
338/**
339 * Gets the current number of available permits.
340 *
341 * **When to use**
342 *
343 * Use to inspect a snapshot of how many permits are currently free.
344 *
345 * **Details**
346 *
347 * Running the returned effect reads the semaphore's current availability.
348 * Taking permits decreases availability, and releasing permits can increase it
349 * up to the semaphore capacity.
350 *
351 * **Gotchas**
352 *
353 * Reading availability does not reserve permits.
354 *
355 * @see {@link capacity} for the fixed total permit capacity
356 * @see {@link release} for returning permits to the shared pool
357 * @see {@link withPermitsIfAvailable} for running only when permits are immediately available
358 *
359 * @category combinators
360 * @since 4.0.0
361 */
362export const available = <K>(self: PartitionedSemaphore<K>): Effect.Effect<number> => self.available
363
364/**
365 * Gets the total capacity.
366 *
367 * **When to use**
368 *
369 * Use to inspect the fixed number of permits configured for the semaphore.
370 *
371 * **Details**
372 *
373 * Capacity is stored when the semaphore is created and does not change as
374 * permits are acquired or released.
375 *
376 * @see {@link available} for the current number of free permits
377 *
378 * @category getters
379 * @since 4.0.0
380 */
381export const capacity = <K>(self: PartitionedSemaphore<K>): number => self.capacity
382
383/**
384 * Returns an effect that acquires the requested number of permits for the
385 * given partition key.
386 *
387 * **When to use**
388 *
389 * Use when you need manual permit acquisition for a partition and want to
390 * control acquisition and release as separate effects.
391 *
392 * **Details**
393 *
394 * If enough permits are available, the effect completes immediately. Otherwise
395 * it waits until released permits are assigned to this partition.
396 *
397 * **Gotchas**
398 *
399 * Requests for more permits than the semaphore capacity never complete.
400 * Requests for zero or a negative number of permits complete without acquiring
401 * anything.
402 *
403 * @see {@link release} for manually returning permits to the shared pool
404 * @see {@link withPermits} for automatic acquire and release around an effect
405 * @see {@link withPermit} for acquiring exactly one permit around an effect
406 *
407 * @category combinators
408 * @since 4.0.0
409 */
410export const take: {
411 <K>(key: K, permits: number): (self: PartitionedSemaphore<K>) => Effect.Effect<void>
412 <K>(self: PartitionedSemaphore<K>, key: K, permits: number): Effect.Effect<void>
413} = dual(3, <K>(self: PartitionedSemaphore<K>, key: K, permits: number): Effect.Effect<void> => self.take(key, permits))
414
415/**
416 * Returns an effect that releases permits back to the shared pool and returns
417 * the current available permit count.
418 *
419 * **When to use**
420 *
421 * Use when you need to return permits acquired with `take` in a lower-level
422 * partitioned permit protocol with explicit release control.
423 *
424 * **Details**
425 *
426 * Released permits are first assigned to waiting partitions in round-robin
427 * order. Only permits not needed by waiters increase the available count,
428 * which is capped at the semaphore capacity.
429 *
430 * @see {@link take} for manual acquisition
431 * @see {@link withPermits} for automatic acquire and release around an effect
432 * @see {@link available} for reading the permit count without releasing
433 *
434 * @category combinators
435 * @since 4.0.0
436 */
437export const release: {
438 (permits: number): <K>(self: PartitionedSemaphore<K>) => Effect.Effect<number>
439 <K>(self: PartitionedSemaphore<K>, permits: number): Effect.Effect<number>
440} = dual(2, <K>(self: PartitionedSemaphore<K>, permits: number): Effect.Effect<number> => self.release(permits))
441
442/**
443 * Runs an effect after acquiring permits for a partition, then releases those
444 * permits when the effect exits.
445 *
446 * **When to use**
447 *
448 * Use to guard weighted partitioned work with automatic permit acquisition and
449 * release around an effect.
450 *
451 * **Details**
452 *
453 * Permit acquisition may wait according to `take` semantics. Once acquired,
454 * the permits are released even if the wrapped effect fails or is interrupted.
455 *
456 * **Gotchas**
457 *
458 * Requests for more permits than the semaphore capacity never complete.
459 * Requests for zero or a negative number of permits run the effect without
460 * acquiring anything.
461 *
462 * @see {@link withPermit} for the single-permit variant
463 * @see {@link withPermitsIfAvailable} for running only when permits are
464 * immediately available
465 * @see {@link take} for manual acquisition
466 * @see {@link release} for manual release
467 *
468 * @category combinators
469 * @since 4.0.0
470 */
471export const withPermits: {
472 <K>(
473 self: PartitionedSemaphore<K>,
474 key: K,
475 permits: number
476 ): <A, E, R>(effect: Effect.Effect<A, E, R>) => Effect.Effect<A, E, R>
477 <K, A, E, R>(
478 self: PartitionedSemaphore<K>,
479 key: K,
480 permits: number,
481 effect: Effect.Effect<A, E, R>
482 ): Effect.Effect<A, E, R>
483} = ((...args: Array<any>) => {
484 if (args.length === 3) {
485 const [self, key, permits] = args
486 return (effect: Effect.Effect<any, any, any>) => self.withPermits(key, permits)(effect)
487 }
488 const [self, key, permits, effect] = args
489 return self.withPermits(key, permits)(effect)
490}) as any
491
492/**
493 * Runs an effect after acquiring one permit for a partition, then releases the
494 * permit when the effect exits.
495 *
496 * **When to use**
497 *
498 * Use to guard partitioned work with exactly one permit and automatic release
499 * when the effect exits.
500 *
501 * **Details**
502 *
503 * This is the single-permit variant of `withPermits`. The permit is released
504 * even if the wrapped effect fails or is interrupted.
505 *
506 * @see {@link withPermits} for acquiring a weighted number of permits
507 * @see {@link withPermitsIfAvailable} for running only when permits are
508 * immediately available
509 * @see {@link take} for manual acquisition
510 * @see {@link release} for manual release
511 *
512 * @category combinators
513 * @since 4.0.0
514 */
515export const withPermit: {
516 <K>(self: PartitionedSemaphore<K>, key: K): <A, E, R>(effect: Effect.Effect<A, E, R>) => Effect.Effect<A, E, R>
517 <K, A, E, R>(
518 self: PartitionedSemaphore<K>,
519 key: K,
520 effect: Effect.Effect<A, E, R>
521 ): Effect.Effect<A, E, R>
522} = ((...args: Array<any>) => {
523 if (args.length === 2) {
524 const [self, key] = args
525 return (effect: Effect.Effect<any, any, any>) => self.withPermit(key)(effect)
526 }
527 const [self, key, effect] = args
528 return self.withPermit(key)(effect)
529}) as any
530
531/**
532 * Runs an effect only when the requested permits can be acquired immediately,
533 * returning the result in `Some`.
534 *
535 * **When to use**
536 *
537 * Use when guarded work should run only if the shared permit pool can provide
538 * the requested permits immediately.
539 *
540 * **Details**
541 *
542 * If the permits are not available, the effect is not run and the result is
543 * `None`. When permits are acquired, they are released after the wrapped
544 * effect completes, fails, or is interrupted. Requests for zero or a negative
545 * number of permits run the effect and return `Some`.
546 *
547 * @see {@link withPermits} for the keyed variant that waits until permits are
548 * available for a partition
549 *
550 * @category combinators
551 * @since 4.0.0
552 */
553export const withPermitsIfAvailable: {
554 <K>(
555 self: PartitionedSemaphore<K>,
556 permits: number
557 ): <A, E, R>(effect: Effect.Effect<A, E, R>) => Effect.Effect<Option.Option<A>, E, R>
558 <K, A, E, R>(
559 self: PartitionedSemaphore<K>,
560 permits: number,
561 effect: Effect.Effect<A, E, R>
562 ): Effect.Effect<Option.Option<A>, E, R>
563} = ((...args: Array<any>) => {
564 if (args.length === 2) {
565 const [self, permits] = args
566 return (effect: Effect.Effect<any, any, any>) => self.withPermitsIfAvailable(permits)(effect)
567 }
568 const [self, permits, effect] = args
569 return self.withPermitsIfAvailable(permits)(effect)
570}) as any