1import { RegistryContext, useAtomValue } from "@effect/atom-react";
2import { chooseLoadBalancedEnvironment } from "@t3tools/client-runtime/load-balancing";
3import { executeAtomQuery } from "@t3tools/client-runtime/state/runtime";
4import type { EnvironmentId } from "@t3tools/contracts";
5import { Atom } from "effect/reactivity";
6import { useContext } from "react";
7
8import { serverEnvironment } from "../state/server";
9
10// Mirrors the freshness window in chooseLoadBalancedEnvironment.
11const HOST_RESOURCES_FRESH_MS = 15_000;
12
13// Keyed by content: callers derive the list every render.
14const hostResourcesAtom = Atom.family((environmentIdsKey: string) =>
15 Atom.make((get) => {
16 const evaluatedAt = Date.now();
17 const resources = (
18 environmentIdsKey === "" ? [] : (environmentIdsKey.split("\n") as EnvironmentId[])
19 ).map((environmentId) => {
20 const result = get(serverEnvironment.hostResources({ environmentId, input: {} }));
21 return {
22 environmentId,
23 resources: result._tag === "Success" ? result.value : null,
24 receivedAt: result._tag === "Success" ? result.timestamp : 0,
25 pending: result._tag === "Initial" || result.waiting,
26 failed: result._tag === "Failure",
27 };
28 });
29 // Re-read when the next fresh snapshot ages out, so a frozen evaluatedAt never selects it.
30 const nextExpiry = Math.min(
31 ...resources
32 .map((resource) => resource.receivedAt + HOST_RESOURCES_FRESH_MS)
33 .filter((expiresAt) => expiresAt >= evaluatedAt),
34 );
35 if (Number.isFinite(nextExpiry)) {
36 const timer = setTimeout(() => get.refreshSelf(), nextExpiry - evaluatedAt + 1);
37 get.addFinalizer(() => clearTimeout(timer));
38 }
39 return { evaluatedAt, resources };
40 }),
41);
42
43/** Only mounted for unresolved automatic drafts, so idle clients do not poll hosts. */
44export function useLoadBalancedEnvironment(
45 environmentIds: readonly EnvironmentId[],
46 weights: Readonly<Record<string, number>>,
47) {
48 const registry = useContext(RegistryContext);
49 const refresh = (ids: readonly EnvironmentId[]) => {
50 for (const environmentId of ids) {
51 registry.refresh(serverEnvironment.hostResources({ environmentId, input: {} }));
52 }
53 };
54 const refreshSelection = async (ids: readonly EnvironmentId[]) => {
55 const candidates = await Promise.all(
56 ids.map(async (environmentId) => {
57 const result = await executeAtomQuery(
58 registry,
59 serverEnvironment.hostResources({ environmentId, input: {} }),
60 { refresh: true, reportFailure: false, signal: AbortSignal.timeout(5_000) },
61 );
62 return {
63 environmentId,
64 resources: result._tag === "Success" ? result.value : null,
65 receivedAt: Date.now(),
66 weight: weights[environmentId] ?? 50,
67 };
68 }),
69 );
70 return chooseLoadBalancedEnvironment(candidates, Date.now()) as EnvironmentId | null;
71 };
72 const { evaluatedAt, resources } = useAtomValue(hostResourcesAtom(environmentIds.join("\n")));
73 const pending = resources.some((resource) => resource.pending);
74 const environmentId = chooseLoadBalancedEnvironment(
75 resources.map((resource) => ({
76 ...resource,
77 weight: weights[resource.environmentId] ?? 50,
78 })),
79 evaluatedAt,
80 ) as EnvironmentId | null;
81 return {
82 refresh,
83 refreshSelection,
84 pending,
85 environmentId,
86 failed: !pending && environmentId === null && resources.some((resource) => resource.failed),
87 };
88}