Fix Query state

packages/effect-fc-next/src/Query.ts

@@ -1,4 +1,4 @@
-import { type Cause, type Context, Duration, Effect, Equal, type Equivalence, Exit, Fiber, Option, Pipeable, Predicate, type Scope, Semaphore, Stream, SubscriptionRef } from "effect"
+import { type Cause, type Context, Duration, Effect, Equal, type Equivalence, Exit, Fiber, Option, Pipeable, Predicate, PubSub, Ref, type Scope, Semaphore, Stream, SubscriptionRef } from "effect"
 import { AsyncResult } from "effect/unstable/reactivity"
 import * as Lens from "./Lens.js"
 import * as QueryClient from "./QueryClient.js"
@@ -231,10 +231,12 @@ extends Pipeable.Class implements Query<K, A, E, R> {
         Scope.Scope | R
     > {
         return Effect.gen({ self: this }, function*() {
-            const state = Lens.fromSubscriptionRef(yield* SubscriptionRef.make(previous))
+            const state = yield* makeQueryStateLens(previous)
 
             const fiber = yield* Effect.forkScoped(Effect.andThen(
-                Lens.update(state, previous => AsyncResult.match(previous.result, {
+                Lens.update<QueryState<K, A, E>, never, never, never, never>(
+                    state,
+                    previous => AsyncResult.match(previous.result, {
                         onInitial: () => ({
                             key: previous.key,
                             result: AsyncResult.initial(true),
@@ -252,19 +254,17 @@ extends Pipeable.Class implements Query<K, A, E, R> {
                                 previousSuccess: result.previousSuccess,
                             }),
                         }),
-                })),
+                    }
+                )),
 
-                Effect.onExit(this.f(previous.key), exit => Effect.all([
-                    Effect.fiberId,
-                    Lens.get(this.fiber),
-                ]).pipe(
-                    Effect.flatMap(([fiberId, fiber]) => Option.match(fiber, {
-                        onSome: v => fiberId === v.id
-                            ? Lens.set(this.fiber, Option.none())
-                            : Effect.void,
-                        onNone: () => Effect.void,
-                    })),
-                    Effect.andThen(Lens.updateAndGet(
+                Effect.onExit(this.f(previous.key), exit => Effect.gen({ self: this }, function*() {
+                    const fiberId = yield* Effect.fiberId
+                    const fiber = yield* Lens.get(this.fiber)
+
+                    if (Option.isSome(fiber) && fiberId === fiber.value.id)
+                        yield* Lens.set(this.fiber, Option.none())
+
+                    const finalState = (yield* Lens.updateAndGet<QueryState<K, A, E>, never, never, never, never>(
                         state,
                         previous => Exit.match(exit, {
                             onSuccess: v => ({
@@ -290,9 +290,11 @@ extends Pipeable.Class implements Query<K, A, E, R> {
                                 }),
                             }),
                         }),
-                    )),
-                    Effect.flatMap(finalState => Lens.set(this.latestFinalState, Option.some(finalState as FinalQueryState<K, A, E>))),
-                )),
+                    )) as FinalQueryState<K, A, E>
+
+                    yield* Lens.set(this.latestFinalState, Option.some(finalState))
+                    yield* PubSub.shutdown(state.pubsub)
+                }))
             ))
 
             yield* Lens.set(this.fiber, Option.some(fiber))
@@ -416,3 +418,44 @@ export const service = <K, A, E = never, R = never>(
     make(options),
     query => Effect.forkScoped(query.run),
 )
+
+
+export class QueryStateLens<in out K, in out A, in out E = never>
+extends Lens.LensImpl<QueryState<K, A, E>, never, never, never, never> {
+    constructor(
+        readonly ref: Ref.Ref<QueryState<K, A, E>>,
+        readonly pubsub: PubSub.PubSub<QueryState<K, A, E>>,
+        readonly semaphore: Semaphore.Semaphore,
+    ) {
+        super()
+    }
+
+    get resolve(): Effect.Effect<Lens.LensImpl.Resolved<QueryState<K, A, E>>, never, never> {
+        return Effect.map(
+            Ref.get(this.ref),
+            value => ({
+                value,
+                commit: next => Effect.flatMap(
+                    next,
+                    value => Effect.andThen(
+                        Ref.set(this.ref, value),
+                        PubSub.publish(this.pubsub, value),
+                    ),
+                ),
+            }),
+        )
+    }
+    get changes() { return Stream.fromPubSub(this.pubsub) }
+    get lock() { return Effect.succeed(this.semaphore.withPermit) }
+}
+
+export const makeQueryStateLens = <K, A, E = never>(
+    initial: QueryState<K, A, E>,
+) => Effect.all([
+    Ref.make(initial),
+    PubSub.unbounded<QueryState<K, A, E>>({ replay: 1 }),
+    Semaphore.make(1),
+]).pipe(
+    Effect.tap(([, pubsub]) => PubSub.publish(pubsub, initial)),
+    Effect.map(([ref, pubsub, semaphore]) => new QueryStateLens(ref, pubsub, semaphore)),
+)