Ensure that all extensions are processed

examples/everything.amp.html

@@ -93,6 +93,7 @@
       border: 1px solid green;
     }
   </style>
+  <script async custom-element="amp-dynamic-css-classes" src="https://cdn.ampproject.org/v0/amp-dynamic-css-classes-0.1.js"></script>
   <script async custom-element="amp-anim" src="https://cdn.ampproject.org/v0/amp-anim-0.1.js"></script>
   <script async custom-element="amp-audio" src="https://cdn.ampproject.org/v0/amp-audio-0.1.js"></script>
   <script async custom-element="amp-carousel" src="https://cdn.ampproject.org/v0/amp-carousel-0.1.js"></script>

src/amp.js

@@ -31,7 +31,6 @@ import {stubElements} from './custom-element';
 import {adopt} from './runtime';
 import {cssText} from '../build/css';
 import {maybeValidate} from './validator-integration';
-import {waitForExtensions} from './render-delaying-extensions';
 
 // We must under all circumstances call makeBodyVisible.
 // It is much better to have AMP tags not rendered than having
@@ -62,7 +61,7 @@ try {
       installGlobalClickListener(window);
 
       maybeValidate(window);
-      makeBodyVisible(document, waitForExtensions(window));
+      makeBodyVisible(document, /* waitForExtensions */ true);
     } catch (e) {
       makeBodyVisible(document);
       throw e;

src/custom-element.js

@@ -108,7 +108,9 @@ export function upgradeOrRegisterElement(win, name, toClass) {
  * @param {!Window} win
  */
 export function stubElements(win) {
-  win.ampExtendedElements = {};
+  if (!win.ampExtendedElements) {
+    win.ampExtendedElements = {};
+  }
   const list = win.document.querySelectorAll('[custom-element]');
   for (let i = 0; i < list.length; i++) {
     const name = list[i].getAttribute('custom-element');
@@ -118,6 +120,10 @@ export function stubElements(win) {
     }
     registerElement(win, name, ElementStub);
   }
+  // Repeat stubbing when HEAD is complete.
+  if (!win.document.body) {
+    dom.waitForBody(win.document, () => stubElements(win));
+  }
 }
 
 
@@ -1171,6 +1177,7 @@ export function resetScheduledElementForTesting(win, elementName) {
   if (win.ampExtendedElements) {
     win.ampExtendedElements[elementName] = null;
   }
+  delete knownElements[elementName];
 }
 
 
@@ -1188,14 +1195,14 @@ export function resetScheduledElementForTesting(win, elementName) {
  * @return {!Promise<*>}
  */
 export function getElementService(win, id, providedByElement) {
-  return Promise.resolve().then(() => {
-    assert(isElementScheduled(win, providedByElement),
-        'Service %s was requested to be provided through %s, ' +
-        'but %s is not loaded in the current page. To fix this ' +
-        'problem load the JavaScript file for %s in this page.',
-        id, providedByElement, providedByElement, providedByElement);
-    return getServicePromise(win, id);
-  });
+  return getElementServiceIfAvailable(win, id, providedByElement).then(
+      service => {
+        return assert(service,
+            'Service %s was requested to be provided through %s, ' +
+            'but %s is not loaded in the current page. To fix this ' +
+            'problem load the JavaScript file for %s in this page.',
+            id, providedByElement, providedByElement, providedByElement);
+      });
 }
 
 /**
@@ -1212,8 +1219,18 @@ export function getElementServiceIfAvailable(win, id, providedByElement) {
   if (s) {
     return s;
   }
-  if (!isElementScheduled(win, providedByElement)) {
-    return Promise.resolve(null);
-  }
-  return getElementService(win, id, providedByElement);
+  // Microtask is necessary to ensure that window.ampExtendedElements has been
+  // initialized.
+  return Promise.resolve().then(() => {
+    if (isElementScheduled(win, providedByElement)) {
+      return getServicePromise(win, id);
+    }
+    // Wait for HEAD to fully form before denying access to the service.
+    return dom.waitForBodyPromise(win.document).then(() => {
+      if (isElementScheduled(win, providedByElement)) {
+        return getServicePromise(win, id);
+      }
+      return null;
+    });
+  });
 }

src/dom.js

@@ -15,6 +15,60 @@
  */
 
 
+/**
+ * Waits until the child element is constructed. Once the child is found, the
+ * callback is executed.
+ * @param {!Element} parent
+ * @param {function(!Element):boolean} checkFunc
+ * @param {function()} callback
+ */
+export function waitForChild(parent, checkFunc, callback) {
+  if (checkFunc(parent)) {
+    callback();
+    return;
+  }
+  const win = parent.ownerDocument.defaultView;
+  if (win.MutationObserver) {
+    const observer = new win.MutationObserver(() => {
+      if (checkFunc(parent)) {
+        observer.disconnect();
+        callback();
+      }
+    });
+    observer.observe(parent, {childList: true});
+  } else {
+    const interval = win.setInterval(() => {
+      if (checkFunc(parent)) {
+        win.clearInterval(interval);
+        callback();
+      }
+    }, /* milliseconds */ 5);
+  }
+}
+
+
+/**
+ * Waits for document's body to be available.
+ * @param {!Document} doc
+ * @param {function()} callback
+ */
+export function waitForBody(doc, callback) {
+  waitForChild(doc.documentElement, () => !!doc.body, callback);
+}
+
+
+/**
+ * Waits for document's body to be available.
+ * @param {!Document} doc
+ * @return {!Promise}
+ */
+export function waitForBodyPromise(doc) {
+  return new Promise(resolve => {
+    waitForBody(doc, resolve);
+  });
+}
+
+
 /**
  * Whether the element is currently contained in the DOM. Polyfills
  * `document.contains()` method when necessary. Notice that according to spec

src/render-delaying-extensions.js

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import {dev} from './log';
 import {getServicePromise} from './service';
 import {timer} from './timer';
 
@@ -64,6 +65,7 @@ export function waitForExtensions(win) {
  */
 export function includedExtensions(win) {
   const document = win.document;
+  dev.assert(document.body);
 
   return EXTENSIONS.filter(extension => {
     return document.querySelector(`[custom-element="${extension}"]`);

src/runtime.js

@@ -17,18 +17,22 @@
 import {BaseElement} from './base-element';
 import {BaseTemplate, registerExtendedTemplate} from './template';
 import {assert} from './asserts';
+import {dev} from './log';
 import {getMode} from './mode';
+import {getService} from './service';
 import {installStyles} from './styles';
 import {installCoreServices} from './amp-core-service';
 import {isExperimentOn, toggleExperiment} from './experiments';
 import {registerElement} from './custom-element';
 import {registerExtendedElement} from './extended-element';
 import {resourcesFor} from './resources';
-import {timer} from './timer';
 import {viewerFor} from './viewer';
 import {viewportFor} from './viewport';
-import {getService} from './service';
+import {waitForBody} from './dom';
+
 
+/** @const @private {string} */
+const TAG = 'runtime';
 
 /** @type {!Array} */
 const elementsForTesting = [];
@@ -130,7 +134,10 @@ export function adopt(global) {
    * @param {GlobalAmp} fn
    */
   global.AMP.push = function(fn) {
-    fn(global.AMP);
+    // Extensions are only processed once HEAD is complete.
+    waitForBody(global.document, () => {
+      fn(global.AMP);
+    });
   };
 
   /**
@@ -143,20 +150,23 @@ export function adopt(global) {
   global.AMP.setTickFunction = () => {};
 
   // Execute asynchronously scheduled elements.
-  for (let i = 0; i < preregisteredElements.length; i++) {
-    const fn = preregisteredElements[i];
-    try {
-      fn(global.AMP);
-    } catch (e) {
-      // Throw errors outside of loop in its own micro task to
-      // avoid on error stopping other extensions from loading.
-      timer.delay(() => {throw e;}, 1);
+  // Extensions are only processed once HEAD is complete.
+  waitForBody(global.document, () => {
+    for (let i = 0; i < preregisteredElements.length; i++) {
+      const fn = preregisteredElements[i];
+      try {
+        fn(global.AMP);
+      } catch (e) {
+        // Throw errors outside of loop in its own micro task to
+        // avoid on error stopping other extensions from loading.
+        dev.error(TAG, 'Extension failed: ', e);
+      }
     }
-  }
-  // Make sure we empty the array of preregistered extensions.
-  // Technically this is only needed for testing, as everything should
-  // go out of scope here, but just making sure.
-  preregisteredElements.length = 0;
+    // Make sure we empty the array of preregistered extensions.
+    // Technically this is only needed for testing, as everything should
+    // go out of scope here, but just making sure.
+    preregisteredElements.length = 0;
+  });
 }
 
 

src/styles.js

@@ -16,6 +16,8 @@
 
 import {setStyles, setStyle} from './style';
 import {platformFor} from './platform';
+import {waitForBody} from './dom';
+import {waitForExtensions} from './render-delaying-extensions';
 
 
 /**
@@ -84,31 +86,26 @@ export function installStyles(doc, cssText, cb, opt_isRuntimeCss, opt_ext) {
  * If the body is not yet available (because our script was loaded
  * synchronously), polls until it is.
  * @param {!Document} doc The document who's body we should make visible.
- * @param {?Promise=} extensionsPromise A loading promise for special extensions
- *     which must load before the body can be made visible
+ * @param {boolean=} opt_waitForExtensions Whether the body visibility should
+ *     be blocked on key extensions being loaded.
  */
-export function makeBodyVisible(doc, extensionsPromise) {
-  let interval;
+export function makeBodyVisible(doc, opt_waitForExtensions) {
   const set = () => {
-    if (doc.body) {
-      setStyles(doc.body, {
-        opacity: 1,
-        visibility: 'visible',
-        animation: 'none',
-      });
-      clearInterval(interval);
-    }
+    setStyles(doc.body, {
+      opacity: 1,
+      visibility: 'visible',
+      animation: 'none',
+    });
   };
-  const poll = () => {
-    interval = setInterval(set, 4);
-    set();
-  };
-
-  if (extensionsPromise) {
-    extensionsPromise.then(poll, poll);
-  } else {
-    poll();
-  }
+  waitForBody(doc, () => {
+    const extensionsPromise = opt_waitForExtensions ?
+        waitForExtensions(doc.defaultView) : null;
+    if (extensionsPromise) {
+      extensionsPromise.then(set, set);
+    } else {
+      set();
+    }
+  });
 }