[perf] make ensureStoreIsUsable scale better (#1362)

The new version of the sync engine is gonna be calling
`ensureStoreIsUsable` on every sync message, so I took some time to make
it scale better. At the moment it operates on a serialized version of
the store, which is expensive and unnecessary. Here I changed it to use
reactive queries for the data it needs, so it only operates on small
bits of data and should not become more expensive as the number of
shapes grows.

### Change Type

<!-- 💡 Indicate the type of change your pull request is. -->
<!-- 🤷‍♀️ If you're not sure, don't select anything -->
<!-- ✂️ Feel free to delete unselected options -->

<!-- To select one, put an x in the box: [x] -->

- [x] `patch` — Bug Fix
- [ ] `minor` — New Feature
- [ ] `major` — Breaking Change

- [ ] `dependencies` — Dependency Update (publishes a `patch` release,
for devDependencies use `internal`)

- [ ] `documentation` — Changes to the documentation only (will not
publish a new version)
- [ ] `tests` — Changes to any testing-related code only (will not
publish a new version)
- [ ] `internal` — Any other changes that don't affect the published
package (will not publish a new version)

### Test Plan

1. Add a step-by-step description of how to test your PR here.
2.

- [ ] Unit Tests
- [ ] Webdriver tests

### Release Notes

- Add a brief release note for your PR here.

packages/tlschema/api-report.md

@@ -88,6 +88,9 @@ export function createAssetValidator<Type extends string, Props extends object>(
 // @public (undocumented)
 export function createCustomShapeId(id: string): TLShapeId;
 
+// @internal (undocumented)
+export function createIntegrityChecker(store: TLStore): () => void;
+
 // @public (undocumented)
 export function createShapeId(): TLShapeId;
 
@@ -336,9 +339,6 @@ export const embedShapeMigrations: Migrations;
 // @public (undocumented)
 export const embedShapeTypeValidator: T.Validator<TLEmbedShape>;
 
-// @internal (undocumented)
-export function ensureStoreIsUsable(store: TLStore): void;
-
 // @internal (undocumented)
 export const fillValidator: T.Validator<"none" | "pattern" | "semi" | "solid">;
 

packages/tlschema/src/TLStore.ts

@@ -99,35 +99,46 @@ function getDefaultPages() {
 }
 
 /** @internal */
-export function ensureStoreIsUsable(store: TLStore): void {
+export function createIntegrityChecker(store: TLStore): () => void {
+	const $pages = store.query.records('page')
+	const $userDocumentSettings = store.query.record('user_document', () => ({
+		userId: { eq: store.props.userId },
+	}))
+
+	const $instanceState = store.query.record('instance', () => ({
+		id: { eq: store.props.instanceId },
+	}))
+
+	const $user = store.query.record('user', () => ({ id: { eq: store.props.userId } }))
+
+	const $userPresences = store.query.records('user_presence')
+	const $instancePageStates = store.query.records('instance_page_state')
+
+	const ensureStoreIsUsable = (): void => {
 		const { userId, instanceId: tabId } = store.props
 		// make sure we have exactly one document
 		if (!store.has(TLDOCUMENT_ID)) {
 			store.put([TLDocument.create({ id: TLDOCUMENT_ID })])
-		return ensureStoreIsUsable(store)
+			return ensureStoreIsUsable()
 		}
 
-	const allRecords = store.allRecords()
-
 		// make sure we have document state for the current user
-	const userDocumentSettings = allRecords
-		.filter(TLUserDocument.isInstance)
-		.find((ud) => ud.userId === userId)
+		const userDocumentSettings = $userDocumentSettings.value
 
 		if (!userDocumentSettings) {
 			store.put([TLUserDocument.create({ userId })])
-		return ensureStoreIsUsable(store)
+			return ensureStoreIsUsable()
 		}
 
 		// make sure there is at least one page
-	const pages = allRecords.filter(TLPage.isInstance).sort(sortByIndex)
+		const pages = $pages.value.sort(sortByIndex)
 		if (pages.length === 0) {
 			store.put(getDefaultPages())
-		return ensureStoreIsUsable(store)
+			return ensureStoreIsUsable()
 		}
 
 		// make sure we have state for the current user's current tab
-	const instanceState = allRecords.filter(TLInstance.isInstance).find((t) => t.id === tabId)
+		const instanceState = $instanceState.value
 		if (!instanceState) {
 			// The tab props are either the the last used tab's props or undefined
 			const propsForNextShape = userDocumentSettings.lastUsedTabId
@@ -147,7 +158,7 @@ export function ensureStoreIsUsable(store: TLStore): void {
 				}),
 			])
 
-		return ensureStoreIsUsable(store)
+			return ensureStoreIsUsable()
 		}
 
 		// make sure the user's currentPageId is still valid
@@ -155,21 +166,19 @@ export function ensureStoreIsUsable(store: TLStore): void {
 		if (!pages.find((p) => p.id === currentPageId)) {
 			currentPageId = pages[0].id!
 			store.put([{ ...instanceState, currentPageId }])
-		return ensureStoreIsUsable(store)
+			return ensureStoreIsUsable()
 		}
 
 		// make sure we have a user state record for the current user
-	if (!allRecords.find((u) => u.id === userId)) {
+		if (!$user.value) {
 			store.put([TLUser.create({ id: userId })])
-		return ensureStoreIsUsable(store)
+			return ensureStoreIsUsable()
 		}
 
-	const userPresences = allRecords
-		.filter(TLUserPresence.isInstance)
-		.filter((r) => r.userId === userId)
+		const userPresences = $userPresences.value.filter((r) => r.userId === userId)
 		if (userPresences.length === 0) {
 			store.put([TLUserPresence.create({ userId, color: getRandomColor() })])
-		return ensureStoreIsUsable(store)
+			return ensureStoreIsUsable()
 		} else if (userPresences.length > 1) {
 			// make sure we don't duplicate user presences
 			store.remove(userPresences.slice(1).map((r) => r.id))
@@ -177,9 +186,9 @@ export function ensureStoreIsUsable(store: TLStore): void {
 
 		// make sure each page has a instancePageState and camera
 		for (const page of pages) {
-		const instancePageStates = allRecords
-			.filter(TLInstancePageState.isInstance)
-			.filter((tps) => tps.pageId === page.id && tps.instanceId === tabId)
+			const instancePageStates = $instancePageStates.value.filter(
+				(tps) => tps.pageId === page.id && tps.instanceId === tabId
+			)
 			if (instancePageStates.length > 1) {
 				// make sure we only have one instancePageState per instance per page
 				store.remove(instancePageStates.slice(1).map((ips) => ips.id))
@@ -189,14 +198,17 @@ export function ensureStoreIsUsable(store: TLStore): void {
 					camera,
 					TLInstancePageState.create({ pageId: page.id, instanceId: tabId, cameraId: camera.id }),
 				])
-			return ensureStoreIsUsable(store)
+				return ensureStoreIsUsable()
 			}
 
 			// make sure the camera exists
 			const camera = store.get(instancePageStates[0].cameraId)
 			if (!camera) {
 				store.put([TLCamera.create({ id: instancePageStates[0].cameraId })])
-			return ensureStoreIsUsable(store)
+				return ensureStoreIsUsable()
 			}
 		}
+	}
+
+	return ensureStoreIsUsable
 }

packages/tlschema/src/createTLSchema.ts

@@ -2,7 +2,7 @@ import { StoreSchema, StoreValidator, createRecordType, defineMigrations } from
 import { T } from '@tldraw/tlvalidate'
 import { Signal } from 'signia'
 import { TLRecord } from './TLRecord'
-import { TLStore, TLStoreProps, ensureStoreIsUsable, onValidationFailure } from './TLStore'
+import { TLStore, TLStoreProps, createIntegrityChecker, onValidationFailure } from './TLStore'
 import { defaultDerivePresenceState } from './defaultDerivePresenceState'
 import { TLAsset } from './records/TLAsset'
 import { TLCamera } from './records/TLCamera'
@@ -112,7 +112,7 @@ export function createTLSchema({
 		{
 			snapshotMigrations: storeMigrations,
 			onValidationFailure,
-			ensureStoreIsUsable,
+			createIntegrityChecker: createIntegrityChecker,
 			derivePresenceState: derivePresenceState ?? defaultDerivePresenceState,
 		}
 	)

packages/tlschema/src/index.ts

@@ -1,7 +1,7 @@
 export { type TLRecord } from './TLRecord'
 export {
 	USER_COLORS,
-	ensureStoreIsUsable,
+	createIntegrityChecker,
 	onValidationFailure,
 	type TLStore,
 	type TLStoreProps,

packages/tlstore/api-report.md

@@ -276,12 +276,12 @@ export class StoreSchema<R extends BaseRecord, P = unknown> {
             createId: any;
         };
     }, options?: StoreSchemaOptions<R, P>): StoreSchema<R, P>;
+    // @internal (undocumented)
+    createIntegrityChecker(store: Store<R, P>): (() => void) | undefined;
     // (undocumented)
     get currentStoreVersion(): number;
     // @internal (undocumented)
     derivePresenceState(store: Store<R, P>): Signal<null | R> | undefined;
-    // @internal (undocumented)
-    ensureStoreIsUsable(store: Store<R, P>): void;
     // (undocumented)
     migratePersistedRecord(record: R, persistedSchema: SerializedSchema, direction?: 'down' | 'up'): MigrationResult<R>;
     // (undocumented)
@@ -308,7 +308,7 @@ export type StoreSchemaOptions<R extends BaseRecord, P> = {
         phase: 'createRecord' | 'initialize' | 'tests' | 'updateRecord';
         recordBefore: null | R;
     }) => R;
-    ensureStoreIsUsable?: (store: Store<R, P>) => void;
+    createIntegrityChecker?: (store: Store<R, P>) => void;
     derivePresenceState?: (store: Store<R, P>) => Signal<null | R>;
 };
 

packages/tlstore/src/lib/Store.ts

@@ -622,9 +622,12 @@ export class Store<R extends BaseRecord = BaseRecord, Props = unknown> {
 		}
 	}
 
+	private _integrityChecker?: () => void | undefined
+
 	/** @internal */
 	ensureStoreIsUsable() {
-		this.schema.ensureStoreIsUsable(this)
+		this._integrityChecker ??= this.schema.createIntegrityChecker(this)
+		this._integrityChecker?.()
 	}
 
 	private _isPossiblyCorrupted = false

packages/tlstore/src/lib/StoreSchema.ts

@@ -48,7 +48,7 @@ export type StoreSchemaOptions<R extends BaseRecord, P> = {
 		recordBefore: R | null
 	}) => R
 	/** @internal */
-	ensureStoreIsUsable?: (store: Store<R, P>) => void
+	createIntegrityChecker?: (store: Store<R, P>) => void
 	/** @internal */
 	derivePresenceState?: (store: Store<R, P>) => Signal<R | null>
 }
@@ -240,8 +240,8 @@ export class StoreSchema<R extends BaseRecord, P = unknown> {
 	}
 
 	/** @internal */
-	ensureStoreIsUsable(store: Store<R, P>): void {
-		this.options.ensureStoreIsUsable?.(store)
+	createIntegrityChecker(store: Store<R, P>): (() => void) | undefined {
+		return this.options.createIntegrityChecker?.(store) ?? undefined
 	}
 
 	/** @internal */