1import * as Context from "../../Context.ts"
2import * as Deferred from "../../Deferred.ts"
3import * as Effect from "../../Effect.ts"
4import * as Exit from "../../Exit.ts"
5import * as Fiber from "../../Fiber.ts"
6import * as MutableHashMap from "../../MutableHashMap.ts"
7import * as MutableRef from "../../MutableRef.ts"
8import * as Scope from "../../Scope.ts"
9
10/**
11 * @internal
12 */
13export class ResourceMap<K, A, E> {
14 readonly lookup: (key: K, scope: Scope.Scope) => Effect.Effect<A, E>
15 readonly entries: BackingMap<K, A, E>
16 readonly isClosed: MutableRef.MutableRef<boolean>
17 constructor(
18 lookup: (key: K, scope: Scope.Scope) => Effect.Effect<A, E>,
19 entries: BackingMap<K, A, E>,
20 isClosed: MutableRef.MutableRef<boolean>
21 ) {
22 this.lookup = lookup
23 this.entries = entries
24 this.isClosed = isClosed
25 }
26
27 static make = Effect.fnUntraced(function*<K, A, E, R>(lookup: (key: K) => Effect.Effect<A, E, R>, options?: {
28 readonly referential?: boolean | undefined
29 }) {
30 const scope = yield* Effect.scope
31 const services = yield* Effect.context<R>()
32 const isClosed = MutableRef.make(false)
33
34 const entries: BackingMap<K, A, E> = options?.referential ?
35 {
36 _tag: "Referential",
37 map: new Map()
38 } :
39 {
40 _tag: "Equal",
41 map: MutableHashMap.empty()
42 }
43
44 yield* Scope.addFinalizerExit(
45 scope,
46 (exit) => {
47 MutableRef.set(isClosed, true)
48 return Effect.forEach(entries.map, ([key, { fiber, scope }]) => {
49 backingDelete(entries, key)
50 return Effect.exit(Effect.andThen(Fiber.interrupt(fiber), Scope.close(scope, exit)))
51 }, { concurrency: "unbounded", discard: true })
52 }
53 )
54
55 return new ResourceMap(
56 (key, scope) => Effect.provide(lookup(key), Context.add(services, Scope.Scope, scope)),
57 entries,
58 isClosed
59 )
60 })
61
62 hasUnsafe(key: K): boolean {
63 return backingGet(this.entries, key) !== undefined
64 }
65
66 keysUnsafe(): Array<K> {
67 return Array.from(this.entries.map, ([key]) => key)
68 }
69
70 get(key: K): Effect.Effect<A, E> {
71 return Effect.uninterruptibleMask((restore) => {
72 if (MutableRef.get(this.isClosed)) {
73 return Effect.interrupt
74 }
75 const existing = backingGet(this.entries, key)
76 if (existing) {
77 return this.awaitEntry(key, existing, restore)
78 }
79 const entry: Entry<A, E> = {
80 scope: Scope.makeUnsafe(),
81 deferred: Deferred.makeUnsafe<A, E>(),
82 fiber: undefined as any,
83 awaiters: 0
84 }
85 backingSet(this.entries, key, entry)
86 // Run the lookup on a detached fiber so it is shared by every caller and
87 // only interrupted once all of them have been interrupted.
88 entry.fiber = Effect.runForkWith(Fiber.getCurrent()!.context)(
89 Effect.onExit(this.lookup(key, entry.scope), (exit) => {
90 if (exit._tag === "Success") {
91 return Deferred.done(entry.deferred, exit)
92 }
93 this.deleteEntry(key, entry)
94 return Effect.andThen(
95 Deferred.done(entry.deferred, exit),
96 Scope.close(entry.scope, exit)
97 )
98 })
99 )
100 return this.awaitEntry(key, entry, restore)
101 })
102 }
103
104 private awaitEntry(
105 key: K,
106 entry: Entry<A, E>,
107 restore: <AX, EX, RX>(effect: Effect.Effect<AX, EX, RX>) => Effect.Effect<AX, EX, RX>
108 ): Effect.Effect<A, E> {
109 if (Deferred.isDoneUnsafe(entry.deferred)) return Deferred.await(entry.deferred)
110 entry.awaiters++
111 return Effect.onExit(restore(Deferred.await(entry.deferred)), () => {
112 if (--entry.awaiters > 0 || Deferred.isDoneUnsafe(entry.deferred)) {
113 return Effect.void
114 }
115 // Detach the abandoned entry first so new callers start a fresh lookup
116 // instead of joining one that is being interrupted.
117 this.deleteEntry(key, entry)
118 return Fiber.interrupt(entry.fiber)
119 })
120 }
121
122 private deleteEntry(key: K, entry: Entry<A, E>): void {
123 // Never delete a replacement entry created after this one was removed.
124 if (backingGet(this.entries, key) === entry) {
125 backingDelete(this.entries, key)
126 }
127 }
128
129 remove(key: K): Effect.Effect<void> {
130 return Effect.suspend(() => {
131 const entry = backingGet(this.entries, key)
132 if (!entry) {
133 return Effect.void
134 }
135 backingDelete(this.entries, key)
136 return Scope.close(entry.scope, Exit.void)
137 })
138 }
139
140 removeIgnore(key: K): Effect.Effect<void> {
141 return Effect.catchCause(this.remove(key), (cause) =>
142 Effect.annotateLogs(Effect.logDebug(cause), {
143 module: "ResourceMap",
144 method: "removeIgnore",
145 key
146 }))
147 }
148}
149
150type BackingMap<K, A, E> = {
151 readonly _tag: "Equal"
152 readonly map: MutableHashMap.MutableHashMap<K, Entry<A, E>>
153} | {
154 readonly _tag: "Referential"
155 readonly map: Map<K, Entry<A, E>>
156}
157
158type Entry<A, E> = {
159 readonly scope: Scope.Closeable
160 readonly deferred: Deferred.Deferred<A, E>
161 // Assigned as soon as the lookup is forked.
162 fiber: Fiber.Fiber<unknown, unknown>
163 awaiters: number
164}
165
166const backingGet = <K, A, E>(map: BackingMap<K, A, E>, key: K): Entry<A, E> | undefined => {
167 if (map._tag === "Equal") {
168 return MutableHashMap.get(map.map, key).valueOrUndefined
169 }
170 return map.map.get(key)
171}
172const backingSet = <K, A, E>(map: BackingMap<K, A, E>, key: K, entry: Entry<A, E>): void => {
173 if (map._tag === "Equal") {
174 MutableHashMap.set(map.map, key, entry)
175 } else {
176 map.map.set(key, entry)
177 }
178}
179const backingDelete = <K, A, E>(map: BackingMap<K, A, E>, key: K): void => {
180 if (map._tag === "Equal") {
181 MutableHashMap.remove(map.map, key)
182 } else {
183 map.map.delete(key)
184 }
185}