581 lines · 19.7 KB · TypeScript
RawDownload
import { EnvironmentId } from "@t3tools/contracts";import * as Context from "effect/Context";import * as Effect from "effect/Effect";import * as Equal from "effect/Equal";import * as Exit from "effect/Exit";import * as Layer from "effect/Layer";import * as Option from "effect/Option";import * as Ref from "effect/Ref";import * as Schema from "effect/Schema";import * as Scope from "effect/Scope";import * as Semaphore from "effect/Semaphore";import * as Stream from "effect/Stream";import * as SubscriptionRef from "effect/SubscriptionRef"; import * as ClientCapabilities from "../platform/capabilities.ts";import {  type ConnectionCatalogEntry,  type ConnectionRegistration,  type PrimaryConnectionRegistration,  SshConnectionProfile,  connectionRegistrationCatalogEntry,} from "./catalog.ts";import * as ConnectionProfileStore from "./profileStore.ts";import * as Connectivity from "./connectivity.ts";import type {  ConnectionAttemptError,  ConnectionTarget,  NetworkStatus,  SupervisorConnectionState,} from "./model.ts";import * as Persistence from "../platform/persistence.ts";import * as EnvironmentSupervisor from "./supervisor.ts";import * as ConnectionDriver from "./driver.ts";import * as ConnectionWakeups from "./wakeups.ts"; const isSshConnectionProfile = Schema.is(SshConnectionProfile); export class EnvironmentNotRegisteredError extends Schema.TaggedErrorClass<EnvironmentNotRegisteredError>()(  "EnvironmentNotRegisteredError",  {    environmentId: EnvironmentId,  },) {  override get message(): string {    return `Environment ${this.environmentId} is not registered.`;  }} export class PlatformEnvironmentRemovalError extends Schema.TaggedErrorClass<PlatformEnvironmentRemovalError>()(  "PlatformEnvironmentRemovalError",  {    environmentId: EnvironmentId,  },) {  override get message(): string {    return `Platform-managed environment ${this.environmentId} cannot be removed.`;  }} export class EnvironmentRegistry extends Context.Service<  EnvironmentRegistry,  {    readonly entries: SubscriptionRef.SubscriptionRef<      ReadonlyMap<EnvironmentId, ConnectionCatalogEntry>    >;    readonly networkStatus: SubscriptionRef.SubscriptionRef<NetworkStatus>;    readonly start: Effect.Effect<void>;    readonly register: (      registration: ConnectionRegistration,    ) => Effect.Effect<void, Persistence.ConnectionPersistenceError>;    readonly registerPlatform: (registration: PrimaryConnectionRegistration) => Effect.Effect<void>;    readonly remove: (      environmentId: EnvironmentId,    ) => Effect.Effect<      void,      | Persistence.ConnectionPersistenceError      | ConnectionAttemptError      | EnvironmentNotRegisteredError      | PlatformEnvironmentRemovalError    >;    readonly removeRelayEnvironments: () => Effect.Effect<      void,      | Persistence.ConnectionPersistenceError      | ConnectionAttemptError      | PlatformEnvironmentRemovalError    >;    readonly retryNow: (environmentId: EnvironmentId) => Effect.Effect<void>;    readonly state: (      environmentId: EnvironmentId,    ) => Effect.Effect<SupervisorConnectionState, EnvironmentNotRegisteredError>;    readonly stateChanges: (      environmentId: EnvironmentId,    ) => Stream.Stream<SupervisorConnectionState, EnvironmentNotRegisteredError>;    readonly run: <A, E, R>(      environmentId: EnvironmentId,      effect: Effect.Effect<A, E, R>,    ) => Effect.Effect<      A,      E | EnvironmentNotRegisteredError,      Exclude<R, EnvironmentSupervisor.EnvironmentSupervisor>    >;    readonly runStream: <A, E, R>(      environmentId: EnvironmentId,      stream: Stream.Stream<A, E, R>,    ) => Stream.Stream<      A,      E | EnvironmentNotRegisteredError,      Exclude<R, EnvironmentSupervisor.EnvironmentSupervisor>    >;    readonly followStream: <A, E, R>(      environmentId: EnvironmentId,      stream: Stream.Stream<A, E, R>,    ) => Stream.Stream<A, E, Exclude<R, EnvironmentSupervisor.EnvironmentSupervisor>>;  }>()("@t3tools/client-runtime/connection/registry/EnvironmentRegistry") {} interface EnvironmentServiceScope {  readonly entry: ConnectionCatalogEntry;  readonly supervisor: EnvironmentSupervisor.EnvironmentSupervisor["Service"];  readonly scope: Scope.Closeable;} export const make = Effect.gen(function* () {  const storage = yield* Persistence.ConnectionTargetStore;  const registrations = yield* Persistence.ConnectionRegistrationStore;  const cache = yield* Persistence.EnvironmentCacheStore;  const ownedDataCleanup = yield* Persistence.EnvironmentOwnedDataCleanup;  const profiles = yield* ConnectionProfileStore.ConnectionProfileStore;  const connectivity = yield* Connectivity.Connectivity;  const driver = yield* ConnectionDriver.ConnectionDriver;  const wakeups = yield* ConnectionWakeups.ConnectionWakeups;  const ssh = yield* ClientCapabilities.SshEnvironmentGateway;  const persistedTargets = yield* storage.list;  const initialEntries = new Map(    yield* Effect.forEach(      persistedTargets,      Effect.fn("EnvironmentRegistry.loadCatalogEntry")(function* (target) {        const profile =          target._tag === "BearerConnectionTarget" || target._tag === "SshConnectionTarget"            ? yield* profiles.get(target.connectionId)            : Option.none();        return [          target.environmentId,          { target, profile } satisfies ConnectionCatalogEntry,        ] as const;      }),      { concurrency: "unbounded" },    ),  );  const entries =    yield* SubscriptionRef.make<ReadonlyMap<EnvironmentId, ConnectionCatalogEntry>>(initialEntries);  const networkStatus = yield* SubscriptionRef.make(yield* connectivity.status);  const serviceScopes = yield* SubscriptionRef.make<    ReadonlyMap<EnvironmentId, EnvironmentServiceScope>  >(new Map());  const platformEnvironmentIds = yield* Ref.make<ReadonlySet<EnvironmentId>>(new Set());  const persistedTargetsByEnvironment = yield* Ref.make<    ReadonlyMap<EnvironmentId, ConnectionTarget>  >(new Map(persistedTargets.map((target) => [target.environmentId, target])));  interface LeaseLock {    readonly semaphore: Semaphore.Semaphore;    readonly users: number;  }   const leaseLocks = yield* Ref.make<ReadonlyMap<EnvironmentId, LeaseLock>>(new Map());  const leaseLocksGuard = yield* Semaphore.make(1);  const started = yield* Ref.make(false);   const withLeaseLock = <A, E, R>(    environmentId: EnvironmentId,    effect: Effect.Effect<A, E, R>,  ): Effect.Effect<A, E, R> =>    Effect.acquireUseRelease(      leaseLocksGuard.withPermits(1)(        Effect.gen(function* () {          const current = yield* Ref.get(leaseLocks);          const existing = current.get(environmentId);          if (existing !== undefined) {            yield* Ref.set(              leaseLocks,              new Map(current).set(environmentId, {                semaphore: existing.semaphore,                users: existing.users + 1,              }),            );            return existing.semaphore;          }          const semaphore = yield* Semaphore.make(1);          yield* Ref.set(leaseLocks, new Map(current).set(environmentId, { semaphore, users: 1 }));          return semaphore;        }),      ),      (semaphore) => semaphore.withPermits(1)(effect),      (semaphore) =>        leaseLocksGuard.withPermits(1)(          Ref.update(leaseLocks, (current) => {            const existing = current.get(environmentId);            if (existing === undefined || existing.semaphore !== semaphore) {              return current;            }            const next = new Map(current);            if (existing.users === 1) {              next.delete(environmentId);            } else {              next.set(environmentId, {                semaphore,                users: existing.users - 1,              });            }            return next;          }),        ),    ).pipe(Effect.withSpan("EnvironmentRegistry.withLeaseLock"));   const getEntry = Effect.fn("EnvironmentRegistry.getEntry")(function* (    environmentId: EnvironmentId,  ) {    const entry = (yield* SubscriptionRef.get(entries)).get(environmentId);    if (entry === undefined) {      return yield* new EnvironmentNotRegisteredError({        environmentId,      });    }    return entry;  });   const closeServiceScope = Effect.fn("EnvironmentRegistry.closeServiceScope")(function* (    environmentId: EnvironmentId,  ) {    const current = yield* SubscriptionRef.get(serviceScopes);    const lease = current.get(environmentId);    if (lease === undefined) {      return;    }    const next = new Map(current);    next.delete(environmentId);    yield* SubscriptionRef.set(serviceScopes, next);    yield* Scope.close(lease.scope, Exit.void);  });   const createServiceScope = Effect.fn("EnvironmentRegistry.createServiceScope")(    (entry: ConnectionCatalogEntry) =>      Effect.uninterruptible(        Effect.gen(function* () {          const environmentId = entry.target.environmentId;          const scope = yield* Scope.make();          const supervisor = yield* EnvironmentSupervisor.make(entry, {            initiallyDesired: false,          }).pipe(            Effect.provideService(Connectivity.Connectivity, connectivity),            Effect.provideService(ConnectionDriver.ConnectionDriver, driver),            Effect.provideService(ConnectionWakeups.ConnectionWakeups, wakeups),            Scope.provide(scope),            Effect.onError(() => Scope.close(scope, Exit.void)),          );          yield* supervisor.connect;          yield* SubscriptionRef.update(serviceScopes, (current) => {            const next = new Map(current);            next.set(environmentId, { entry, supervisor, scope });            return next;          });          return supervisor;        }),      ),  );   const acquireSupervisor = Effect.fn("EnvironmentRegistry.acquireSupervisor")(function* (    environmentId: EnvironmentId,  ) {    return yield* withLeaseLock(      environmentId,      Effect.gen(function* () {        const entry = yield* getEntry(environmentId);        const existing = (yield* SubscriptionRef.get(serviceScopes)).get(environmentId);        if (existing !== undefined) {          if (Equal.equals(existing.entry, entry)) {            return existing.supervisor;          }          yield* closeServiceScope(environmentId);        }        return yield* createServiceScope(entry);      }),    );  });   const run: EnvironmentRegistry["Service"]["run"] = Effect.fn("EnvironmentRegistry.run")(    function* <A, E, R>(environmentId: EnvironmentId, effect: Effect.Effect<A, E, R>) {      const supervisor = yield* acquireSupervisor(environmentId);      return yield* Effect.provideService(        effect,        EnvironmentSupervisor.EnvironmentSupervisor,        supervisor,      );    },  );   const runStream: EnvironmentRegistry["Service"]["runStream"] = <A, E, R>(    environmentId: EnvironmentId,    stream: Stream.Stream<A, E, R>,  ) =>    Stream.unwrap(      acquireSupervisor(environmentId).pipe(        Effect.map((supervisor) =>          Stream.provideService(stream, EnvironmentSupervisor.EnvironmentSupervisor, supervisor),        ),      ),    );   const followStream: EnvironmentRegistry["Service"]["followStream"] = <A, E, R>(    environmentId: EnvironmentId,    stream: Stream.Stream<A, E, R>,  ) =>    Stream.concat(      Stream.fromEffect(SubscriptionRef.get(entries)),      SubscriptionRef.changes(entries),    ).pipe(      Stream.map((current) => Option.fromUndefinedOr(current.get(environmentId))),      Stream.changes,      Stream.switchMap(        Option.match({          onNone: () => Stream.empty,          onSome: () =>            Stream.unwrap(              acquireSupervisor(environmentId).pipe(                Effect.match({                  onFailure: () => Stream.empty,                  onSuccess: (supervisor) =>                    Stream.provideService(                      stream,                      EnvironmentSupervisor.EnvironmentSupervisor,                      supervisor,                    ),                }),              ),            ),        }),      ),    );   const start = Effect.gen(function* () {    if (yield* Ref.getAndSet(started, true)) {      return;    }    yield* Effect.forEach(      persistedTargets,      (target) =>        acquireSupervisor(target.environmentId).pipe(          Effect.catchTag("EnvironmentNotRegisteredError", () => Effect.void),        ),      {        concurrency: "unbounded",        discard: true,      },    );  }).pipe(Effect.withSpan("EnvironmentRegistry.start"));   const installEntryLocked = Effect.fn("EnvironmentRegistry.installEntryLocked")(function* (    entry: ConnectionCatalogEntry,    options?: { readonly retainEquivalentRuntime?: boolean },  ) {    const target = entry.target;    const previous = (yield* SubscriptionRef.get(entries)).get(target.environmentId);    const existingScope = (yield* SubscriptionRef.get(serviceScopes)).get(target.environmentId);    if (      options?.retainEquivalentRuntime === true &&      previous !== undefined &&      Equal.equals(previous, entry) &&      existingScope !== undefined &&      Equal.equals(existingScope.entry, entry)    ) {      return;    }     yield* closeServiceScope(target.environmentId);    yield* SubscriptionRef.update(entries, (current) => {      const next = new Map(current);      next.set(target.environmentId, entry);      return next;    });    yield* createServiceScope(entry);  });   const register = Effect.fn("EnvironmentRegistry.register")(function* (    registration: ConnectionRegistration,  ) {    const entry = connectionRegistrationCatalogEntry(registration);    const environmentId = entry.target.environmentId;    yield* withLeaseLock(      environmentId,      Effect.gen(function* () {        if ((yield* Ref.get(platformEnvironmentIds)).has(environmentId)) {          return;        }        yield* registrations.register(registration);        yield* Ref.update(persistedTargetsByEnvironment, (current) => {          const next = new Map(current);          next.set(environmentId, registration.target);          return next;        });        yield* installEntryLocked(entry);      }),    );  });   const registerPlatform = Effect.fn("EnvironmentRegistry.registerPlatform")(function* (    registration: PrimaryConnectionRegistration,  ) {    const entry = connectionRegistrationCatalogEntry(registration);    const target = entry.target;    yield* withLeaseLock(      target.environmentId,      Effect.gen(function* () {        yield* Ref.update(platformEnvironmentIds, (current) => {          const next = new Set(current);          next.add(target.environmentId);          return next;        });         const persistedTarget = (yield* Ref.get(persistedTargetsByEnvironment)).get(          target.environmentId,        );        if (persistedTarget !== undefined) {          yield* registrations.remove(persistedTarget).pipe(            Effect.tap(() =>              Ref.update(persistedTargetsByEnvironment, (current) => {                const next = new Map(current);                next.delete(target.environmentId);                return next;              }),            ),            Effect.catch((error) =>              Effect.logWarning(                "Could not remove a persisted registration shadowed by the primary environment.",                {                  environmentId: target.environmentId,                  error,                },              ),            ),          );        }         yield* installEntryLocked(entry, { retainEquivalentRuntime: true });      }),    );  });   const remove = Effect.fn("EnvironmentRegistry.remove")(function* (environmentId: EnvironmentId) {    return yield* withLeaseLock(      environmentId,      Effect.gen(function* () {        if ((yield* Ref.get(platformEnvironmentIds)).has(environmentId)) {          return yield* new PlatformEnvironmentRemovalError({            environmentId,          });        }        const target = (yield* getEntry(environmentId)).target;        const profile =          target._tag === "BearerConnectionTarget" || target._tag === "SshConnectionTarget"            ? yield* profiles.get(target.connectionId)            : Option.none();         yield* registrations.remove(target);        yield* Ref.update(persistedTargetsByEnvironment, (current) => {          const next = new Map(current);          next.delete(environmentId);          return next;        });        yield* closeServiceScope(environmentId);        yield* SubscriptionRef.update(entries, (current) => {          const next = new Map(current);          next.delete(environmentId);          return next;        });        yield* Effect.all(          [            cache.clear(environmentId).pipe(              Effect.catch((error) =>                Effect.logWarning("Could not clear cached environment data after removal.", {                  environmentId,                  error,                }),              ),            ),            ownedDataCleanup.clear(environmentId),          ],          { concurrency: "unbounded", discard: true },        );         if (          target._tag === "SshConnectionTarget" &&          Option.isSome(profile) &&          isSshConnectionProfile(profile.value)        ) {          yield* ssh.disconnect(profile.value.target).pipe(            Effect.tapError((error) =>              Effect.logWarning("Could not disconnect the managed SSH environment.", {                environmentId,                error,              }),

Lines 1–500 of 581.