FrameGraph: Fix some memory leaks and a few errors in NRGE + switch build methods to async (#17487)

Author: Popov72

packages/dev/core/src/FrameGraph/Node/nodeRenderGraph.ts

@@ -11,6 +11,7 @@ import type {
     WritableObject,
     IShadowLight,
     INodeRenderGraphCustomBlockDescription,
+    Immutable,
 } from "core/index";
 import { Observable } from "../../Misc/observable";
 import { NodeRenderGraphOutputBlock } from "./Blocks/outputBlock";
@@ -83,7 +84,12 @@ export class NodeRenderGraph {
     public attachedBlocks: NodeRenderGraphBlock[] = [];
 
     /**
-     * Observable raised when the node render graph is built
+     * Observable raised before the node render graph is built
+     */
+    public onBeforeBuildObservable = new Observable<FrameGraph>();
+
+    /**
+     * Observable raised after the node render graph is built
      * Note that this is the same observable as the one in the underlying FrameGraph!
      */
     public get onBuildObservable() {
@@ -136,6 +142,13 @@ export class NodeRenderGraph {
         return this._scene;
     }
 
+    /**
+     * Gets the options used to create this node render graph
+     */
+    public get options(): Immutable<INodeRenderGraphCreateOptions> {
+        return this._options;
+    }
+
     /**
      * Creates a new node render graph
      * @param name defines the name of the node render graph
@@ -163,9 +176,7 @@ export class NodeRenderGraph {
 
         if (options.rebuildGraphOnEngineResize) {
             this._resizeObserver = this._engine.onResizeObservable.add(async () => {
-                this.build();
-
-                await this.whenReadyAsync();
+                await this.buildAsync(false, true, false);
             });
         }
     }
@@ -285,14 +296,29 @@ export class NodeRenderGraph {
     }
 
     /**
-     * Build the final list of blocks that will be executed by the "execute" method
+     * @deprecated Use buildAsync instead
      * @param dontBuildFrameGraph If the underlying frame graph should not be built (default: false)
      */
-    public build(dontBuildFrameGraph = false) {
+    public build(dontBuildFrameGraph = false): void {
+        void this.buildAsync(dontBuildFrameGraph, false, false);
+    }
+
+    /**
+     * Build the final list of blocks that will be executed by the "execute" method.
+     * It also builds the underlying frame graph unless specified otherwise.
+     * @param dontBuildFrameGraph If the underlying frame graph should not be built (default: false)
+     * @param waitForReadiness If the method should wait for the frame graph to be ready before resolving (default: true). Note that this parameter has no effect if "dontBuildFrameGraph" is true.
+     * @param setAsSceneFrameGraph If the built frame graph must be set as the scene's frame graph (default: true)
+     */
+    public async buildAsync(dontBuildFrameGraph = false, waitForReadiness = true, setAsSceneFrameGraph = true): Promise<void> {
         if (!this.outputBlock) {
             throw new Error("You must define the outputBlock property before building the node render graph");
         }
 
+        if (setAsSceneFrameGraph) {
+            this._scene.frameGraph = this._frameGraph;
+        }
+
         this._initializeBlock(this.outputBlock);
 
         this._frameGraph.clear();
@@ -306,6 +332,8 @@ export class NodeRenderGraph {
             this._autoFillExternalInputs();
         }
 
+        this.onBeforeBuildObservable.notifyObservers(this._frameGraph);
+
         // Make sure that one of the object renderer is flagged as the main object renderer
         const objectRendererBlocks = this.getBlocksByPredicate<NodeRenderGraphBaseObjectRendererBlock>((block) => block instanceof NodeRenderGraphBaseObjectRendererBlock);
         if (objectRendererBlocks.length > 0 && !objectRendererBlocks.find((block) => block.isMainObjectRenderer)) {
@@ -316,7 +344,7 @@ export class NodeRenderGraph {
             this.outputBlock.build(state);
 
             if (!dontBuildFrameGraph) {
-                this._frameGraph.build();
+                await this._frameGraph.buildAsync(waitForReadiness);
             }
         } finally {
             this._buildId = NodeRenderGraph._BuildIdGenerator++;
@@ -368,16 +396,14 @@ export class NodeRenderGraph {
      * Returns a promise that resolves when the node render graph is ready to be executed
      * This method must be called after the graph has been built (NodeRenderGraph.build called)!
      * @param timeStep Time step in ms between retries (default is 16)
-     * @param maxTimeout Maximum time in ms to wait for the graph to be ready (default is 5000)
+     * @param maxTimeout Maximum time in ms to wait for the graph to be ready (default is 10000)
      * @returns The promise that resolves when the graph is ready
      */
     // eslint-disable-next-line @typescript-eslint/promise-function-async, no-restricted-syntax
-    public whenReadyAsync(timeStep = 16, maxTimeout = 5000): Promise<void> {
+    public async whenReadyAsync(timeStep = 16, maxTimeout = 10000): Promise<void> {
         this._frameGraph.pausedExecution = true;
-        // eslint-disable-next-line github/no-then
-        return this._frameGraph.whenReadyAsync(timeStep, maxTimeout).then(() => {
-            this._frameGraph.pausedExecution = false;
-        });
+        await this._frameGraph.whenReadyAsync(timeStep, maxTimeout);
+        this._frameGraph.pausedExecution = false;
     }
 
     /**
@@ -667,6 +693,7 @@ export class NodeRenderGraph {
 
     /**
      * Makes a duplicate of the current node render graph.
+     * Note that you should call buildAsync() on the returned graph to make it usable.
      * @param name defines the name to use for the new node render graph
      * @returns the new node render graph
      */
@@ -678,7 +705,6 @@ export class NodeRenderGraph {
 
         clone.parseSerializedObject(serializationObject);
         clone._buildId = this._buildId;
-        clone.build();
 
         return clone;
     }
@@ -747,11 +773,11 @@ export class NodeRenderGraph {
      * @param nodeRenderGraphOptions defines options to use when creating the node render graph
      * @returns a new NodeRenderGraph
      */
-    public static CreateDefault(name: string, scene: Scene, nodeRenderGraphOptions?: INodeRenderGraphCreateOptions): NodeRenderGraph {
+    public static async CreateDefaultAsync(name: string, scene: Scene, nodeRenderGraphOptions?: INodeRenderGraphCreateOptions): Promise<NodeRenderGraph> {
         const renderGraph = new NodeRenderGraph(name, scene, nodeRenderGraphOptions);
 
         renderGraph.setToDefault();
-        renderGraph.build();
+        await renderGraph.buildAsync(false, true, false);
 
         return renderGraph;
     }
@@ -769,7 +795,7 @@ export class NodeRenderGraph {
 
         renderGraph.parseSerializedObject(source);
         if (!skipBuild) {
-            renderGraph.build();
+            void renderGraph.buildAsync();
         }
 
         return renderGraph;
@@ -793,12 +819,12 @@ export class NodeRenderGraph {
         skipBuild: boolean = true
     ): Promise<NodeRenderGraph> {
         if (snippetId === "_BLANK") {
-            return Promise.resolve(NodeRenderGraph.CreateDefault("blank", scene, nodeRenderGraphOptions));
+            return NodeRenderGraph.CreateDefaultAsync("blank", scene, nodeRenderGraphOptions);
         }
 
         return new Promise((resolve, reject) => {
             const request = new WebRequest();
-            request.addEventListener("readystatechange", () => {
+            request.addEventListener("readystatechange", async () => {
                 if (request.readyState == 4) {
                     if (request.status == 200) {
                         const snippet = JSON.parse(JSON.parse(request.responseText).jsonPayload);
@@ -813,7 +839,7 @@ export class NodeRenderGraph {
 
                         try {
                             if (!skipBuild) {
-                                nodeRenderGraph.build();
+                                await nodeRenderGraph.buildAsync();
                             }
                             resolve(nodeRenderGraph);
                         } catch (err) {

packages/dev/core/src/FrameGraph/Passes/pass.ts

@@ -37,4 +37,7 @@ export class FrameGraphPass<T extends FrameGraphContext> implements IFrameGraphP
     public _isValid(): Nullable<string> {
         return this._executeFunc !== undefined ? null : "Execute function is not set (call setExecuteFunc to set it)";
     }
+
+    /** @internal */
+    public _dispose() {}
 }

packages/dev/core/src/FrameGraph/Passes/renderPass.ts

@@ -146,4 +146,9 @@ export class FrameGraphRenderPass extends FrameGraphPass<FrameGraphRenderContext
               ? null
               : "Render target and render target depth cannot both be undefined.";
     }
+
+    /** @internal */
+    public override _dispose() {
+        this._frameGraphRenderTarget?.dispose();
+    }
 }

packages/dev/core/src/FrameGraph/frameGraph.ts

@@ -34,6 +34,7 @@ export class FrameGraph implements IDisposable {
     private readonly _tasks: FrameGraphTask[] = [];
     private readonly _passContext: FrameGraphContext;
     private readonly _renderContext: FrameGraphRenderContext;
+    private readonly _initAsyncPromises: Promise<void>[] = [];
     private _currentProcessedTask: FrameGraphTask | null = null;
     private _whenReadyAsyncCancel: Nullable<() => void> = null;
 
@@ -147,6 +148,7 @@ export class FrameGraph implements IDisposable {
         }
 
         this._tasks.push(task);
+        this._initAsyncPromises.push(task.initAsync());
     }
 
     /**
@@ -203,14 +205,37 @@ export class FrameGraph implements IDisposable {
         return pass;
     }
 
+    /** @internal */
+    public async _whenAsynchronousInitializationDoneAsync(): Promise<void> {
+        if (this._initAsyncPromises.length > 0) {
+            await Promise.all(this._initAsyncPromises);
+            this._initAsyncPromises.length = 0;
+        }
+    }
+
+    /**
+     * @deprecated Use buildAsync instead
+     */
+    public build(): void {
+        void this.buildAsync(false);
+    }
+
+    private _built = false; // TODO: to be removed when build() is removed
+
     /**
      * Builds the frame graph.
      * This method should be called after all tasks have been added to the frame graph (FrameGraph.addTask) and before the graph is executed (FrameGraph.execute).
+     * @param waitForReadiness If true, the method will wait for the frame graph to be ready before returning (default is true)
      */
-    public build(): void {
+    public async buildAsync(waitForReadiness = true): Promise<void> {
         this.textureManager._releaseTextures(false);
 
+        this.pausedExecution = true;
+        this._built = false;
+
         try {
+            await this._whenAsynchronousInitializationDoneAsync();
+
             for (const task of this._tasks) {
                 task._reset();
 
@@ -233,12 +258,21 @@ export class FrameGraph implements IDisposable {
                 task.onTexturesAllocatedObservable.notifyObservers(this._renderContext);
             }
 
+            this._built = true;
+
             this.onBuildObservable.notifyObservers(this);
+
+            if (waitForReadiness) {
+                await this.whenReadyAsync();
+            }
         } catch (e) {
             this._tasks.length = 0;
             this._currentProcessedTask = null;
             this.textureManager._isRecordingTask = false;
             throw e;
+        } finally {
+            this.pausedExecution = false;
+            this._built = true;
         }
     }
 
@@ -256,14 +290,28 @@ export class FrameGraph implements IDisposable {
     }
 
     /**
-     * Returns a promise that resolves when the frame graph is ready to be executed
-     * This method must be called after the graph has been built (FrameGraph.build called)!
+     * Returns a promise that resolves when the frame graph is ready to be executed.
+     * In general, calling “await buildAsync()” should suffice, as this function also waits for readiness by default.
      * @param timeStep Time step in ms between retries (default is 16)
-     * @param maxTimeout Maximum time in ms to wait for the graph to be ready (default is 5000)
+     * @param maxTimeout Maximum time in ms to wait for the graph to be ready (default is 10000)
      * @returns The promise that resolves when the graph is ready
      */
-    public async whenReadyAsync(timeStep = 16, maxTimeout = 5000): Promise<void> {
+    public async whenReadyAsync(timeStep = 16, maxTimeout = 10000): Promise<void> {
         let firstNotReadyTask: FrameGraphTask | null = null;
+
+        // TODO: to be removed when build() is removed
+        await new Promise((resolve) => {
+            const checkBuilt = () => {
+                if (this._built) {
+                    resolve(void 0);
+                    return;
+                }
+                setTimeout(checkBuilt, 16);
+            };
+            checkBuilt();
+        });
+        // END TODO
+
         return await new Promise((resolve, reject) => {
             this._whenReadyAsyncCancel = _RetryWithInterval(
                 () => {

packages/dev/core/src/FrameGraph/frameGraphRenderTarget.ts

@@ -94,4 +94,8 @@ export class FrameGraphRenderTarget {
 
         return this._renderTargetDepth === other._renderTargetDepth;
     }
+
+    public dispose() {
+        this._renderTargetWrapper?.dispose();
+    }
 }

packages/dev/core/src/FrameGraph/frameGraphTask.ts

@@ -66,6 +66,16 @@ export abstract class FrameGraphTask {
      */
     public abstract record(skipCreationOfDisabledPasses?: boolean): void;
 
+    /**
+     * This function is called once after the task has been added to the frame graph and before the frame graph is built for the first time.
+     * This allows you to initialize asynchronous resources, which is not possible in the constructor.
+     * @returns A promise that resolves when the initialization is complete.
+     */
+    // eslint-disable-next-line @typescript-eslint/promise-function-async, no-restricted-syntax
+    public initAsync(): Promise<void> {
+        return Promise.resolve();
+    }
+
     /**
      * An observable that is triggered after the textures have been allocated.
      */
@@ -112,6 +122,12 @@ export abstract class FrameGraphTask {
 
     /** @internal */
     public _reset() {
+        for (const pass of this._passes) {
+            pass._dispose();
+        }
+        for (const pass of this._passesDisabled) {
+            pass._dispose();
+        }
         this._passes.length = 0;
         this._passesDisabled.length = 0;
     }

packages/dev/core/src/FrameGraph/frameGraphTypes.ts

@@ -97,4 +97,7 @@ export interface IFrameGraphPass {
 
     /** @internal */
     _isValid(): Nullable<string>;
+
+    /** @internal */
+    _dispose(): void;
 }

packages/dev/core/src/Misc/screenshotTools.ts

@@ -714,15 +714,7 @@ export async function CreateScreenshotForFrameGraphAsync(
      */
     engine.getCaps().parallelShaderCompile = undefined;
 
-    frameGraph.build();
-
-    // We don't want the frame graph to render while waiting for whenReadyAsync to complete
-    frameGraph.pausedExecution = true;
-
-    await frameGraph.whenReadyAsync();
-
-    // eslint-disable-next-line require-atomic-updates
-    frameGraph.pausedExecution = false;
+    await frameGraph.buildAsync();
 
     const numberOfFrames = numberOfFramesToRender ?? (textureManager.hasHistoryTextures ? 32 : 1);
 
@@ -752,9 +744,7 @@ export async function CreateScreenshotForFrameGraphAsync(
 
     textureManager.resetBackBufferTextures();
 
-    frameGraph.build();
-
-    await frameGraph.whenReadyAsync();
+    await frameGraph.buildAsync();
 
     // eslint-disable-next-line require-atomic-updates
     frameGraph.pausedExecution = pausedExecution;