Render fallbacks of nested Suspenses when suspend [major]

lib/renderer.js

@@ -47,6 +47,7 @@ class PartialRenderer extends ReactDOMServerRenderer {
 
 		// Init lazy node awaiting counter
 		this.numAwaiting = 0;
+		this.fallbacksQueue = [];
 
 		// Init interrupt signal
 		this.interrupt = null;
@@ -75,6 +76,9 @@ class PartialRenderer extends ReactDOMServerRenderer {
 
 		// Render
 		this.cycle();
+
+		// Process fallbacks requiring render
+		this.drainFallbacksQueue();
 	}
 
 	cycle() {
@@ -88,10 +92,7 @@ class PartialRenderer extends ReactDOMServerRenderer {
 			return;
 		}
 
-		if (this.numAwaiting !== 0) return;
-
-		const {suspenseNode} = this;
-		if (suspenseNode && suspenseNode.suspended) return;
+		if (this.numAwaiting !== 0 || this.fallbacksQueue.length !== 0) return;
 
 		// Finished processing
 		this.destroy();
@@ -105,7 +106,17 @@ class PartialRenderer extends ReactDOMServerRenderer {
 
 	errored(err) {
 		// Abort all promises
-		this.abortDescendents(this.tree);
+		walkTree(this.tree, node => {
+			const {type} = node;
+			if (type === TYPE_SUSPENSE) {
+				if (node.suspended) return walkTree.STOP;
+			} else if (type === TYPE_PROMISE && !node.resolved) {
+				this.abort(node);
+			}
+		});
+
+		// Clear fallbacks queue
+		this.fallbacksQueue.length = 0;
 
 		// Destroy + callback with error
 		// NB Do not leak `this` to callback
@@ -267,6 +278,7 @@ class PartialRenderer extends ReactDOMServerRenderer {
 		node.suspendedAbove = parentSuspense ?
 			parentSuspense.suspended || parentSuspense.suspendedAbove :
 			false;
+		node.containsLazy = false;
 		this.suspenseNode = node;
 
 		this.node = node;
@@ -292,6 +304,9 @@ class PartialRenderer extends ReactDOMServerRenderer {
 		const {stack} = this;
 		const frame = this.stackPopOriginal.call(stack);
 
+		// Flag suspense as containing lazy
+		if (!suspenseNode.containsLazy) suspenseNode.containsLazy = true;
+
 		// If above suspense boundary suspended, suspend this too
 		if (suspenseNode.suspendedAbove && !suspenseNode.suspended) suspenseNode.suspended = true;
 
@@ -374,10 +389,12 @@ class PartialRenderer extends ReactDOMServerRenderer {
 	rerender(node, element) {
 		// Step into node and reinstate stack state with element on stack ready to render
 		this.node = node;
-		const suspenseNode = node.parentSuspense;
-		this.suspenseNode = suspenseNode;
+		this.suspenseNode = node.parentSuspense;
 		this.restoreStack(node.stackState, element);
 
+		// Clear stack state on node
+		node.stackState = null;
+
 		// Render element
 		this.cycle();
 
@@ -387,18 +404,31 @@ class PartialRenderer extends ReactDOMServerRenderer {
 		// Clear contexts added in `.restoreStack()`
 		this.resetProviders();
 
-		// If suspended, render fallback of suspense boundary
-		if (!suspenseNode || !suspenseNode.suspended) return;
+		// Process fallbacks requiring render
+		this.drainFallbacksQueue();
+	}
 
-		// Convert node to fallback and discard children
-		const {fallback} = suspenseNode;
-		this.convertNodeToFallback(suspenseNode);
+	drainFallbacksQueue() {
+		// Process fallbacks queue - convert nodes to fallbacks.
+		// Find first fallback that requires rendering, and render it.
+		const {fallbacksQueue} = this;
+		let suspenseNode, fallback;
+		while (true) { // eslint-disable-line no-constant-condition
+			if (fallbacksQueue.length === 0) return;
 
-		// Render fallback
-		if (isRenderableElement(fallback)) this.rerender(suspenseNode, fallback);
+			// Convert node to fallback and discard children
+			suspenseNode = fallbacksQueue.pop();
+			fallback = suspenseNode.fallback;
+			this.convertNodeToFallback(suspenseNode);
+
+			// Stop when found a fallback that requires rendering
+			if (isRenderableElement(fallback)) break;
 
-		// Clear frame to free memory
-		suspenseNode.frame = null;
+			suspenseNode.stackState = null;
+		}
+
+		// Render fallback
+		this.rerender(suspenseNode, fallback);
 	}
 
 	rejected(node, err) {
@@ -408,6 +438,7 @@ class PartialRenderer extends ReactDOMServerRenderer {
 
 		// Promise is no longer awaited
 		node.resolved = true;
+		node.stackState = null;
 		this.numAwaiting--;
 
 		// Abort with error
@@ -445,12 +476,54 @@ class PartialRenderer extends ReactDOMServerRenderer {
 		// Abort this promise
 		followAndAbort(promise);
 
-		// Abort all promises within boundary
+		// Suspend suspense and add to fallbacks queue if it is outside of current render cycle
 		const {suspenseNode} = this;
-		this.abortDescendents(suspenseNode);
-
-		// Suspend
 		suspenseNode.suspended = true;
+		if (!suspenseNode.frame) this.fallbacksQueue.push(suspenseNode);
+
+		// Abort all promises within Suspense and trigger fallbacks of nested Suspenses
+		for (let child of suspenseNode.children) {
+			this.suspendDescendents(child);
+		}
+	}
+
+	/*
+	 * A Suspense boundary above this node has been suspended.
+	 * Traverse tree and abort all promises.
+	 * Suspend any suspense boundaries nested within this tree which contain lazy nodes.
+	 * Suspense boundaries within lazy nodes are skipped, since these lazy nodes will no
+	 * longer be loaded, and therefore the suspense boundaries within them will not
+	 * be rendered.
+	 * Add suspense nodes requiring fallback to be rendered to the fallbacks queue.
+	 */
+	suspendDescendents(node) {
+		walkTree(node, (node, inLazy) => {
+			const {type} = node;
+			if (type === TYPE_SUSPENSE) {
+				if (node.suspended) return walkTree.STOP;
+
+				if (!inLazy) {
+					if (node.containsLazy) {
+						node.suspended = true;
+						this.fallbacksQueue.push(node);
+					} else {
+						node.suspendedAbove = true;
+					}
+				}
+			} else if (type === TYPE_PROMISE) {
+				if (!node.resolved) this.abort(node);
+				inLazy = true;
+			}
+
+			return inLazy;
+		}, false);
+	}
+
+	abort(node) {
+		node.resolved = true;
+		node.stackState = null;
+		abort(node.promise);
+		this.numAwaiting--;
 	}
 
 	createChildWithStackState(type, frame) {
@@ -530,6 +603,8 @@ class PartialRenderer extends ReactDOMServerRenderer {
 		// Step down suspense stack
 		this.suspenseNode = node.parentSuspense;
 
+		// TODO Clear stack state of Suspense nodes when definitively resolved
+
 		if (!node.suspended) return;
 
 		// Was suspended
@@ -623,27 +698,6 @@ class PartialRenderer extends ReactDOMServerRenderer {
 		this.clearProviders();
 		this.contextIndex = -1;
 	}
-
-	/**
-	 * Abort all descendents (children, children's children etc) of a node.
-	 * If promise on lazy component has `.abort()` method, it is called.
-	 * @param {Object} tree - Starting node
-	 * @returns {undefined}
-	 */
-	abortDescendents(tree) {
-		walkTree(tree, node => {
-			const {type} = node;
-			if (type === TYPE_TEXT) return false;
-
-			if (type === TYPE_PROMISE && !node.resolved) {
-				node.resolved = true;
-				this.numAwaiting--;
-				abort(node.promise);
-			}
-
-			return true;
-		});
-	}
 }
 
 module.exports = PartialRenderer;

lib/treeToHtml.js

@@ -23,13 +23,12 @@ function treeToHtml(tree, makeStaticMarkup) {
 		lastText = false;
 
 	walkTree(tree, node => {
-		if (node.type !== TYPE_TEXT) return true;
+		if (node.type !== TYPE_TEXT) return;
 
 		const thisOut = node.out;
 		if (lastText && thisOut.slice(0, 1) !== '<') out += '<!-- -->';
 		out += thisOut;
 		if (!makeStaticMarkup) lastText = thisOut.slice(-1) !== '>';
-		return false;
 	});
 
 	return out;

lib/walkTree.js

@@ -5,22 +5,32 @@
 
 'use strict';
 
+// Constants
+const STOP = {};
+
 // Exports
 
 /**
  * Walk tree and call `fn()` on each node.
- * Stop walking branch if `fn()` returns `false`.
+ * `fn()` can return state which is passed to its child nodes.
+ * Stop walking branch if `fn()` returns `walkTree.STOP`.
  * @param {Object} node - Boundary tree node
  * @param {Function} fn - Function to call on each node
+ * @param {*} [state] - Optional state to
  * @returns {undefined}
  */
-function walkTree(node, fn) {
-	const continueWalking = fn(node);
-	if (!continueWalking) return;
+function walkTree(node, fn, state) {
+	state = fn(node, state);
+	if (state === STOP) return;
+
+	const {children} = node;
+	if (!children) return;
 
-	for (let child of node.children) {
-		walkTree(child, fn);
+	for (let child of children) {
+		walkTree(child, fn, state);
 	}
 }
 
+walkTree.STOP = STOP;
+
 module.exports = walkTree;