convex/src/browser/sync/local_state.ts

Client for the Convex Cloud

import { convexToJson } from "@convex-dev/common";
import {
  AddQuery,
  RemoveQuery,
  QueryId,
  QuerySetModification,
  QuerySetVersion,
  IdentityVersion,
  Authenticate,
} from "./protocol.js";
import {
  canonicalizeUdfPath,
  QueryToken,
  serializePathAndArgs,
} from "./udf_path_utils.js";

type LocalQuery = {
  id: QueryId;
  canonicalizedUdfPath: string;
  args: any[];
  numSubscribers: number;
};

export class LocalSyncState {
  private nextQueryId: QueryId;
  private querySetVersion: QuerySetVersion;
  private readonly querySet: Map<QueryToken, LocalQuery>;
  private readonly queryIdToToken: Map<QueryId, QueryToken>;
  private identityVersion: IdentityVersion;
  private auth?: { tokenType: "Admin" | "User"; value: string };

  constructor() {
    this.nextQueryId = 0;
    this.querySetVersion = 0;
    this.identityVersion = 0;
    this.querySet = new Map();
    this.queryIdToToken = new Map();
  }

  subscribe(
    udfPath: string,
    args: any[]
  ): {
    queryToken: QueryToken;
    modification: QuerySetModification | null;
    unsubscribe: () => QuerySetModification | null;
  } {
    const canonicalizedUdfPath = canonicalizeUdfPath(udfPath);
    const queryToken = serializePathAndArgs(canonicalizedUdfPath, args);

    const existingEntry = this.querySet.get(queryToken);

    if (existingEntry !== undefined) {
      existingEntry.numSubscribers += 1;
      return {
        queryToken,
        modification: null,
        unsubscribe: () => this.removeSubscriber(queryToken),
      };
    } else {
      const queryId = this.nextQueryId++;
      const query: LocalQuery = {
        id: queryId,
        canonicalizedUdfPath,
        args,
        numSubscribers: 1,
      };
      this.querySet.set(queryToken, query);
      this.queryIdToToken.set(queryId, queryToken);

      const baseVersion = this.querySetVersion;
      const newVersion = ++this.querySetVersion;

      const add: AddQuery = {
        type: "Add",
        queryId,
        udfPath: canonicalizedUdfPath,
        args: args.map(convexToJson),
      };
      const modification: QuerySetModification = {
        type: "ModifyQuerySet",
        baseVersion,
        newVersion,
        modifications: [add],
      };
      return {
        queryToken,
        modification,
        unsubscribe: () => this.removeSubscriber(queryToken),
      };
    }
  }

  queryId(udfPath: string, args: any[]): QueryId | null {
    const canonicalizedUdfPath = canonicalizeUdfPath(udfPath);
    const queryToken = serializePathAndArgs(canonicalizedUdfPath, args);
    const existingEntry = this.querySet.get(queryToken);
    if (existingEntry !== undefined) {
      return existingEntry.id;
    }
    return null;
  }

  setAuth(value: string): Authenticate {
    this.auth = {
      tokenType: "User",
      value: value,
    };
    const baseVersion = this.identityVersion++;
    return {
      type: "Authenticate",
      baseVersion: baseVersion,
      ...this.auth,
    };
  }

  setAdminAuth(value: string): Authenticate {
    this.auth = {
      tokenType: "Admin",
      value: value,
    };
    const baseVersion = this.identityVersion++;
    return {
      type: "Authenticate",
      baseVersion: baseVersion,
      ...this.auth,
    };
  }

  clearAuth(): Authenticate {
    this.auth = undefined;
    const baseVersion = this.identityVersion++;
    return {
      type: "Authenticate",
      tokenType: "None",
      baseVersion: baseVersion,
    };
  }

  queryPath(queryId: QueryId): string | null {
    const pathAndArgs = this.queryIdToToken.get(queryId);
    if (pathAndArgs) {
      return this.querySet.get(pathAndArgs)!.canonicalizedUdfPath;
    }
    return null;
  }

  queryArgs(queryId: QueryId): any[] | null {
    const pathAndArgs = this.queryIdToToken.get(queryId);
    if (pathAndArgs) {
      return this.querySet.get(pathAndArgs)!.args;
    }
    return null;
  }

  queryToken(queryId: QueryId): string | null {
    return this.queryIdToToken.get(queryId) ?? null;
  }

  restart(): [QuerySetModification, Authenticate?] {
    const modifications = [];
    for (const localQuery of this.querySet.values()) {
      const add: AddQuery = {
        type: "Add",
        queryId: localQuery.id,
        udfPath: localQuery.canonicalizedUdfPath,
        args: localQuery.args.map(convexToJson),
      };
      modifications.push(add);
    }
    this.querySetVersion = 1;
    const querySet: QuerySetModification = {
      type: "ModifyQuerySet",
      baseVersion: 0,
      newVersion: 1,
      modifications,
    };
    // If there's no auth, no need to send an update as the server will also start with an unknown identity.
    if (!this.auth) {
      this.identityVersion = 0;
      return [querySet, undefined];
    }
    const authenticate: Authenticate = {
      type: "Authenticate",
      baseVersion: 0,
      ...this.auth,
    };
    this.identityVersion = 1;
    return [querySet, authenticate];
  }

  private removeSubscriber(
    queryToken: QueryToken
  ): QuerySetModification | null {
    const localQuery = this.querySet.get(queryToken)!;

    if (localQuery.numSubscribers > 1) {
      localQuery.numSubscribers -= 1;
      return null;
    } else {
      this.querySet.delete(queryToken);
      this.queryIdToToken.delete(localQuery.id);
      const baseVersion = this.querySetVersion;
      const newVersion = ++this.querySetVersion;
      const remove: RemoveQuery = {
        type: "Remove",
        queryId: localQuery.id,
      };
      return {
        type: "ModifyQuerySet",
        baseVersion,
        newVersion,
        modifications: [remove],
      };
    }
  }
}