SchemaView performance improvements

core/backend/src/IModelDb.ts

@@ -70,7 +70,7 @@ import { createNoOpLockControl } from "./internal/NoLocks";
 import { IModelDbFonts } from "./IModelDbFonts";
 import { createIModelDbFonts } from "./internal/IModelDbFontsImpl";
 import { _activeTxn, _cache, _close, _hubAccess, _implicitTxn, _instanceKeyCache, _nativeDb, _releaseAllLocks, _resetIModelDb } from "./internal/Symbols";
-import { ECSpecVersion, ECVersion, SchemaContext, SchemaJsonLocater, SchemaView } from "@itwin/ecschema-metadata";
+import { ECSpecVersion, ECVersion, SchemaContext, SchemaJsonLocater, SchemaManifest, type SchemaManifestReferenceRow, type SchemaManifestSchemaRow, SchemaView } from "@itwin/ecschema-metadata";
 import { SchemaMap } from "./Schema";
 import { ElementLRUCache, InstanceKeyLRUCache } from "./internal/ElementLRUCache";
 import { IModelIncrementalSchemaLocater } from "./IModelIncrementalSchemaLocater";
@@ -418,6 +418,25 @@ export interface CloseIModelArgs {
   optimize?: boolean;
 }
 
+/** Options for [[IModelDb.getSchemaView]].
+ * @beta
+ */
+export interface GetSchemaViewArgs {
+  /** When provided, return a view incrementally loaded with only these schemas plus their
+   * transitive reference closure, instead of every schema in the iModel.
+   *
+   * The view accumulates: it is one instance reused across calls, so a later request with different
+   * schemas - or a later call with no filter at all - merges any still-missing schemas into the same
+   * view, and schemas requested earlier remain available. Schemas that are not loaded (and that the
+   * request did not pull in) are simply absent - `findClass` and friends return `undefined` for
+   * them, the same as for a schema the iModel does not contain.
+   *
+   * Names the iModel does not contain are ignored. Omitting this option ensures every schema is
+   * loaded, identical to calling `getSchemaView()` with no arguments.
+   */
+  readonly schemas?: readonly string[];
+}
+
 /** An iModel database file. The database file can either be a briefcase or a snapshot.
  * @see [Accessing iModels]($docs/learning/backend/AccessingIModels.md)
  * @see [About IModelDb]($docs/learning/backend/IModelDb.md)
@@ -446,7 +465,20 @@ export abstract class IModelDb extends IModel {
   private _jsClassMap?: EntityJsClassMap;
   private _schemaMap?: SchemaMap;
   private _schemaContext?: SchemaContext;
-  private _schemasPromise?: Promise<SchemaView>;
+  // Single, accumulating schema view ("husk") - we never hold more than one SchemaView. The strategy
+  // is fixed by the first getSchemaView call: a no-filter first call loads every schema as one full
+  // `PRAGMA schema_view` blob and sets `_allSchemasLoaded` (the manifest and loaded-name set are never
+  // touched); a filtered first call loads only the requested reference closure as a
+  // `PRAGMA schema_view_fragment` blob and uses the manifest + `_loadedSchemaNames` to extend the same
+  // view on later calls. `_schemaManifestPromise` is the cheap reference graph (incremental mode only);
+  // `_loadedSchemaNames` is the synchronous "already loaded" gate (incremental mode only);
+  // `_allSchemasLoaded` short-circuits once everything is in; `_schemaLoad` serializes loads so
+  // overlapping requests cannot double-merge a schema.
+  private _schemaManifestPromise?: Promise<SchemaManifest>;
+  private _schemaHusk?: SchemaView;
+  private readonly _loadedSchemaNames = new Set<string>();
+  private _allSchemasLoaded = false;
+  private _schemaLoad: Promise<void> = Promise.resolve();
   /** @deprecated in 5.0.0 - will not be removed until after 2026-06-13. Use [[fonts]]. */
   protected _fontMap?: FontMap; // eslint-disable-line @typescript-eslint/no-deprecated
   private readonly _fonts: IModelDbFonts = createIModelDbFonts(this);
@@ -1155,11 +1187,14 @@ export abstract class IModelDb extends IModel {
       this._jsClassMap = undefined;
       this._schemaMap = undefined;
       this._schemaContext = undefined;
-      if (this._schemasPromise) {
-        const old = this._schemasPromise;
-        this._schemasPromise = undefined;
-        old.then((view) => view.markOutdated()).catch(() => {});
-      }
+      // Drop the schema view and its manifest so the next request rebuilds against the new schema
+      // state. The loaded-set must be cleared too, or a stale entry would suppress a reload.
+      if (this._schemaHusk)
+        this._schemaHusk.markOutdated();
+      this._schemaHusk = undefined;
+      this._schemaManifestPromise = undefined;
+      this._loadedSchemaNames.clear();
+      this._allSchemasLoaded = false;
       this[_nativeDb].clearECDbCache();
     }
     this.elements[_cache].clear();
@@ -1736,45 +1771,144 @@ export abstract class IModelDb extends IModel {
    * It is the recommended default for runtime read-only metadata access and is significantly
    * faster and lower-memory than [[schemaContext]]. Use [[schemaContext]] for schema authoring,
    * custom-attribute deserialization, or anywhere you need the full ecschema-metadata object graph.
+   *
+   * Pass `args.schemas` to load only a subset (plus its reference closure) instead of every schema.
+   * The view accumulates: a later call merges any still-missing schemas into the same instance; see
+   * [[GetSchemaViewArgs]].
    * @beta
    */
-  public async getSchemaView(): Promise<SchemaView> {
-    if (this._schemasPromise) {
-      const ctx = await this._schemasPromise;
-      if (!ctx.isOutdated)
-        return ctx;
-    }
-    // Capture the in-flight promise locally so the rejection handler only clears
-    // `_schemasPromise` if it still points at this build. A concurrent invalidation +
-    // re-fetch could otherwise replace the field before our hydrate fails, and a naive
-    // `_schemasPromise = undefined` would clobber that newer reference.
-    const inflight = this._hydrateSchemas();
-    this._schemasPromise = inflight;
-    inflight.catch(() => {
-      if (this._schemasPromise === inflight)
-        this._schemasPromise = undefined;
-    });
-    return inflight;
-  }
-
-  private async _hydrateSchemas(): Promise<SchemaView> {
-    // PRAGMA returns exactly one row with format, formatVersion, data (binary), schemaToken.
-    // Important: only call reader.next() once - do NOT use `for await` on PRAGMA results.
-    // ConcurrentQuery wraps regular ECSQL in LIMIT/OFFSET for pagination but skips this for
-    // PRAGMAs. If the serialized result exceeds the memory threshold, the response is marked
-    // "Partial", and a `for await` loop would re-issue the same PRAGMA forever since PRAGMAs
-    // don't support OFFSET-based pagination.
-    // This implementation uses the non-pinned version of the pragma other than frontend - because backend
-    // is always strictly coupled with the native code.
-    const reader = this.createQueryReader("PRAGMA schema_view");
+  public async getSchemaView(args?: GetSchemaViewArgs): Promise<SchemaView> {
+    const schemas = args?.schemas;
+
+    // Synchronous fast path: the single view already holds everything this caller needs. A no-filter
+    // request is satisfied only once every schema is loaded; a filtered request as soon as each
+    // requested name is present. No I/O, no manifest, no flyweight allocation.
+    if (this._schemaHusk !== undefined &&
+      (this._allSchemasLoaded || (schemas !== undefined && schemas.every((name) => this._loadedSchemaNames.has(name.toLowerCase())))))
+      return this._schemaHusk;
+
+    // Otherwise serialize onto any in-flight load so concurrent callers cannot double-merge a schema.
+    // What is actually missing is recomputed inside the continuation (an earlier queued load may have
+    // covered it). Failures are isolated so one rejected load does not poison later callers.
+    const loadSchemasPromise = this._schemaLoad.then(async () => this._ensureSchemasLoaded(schemas));
+    // The next caller chains its load onto `_schemaLoad`, so our stored promise must never reject.
+    // The current caller gets `loadSchemasPromise` which can reject, but our stored promise always resolves.
+    this._schemaLoad = loadSchemasPromise.then(() => undefined, () => undefined);
+    return loadSchemasPromise;
+  }
+
+  /** The body of [[getSchemaView]], run serialized behind `_schemaLoad`. Ensures the requested schemas
+   * (or all schemas, when no filter is given) are loaded into the single accumulating view and returns
+   * it.
+   *
+   * The strategy is fixed by the *first* call:
+   *  - First call with no filter -> load every schema as one full `schema_view` blob. The manifest and
+   *    loaded-name tracking are never needed; `_allSchemasLoaded` is set and all later calls short-circuit.
+   *  - First call with a filter -> load only the requested reference closure as a `schema_view_fragment`
+   *    blob, and keep the manifest + loaded-name set to extend the *same* view on later calls. A later
+   *    no-filter call in this mode completes the view by fetching all remaining schemas as one fragment.
+   */
+  private async _ensureSchemasLoaded(schemas: readonly string[] | undefined): Promise<SchemaView> {
+    // Re-check the fast-path condition now that earlier queued loads have settled.
+    if (this._schemaHusk !== undefined &&
+      (this._allSchemasLoaded || (schemas !== undefined && schemas.every((name) => this._loadedSchemaNames.has(name.toLowerCase())))))
+      return this._schemaHusk;
+
+    const isFirstLoad = this._schemaHusk === undefined;
+    const husk = this._schemaHusk ?? (this._schemaHusk = SchemaView.createMergeable());
+
+    // Full strategy: the very first request wants everything. Fetch the whole iModel as one blob (one
+    // round trip, best cross-schema dedup) and skip the manifest and closure walk entirely.
+    if (isFirstLoad && schemas === undefined) {
+      const { data, token } = await this._fetchSchemaBlob("PRAGMA schema_view");
+      husk.mergeFragment(data);
+      husk.setSchemaToken(token);
+      this._allSchemasLoaded = true;
+      return husk;
+    }
+
+    // Incremental strategy: a filter was given, or the first call established this mode and the caller
+    // now wants more. The manifest and loaded-name set are only ever used here. Compute the
+    // still-missing reference closure and fetch just those schemas as a fragment.
+    const manifest = await this._getSchemaManifest();
+    const loadedIndices = new Set<number>();
+    for (const name of this._loadedSchemaNames) {
+      const entry = manifest.findByName(name);
+      if (entry !== undefined)
+        loadedIndices.add(entry.index);
+    }
+
+    // No filter in incremental mode means "load whatever is left" - the closure of every schema.
+    const requested = schemas ?? manifest.getAvailableSchemaNames();
+    const order = manifest.computeLoadOrder(requested, loadedIndices);
+    if (order.length > 0) {
+      const ids = order.map((entry) => entry.ecInstanceId);
+      const idList = ids.map((id) => Math.trunc(id).toString(10)).join(",");
+      const { data, token } = await this._fetchSchemaBlob(`PRAGMA schema_view_fragment('${idList}')`);
+      husk.mergeFragment(data);
+      husk.setSchemaToken(token);
+      // Record every closure entry as loaded - including *excluded* schemas (e.g. CoreCustomAttributes)
+      // that the writer emits no rows for and so never appear in the view. Tracking names (not view
+      // contents) is what lets a later request's gate and closure prune them instead of re-fetching.
+      for (const entry of order)
+        this._loadedSchemaNames.add(entry.name.toLowerCase());
+    }
+    if (schemas === undefined || this._loadedSchemaNames.size === manifest.schemaCount)
+      this._allSchemasLoaded = true;
+    return husk;
+  }
+
+  /** Load (once) the cheap reference graph of every schema in the iModel from ECDbMeta. No schema
+   * data is hydrated - just names, versions, ids, and the reference edges. */
+  private async _getSchemaManifest(): Promise<SchemaManifest> {
+    if (this._schemaManifestPromise === undefined) {
+      this._schemaManifestPromise = (async (): Promise<SchemaManifest> => {
+        const schemaRows: SchemaManifestSchemaRow[] = [];
+        const schemaSql = "SELECT ECInstanceId as id, Name as name, VersionMajor as versionMajor, VersionWrite as versionWrite, VersionMinor as versionMinor FROM meta.ECSchemaDef";
+        for await (const row of this.createQueryReader(schemaSql, undefined, { rowFormat: QueryRowFormat.UseECSqlPropertyNames })) {
+          schemaRows.push({
+            id: Number(row.id), // ECInstanceId comes back as a hex Id64String; schemas are few with small ids
+            name: row.name,
+            versionMajor: row.versionMajor,
+            versionWrite: row.versionWrite,
+            versionMinor: row.versionMinor,
+          });
+        }
+
+        const referenceRows: SchemaManifestReferenceRow[] = [];
+        const referenceSql = "SELECT SourceECInstanceId as sourceId, TargetECInstanceId as targetId FROM meta.SchemaHasSchemaReferences";
+        for await (const row of this.createQueryReader(referenceSql, undefined, { rowFormat: QueryRowFormat.UseECSqlPropertyNames })) {
+          referenceRows.push({ schemaId: Number(row.sourceId), referencedSchemaId: Number(row.targetId) });
+        }
+
+        return SchemaManifest.fromRows(schemaRows, referenceRows);
+      })();
+      this._schemaManifestPromise.catch(() => {
+        this._schemaManifestPromise = undefined; // let a later call retry on failure
+      });
+    }
+    return this._schemaManifestPromise;
+  }
+
+  /** Fetch one schema-view blob (full or fragment) and return its data and cache-invalidation token.
+   * Both `PRAGMA schema_view` and `PRAGMA schema_view_fragment` return a single row with the same
+   * columns. Ids in a fragment pragma originate from our own manifest and are formatted strictly as
+   * decimal digits (native re-validates) as defense in depth. */
+  private async _fetchSchemaBlob(pragma: string): Promise<{ data: Uint8Array, token: string }> {
+    // Only call reader.next() once - do NOT use `for await` on PRAGMA results. ConcurrentQuery wraps
+    // regular ECSQL in LIMIT/OFFSET for pagination but skips this for PRAGMAs; if the serialized result
+    // exceeds the memory threshold the response is marked "Partial", and a `for await` loop would
+    // re-issue the same PRAGMA forever since PRAGMAs don't support OFFSET-based pagination. The backend
+    // uses the non-pinned pragma (unlike the frontend) because it is strictly coupled with native code.
+    const reader = this.createQueryReader(pragma);
     const result = await reader.next();
     if (result.done)
-      throw new IModelError(DbResult.BE_SQLITE_ERROR, "PRAGMA schema_view returned no rows");
+      throw new IModelError(DbResult.BE_SQLITE_ERROR, `${pragma} returned no rows`);
     const data = result.value.data as Uint8Array | undefined;
     const token = result.value.schemaToken as string | undefined;
     if (data === undefined || data === null)
-      throw new IModelError(DbResult.BE_SQLITE_ERROR, "PRAGMA schema_view returned null data column");
-    return SchemaView.fromBinary(data, token ?? "");
+      throw new IModelError(DbResult.BE_SQLITE_ERROR, `${pragma} returned null data column`);
+    return { data, token: token ?? "" };
   }
 
   /** Get the linkTableRelationships for this IModel */