webstudio/apps/builder/app/shared/sync-client.ts

import { createNanoEvents } from "nanoevents";
import { createId } from "@webstudio-is/sdk";
import type { Change, Store } from "immerhin";
import type { WritableAtom } from "nanostores";
import type {
  Transaction,
  RevertedTransaction,
  SyncMessage,
  SyncEmitter,
} from "@webstudio-is/sync-client";

export interface SyncStorage {
  name: string;
  sendTransaction(transaction: Transaction): void;
  subscribe(setState: (state: unknown) => void, signal: AbortSignal): void;
}

const subscribeStorages = (
  storages: SyncStorage[],
  setState: (state: unknown) => void,
  signal: AbortSignal
) => {
  if (storages.length === 0) {
    setState(undefined);
    return;
  }
  storages[0].subscribe((value) => {
    if (value === undefined) {
      // invoke the next storage
      subscribeStorages(storages.slice(1), setState, signal);
      return;
    }
    setState(value);
  }, signal);
};

interface SyncObject {
  name: string;
  getState(): unknown;
  setState(state: unknown): void;
  applyTransaction(transaction: Transaction): void;
  revertTransaction(transaction: RevertedTransaction): void;
  subscribe(
    sendTransaction: (transaction: Transaction) => void,
    signal: AbortSignal
  ): void;
}

export class ImmerhinSyncObject implements SyncObject {
  name: string;
  store: Store;
  transformOnSend?: (changes: Change[]) => Change[];
  transformOnReceive?: (changes: Change[]) => Change[];
  onRevert?: (changes: Change[]) => void;
  constructor(
    name: string,
    store: Store,
    transform?: {
      onSend?: (changes: Change[]) => Change[];
      onReceive?: (changes: Change[]) => Change[];
      onRevert?: (changes: Change[]) => void;
    }
  ) {
    this.name = name;
    this.store = store;
    this.transformOnSend = transform?.onSend;
    this.transformOnReceive = transform?.onReceive;
    this.onRevert = transform?.onRevert;
  }
  getState() {
    const state = new Map<string, unknown>();
    for (const [namespace, $store] of this.store.containers) {
      state.set(namespace, $store.get());
    }
    return state;
  }
  setState(state: Map<string, unknown>) {
    for (const [namespace, $store] of this.store.containers) {
      // catch errors triggered by CSP configuration when user put iframe onto canvas
      try {
        // Immer cannot handle Map instances from another realm.
        // Use `clone` to recreate the data with the current realm's classes.
        // This works because the structured clone algorithm skips prototype chains; classes must be defined in both realms.
        $store.set(structuredClone(state.get(namespace)));
      } catch (error) {
        console.error(error);
      }
    }
  }
  applyTransaction(transaction: Transaction<Change[]>) {
    // Immer cannot handle values whose prototypes come from another realm.
    // Recreate received patches with the current realm's classes before applying them.
    const receivedPayload = structuredClone(transaction.payload);
    const payload = this.transformOnReceive
      ? this.transformOnReceive(receivedPayload)
      : receivedPayload;
    this.store.addTransaction(transaction.id, payload, "remote");
  }
  revertTransaction(transaction: RevertedTransaction) {
    const revertedChanges = this.store.transactionManager.currentStack
      .find(({ id }) => id === transaction.id)
      ?.getChanges();
    this.store.revertTransaction(transaction.id);
    if (revertedChanges !== undefined) {
      this.onRevert?.(revertedChanges);
    }
  }
  subscribe(
    sendTransaction: (transaction: Transaction<Change[]>) => void,
    signal: AbortSignal
  ) {
    const unsubscribe = this.store.subscribe((id, payload, source) => {
      if (source === "remote") {
        return;
      }
      const outgoing = this.transformOnSend
        ? this.transformOnSend(payload)
        : payload;
      sendTransaction({ id, object: this.name, payload: outgoing });
    });
    signal.addEventListener("abort", unsubscribe);
  }
}

export class NanostoresSyncObject implements SyncObject {
  name: string;
  store: WritableAtom<unknown>;
  // track the source of "store.set" to avoid cyclic updates
  operation: "local" | "state" | "add" = "local";
  constructor(name: string, store: WritableAtom<unknown>) {
    this.name = name;
    this.store = store;
  }
  getState() {
    return this.store.get();
  }
  setState(state: unknown) {
    this.operation = "state";
    this.store.set(state);
    this.operation = "local";
  }
  applyTransaction(transaction: Transaction) {
    this.operation = "add";
    // catch errors triggered by CSP configuration when user put iframe onto canvas
    try {
      // `instanceof` checks do not work with instances like Map, File, etc., from another realm.
      // Use `clone` to recreate the data with the current realm's classes.
      // This works because the structured clone algorithm skips prototype chains; classes must be defined in both realms.
      this.store.set(structuredClone(transaction.payload));
    } catch (error) {
      console.error(error);
    }
    this.operation = "local";
  }
  revertTransaction(_transaction: RevertedTransaction) {
    // @todo store the list of transactions
  }
  subscribe(
    sendTransaction: (transaction: Transaction) => void,
    signal: AbortSignal
  ) {
    const unsubscribe = this.store.listen((payload) => {
      if (this.operation !== "local") {
        return;
      }
      const transaction = {
        id: createId("nano"),
        object: this.name,
        payload,
      };
      sendTransaction(transaction);
    });
    signal.addEventListener("abort", unsubscribe);
  }
}

export class SyncObjectPool implements SyncObject {
  name = "SyncObjectPool";
  objects = new Map<string, SyncObject>();
  constructor(objects: SyncObject[]) {
    for (const object of objects) {
      this.objects.set(object.name, object);
    }
  }
  getState() {
    const state = new Map<string, unknown>();
    for (const [name, object] of this.objects) {
      state.set(name, object.getState());
    }
    return state;
  }
  setState(state: Map<string, unknown>) {
    for (const [name, object] of this.objects) {
      // merge into existing state to partially restore from storage
      if (state.has(name)) {
        object.setState(state.get(name) as never);
      }
    }
  }
  applyTransaction(transaction: Transaction) {
    this.objects.get(transaction.object)?.applyTransaction(transaction);
  }
  revertTransaction(transaction: RevertedTransaction) {
    this.objects.get(transaction.object)?.revertTransaction(transaction);
  }
  subscribe(
    sendTransaction: (transaction: Transaction) => void,
    signal: AbortSignal
  ) {
    for (const object of this.objects.values()) {
      object.subscribe(sendTransaction, signal);
    }
  }
}

export class NanoEventsSyncEmitter implements SyncEmitter {
  private inner = createNanoEvents<{ message: (m: SyncMessage) => void }>();
  emit(message: SyncMessage): void {
    this.inner.emit("message", message);
  }
  on(handler: (message: SyncMessage) => void): () => void {
    return this.inner.on("message", handler);
  }
}

type SyncClientOptions = {
  role: "leader" | "follower";
  object: SyncObject;
  emitter?: SyncEmitter;
  storages?: SyncStorage[];
};

export class SyncClient {
  clientId = createId("nano");
  role: SyncClientOptions["role"];
  object: SyncObject;
  storages: SyncStorage[];
  emitter: SyncEmitter;
  connection: "disconnected" | "connecting" | "connected" = "disconnected";

  constructor(options: SyncClientOptions) {
    this.role = options.role;
    this.object = options.object;
    this.storages = options.storages ?? [];
    this.emitter = options.emitter ?? new NanoEventsSyncEmitter();
  }

  lead() {
    this.role = "leader";
    this.emitter.emit({
      clientId: this.clientId,
      type: "state",
      state: this.object.getState(),
    });
  }

  connect({ signal, onReady }: { signal: AbortSignal; onReady?: () => void }) {
    const off = this.emitter.on((message) => {
      // ignore own messages
      if (this.clientId === message.clientId) {
        return;
      }
      if (message.type === "connect") {
        if (this.role !== "leader") {
          return;
        }
        this.emitter.emit({
          clientId: this.clientId,
          type: "state",
          state: this.object.getState(),
        });
      }
      if (message.type === "state") {
        if (this.connection === "connected") {
          return;
        }
        this.connection = "connected";
        this.object.setState(message.state);
      }
      if (message.type === "apply") {
        this.object.applyTransaction(message.transaction);
        if (this.role === "leader") {
          for (const storage of this.storages) {
            storage.sendTransaction(message.transaction);
          }
        }
      }
      if (message.type === "revert") {
        this.object.revertTransaction(message.transaction);
      }
    });
    signal.addEventListener("abort", off);

    this.object.subscribe((transaction) => {
      this.emitter.emit({
        clientId: this.clientId,
        type: "apply",
        transaction,
      });
      if (this.role === "leader") {
        for (const storage of this.storages) {
          storage.sendTransaction(transaction);
        }
      }
    }, signal);

    this.emitter.emit({
      clientId: this.clientId,
      type: "connect",
    });
    if (this.connection === "disconnected") {
      this.connection = "connecting";
    }
    if (this.role === "leader") {
      subscribeStorages(
        this.storages,
        (state) => {
          // fallback to default object state
          state ??= this.object.getState();
          this.object.setState(state);
          this.emitter.emit({
            clientId: this.clientId,
            type: "state",
            state,
          });
          onReady?.();
        },
        signal
      );
    }
  }
}