PERF: Move highlightjs to a background worker, and add result cache (#10191)

Syntax highlighting is a CPU-intensive process which we run a lot while rendering posts and while using the composer preview. Moving it to a background worker releases the main thread to the browser, which makes the UX much smoother.

app/assets/javascripts/admin/components/highlighted-code.js

@@ -1,11 +1,37 @@
 import Component from "@ember/component";
-import { on, observes } from "discourse-common/utils/decorators";
-import highlightSyntax from "discourse/lib/highlight-syntax";
+import { highlightText } from "discourse/lib/highlight-syntax";
+import { escapeExpression } from "discourse/lib/utilities";
+import discourseComputed from "discourse-common/utils/decorators";
+import { htmlSafe } from "@ember/template";
 
 export default Component.extend({
-  @on("didInsertElement")
-  @observes("code")
-  _refresh: function() {
-    highlightSyntax($(this.element));
+  didReceiveAttrs() {
+    this._super(...arguments);
+    if (this.code === this.previousCode) return;
+
+    this.set("previousCode", this.code);
+    this.set("highlightResult", null);
+    const toHighlight = this.code;
+    highlightText(escapeExpression(toHighlight), this.lang).then(
+      ({ result }) => {
+        if (toHighlight !== this.code) return; // Code has changed since highlight was requested
+        this.set("highlightResult", result);
+      }
+    );
+  },
+
+  @discourseComputed("code", "highlightResult")
+  displayCode(code, highlightResult) {
+    if (highlightResult) return htmlSafe(highlightResult);
+    return code;
+  },
+
+  @discourseComputed("highlightResult", "lang")
+  codeClasses(highlightResult, lang) {
+    const classes = [];
+    if (lang) classes.push(lang);
+    if (highlightResult) classes.push("hljs");
+
+    return classes.join(" ");
   }
 });

app/assets/javascripts/admin/models/theme.js

@@ -7,8 +7,8 @@ import discourseComputed from "discourse-common/utils/decorators";
 import { popupAjaxError } from "discourse/lib/ajax-error";
 import { ajax } from "discourse/lib/ajax";
 import { escapeExpression } from "discourse/lib/utilities";
-import highlightSyntax from "discourse/lib/highlight-syntax";
 import { url } from "discourse/lib/computed";
+import highlightSyntax from "discourse/lib/highlight-syntax";
 
 const THEME_UPLOAD_VAR = 2;
 const FIELDS_IDS = [0, 1, 5];
@@ -321,7 +321,7 @@ const Theme = RestModel.extend({
             }
           }
         );
-        highlightSyntax();
+        highlightSyntax(document.querySelector(".bootbox.modal"));
       } else {
         return this.save({ remote_update: true }).then(() =>
           this.set("changed", false)

app/assets/javascripts/admin/templates/components/highlighted-code.hbs

@@ -1 +1 @@
-<pre><code class={{lang}}>{{code}}</code></pre>
+<pre><code class={{codeClasses}}>{{displayCode}}</code></pre>

app/assets/javascripts/discourse/app/initializers/post-decorations.js

@@ -1,15 +1,15 @@
-import highlightSyntax from "discourse/lib/highlight-syntax";
 import lightbox from "discourse/lib/lightbox";
 import { setupLazyLoading } from "discourse/lib/lazy-load-images";
 import { setTextDirections } from "discourse/lib/text-direction";
 import { withPluginApi } from "discourse/lib/plugin-api";
+import highlightSyntax from "discourse/lib/highlight-syntax";
 
 export default {
   name: "post-decorations",
   initialize(container) {
     withPluginApi("0.1", api => {
       const siteSettings = container.lookup("site-settings:main");
-      api.decorateCooked(highlightSyntax, {
+      api.decorateCookedElement(highlightSyntax, {
         id: "discourse-syntax-highlighting"
       });
       api.decorateCookedElement(lightbox, { id: "discourse-lightbox" });

app/assets/javascripts/discourse/app/lib/highlight-syntax.js

@@ -1,40 +1,181 @@
-/*global hljs:true */
-let _moreLanguages = [];
-
+import { Promise } from "rsvp";
+import { getURLWithCDN } from "discourse-common/lib/get-url";
+import { next, schedule } from "@ember/runloop";
 import loadScript from "discourse/lib/load-script";
+import { isTesting } from "discourse-common/config/environment";
 
-export default function highlightSyntax($elem) {
-  const selector = Discourse.SiteSettings.autohighlight_all_code
-      ? "pre code"
-      : "pre code[class]",
-    path = Discourse.HighlightJSPath;
+let highlightJsUrl;
+let highlightJsWorkerUrl;
 
-  if (!path) {
-    return;
-  }
+const _moreLanguages = [];
+let _worker = null;
+let _workerPromise = null;
+const _pendingResolution = {};
+let _counter = 0;
+let _cachedResultsMap = new Map();
 
-  $(selector, $elem).each(function(i, e) {
-    // Large code blocks can cause crashes or slowdowns
-    if (e.innerHTML.length > 30000) {
-      return;
-    }
+const CACHE_SIZE = 100;
 
-    $(e).removeClass("lang-auto");
-    loadScript(path).then(() => {
-      customHighlightJSLanguages();
-      hljs.highlightBlock(e);
-    });
-  });
+export function setupHighlightJs(args) {
+  highlightJsUrl = args.highlightJsUrl;
+  highlightJsWorkerUrl = args.highlightJsWorkerUrl;
 }
 
 export function registerHighlightJSLanguage(name, fn) {
   _moreLanguages.push({ name: name, fn: fn });
 }
 
-function customHighlightJSLanguages() {
-  _moreLanguages.forEach(l => {
-    if (hljs.getLanguage(l.name) === undefined) {
-      hljs.registerLanguage(l.name, l.fn);
+export default function highlightSyntax(elem, { autoHighlight = false } = {}) {
+  const selector = autoHighlight ? "pre code" : "pre code[class]";
+
+  elem.querySelectorAll(selector).forEach(e => highlightElement(e));
+}
+
+function highlightElement(e) {
+  e.classList.remove("lang-auto");
+  let lang = null;
+  e.classList.forEach(c => {
+    if (c.startsWith("lang-")) {
+      lang = c.slice("lang-".length);
+    }
+  });
+
+  const requestString = e.textContent;
+  highlightText(e.textContent, lang).then(({ result, fromCache }) => {
+    const doRender = () => {
+      // Ensure the code hasn't changed since highlighting was triggered:
+      if (requestString !== e.textContent) return;
+
+      e.innerHTML = result;
+      e.classList.add("hljs");
+    };
+
+    if (fromCache) {
+      // This happened synchronously, we can safely add rendering
+      // to the end of the current Runloop
+      schedule("afterRender", null, doRender);
+    } else {
+      // This happened async, we are probably not in a runloop
+      // If we call `schedule`, a new runloop will be triggered immediately
+      // So schedule rendering to happen in the next runloop
+      next(() => schedule("afterRender", null, doRender));
     }
   });
 }
+
+export function highlightText(text, language) {
+  // Large code blocks can cause crashes or slowdowns
+  if (text.length > 30000) {
+    return Promise.resolve({ result: text, fromCache: true });
+  }
+
+  return getWorker().then(w => {
+    let result;
+    if ((result = _cachedResultsMap.get(cacheKey(text, language)))) {
+      return Promise.resolve({ result, fromCache: true });
+    }
+
+    let resolve;
+    const promise = new Promise(f => (resolve = f));
+
+    w.postMessage({
+      type: "highlight",
+      id: _counter,
+      text,
+      language
+    });
+
+    _pendingResolution[_counter] = {
+      promise,
+      resolve,
+      text,
+      language
+    };
+
+    _counter++;
+
+    return promise;
+  });
+}
+
+function getWorker() {
+  if (_worker) return Promise.resolve(_worker);
+  if (_workerPromise) return _workerPromise;
+
+  const w = new Worker(highlightJsWorkerUrl);
+  w.onmessage = onWorkerMessage;
+  w.postMessage({
+    type: "loadHighlightJs",
+    path: fullHighlightJsUrl()
+  });
+
+  _workerPromise = setupCustomLanguages(w).then(() => (_worker = w));
+  return _workerPromise;
+}
+
+function setupCustomLanguages(worker) {
+  if (_moreLanguages.length === 0) return Promise.resolve();
+  // To build custom language definitions we need to have hljs loaded
+  // Plugins/themes can't run code in a worker, so we have to load hljs in the main thread
+  // But the actual highlighting will still be done in the worker
+
+  return loadScript(highlightJsUrl).then(() => {
+    _moreLanguages.forEach(({ name, fn }) => {
+      const definition = fn(window.hljs);
+      worker.postMessage({
+        type: "registerLanguage",
+        definition,
+        name
+      });
+    });
+  });
+}
+
+function onWorkerMessage(message) {
+  const id = message.data.id;
+  const request = _pendingResolution[id];
+  delete _pendingResolution[id];
+  request.resolve({ result: message.data.result, fromCache: false });
+
+  cacheResult({
+    text: request.text,
+    language: request.language,
+    result: message.data.result
+  });
+}
+
+function cacheResult({ text, language, result }) {
+  _cachedResultsMap.set(cacheKey(text, language), result);
+  while (_cachedResultsMap.size > CACHE_SIZE) {
+    _cachedResultsMap.delete(_cachedResultsMap.entries().next().value[0]);
+  }
+}
+
+function cacheKey(text, lang) {
+  return `${lang}:${text}`;
+}
+
+function fullHighlightJsUrl() {
+  let hljsUrl = getURLWithCDN(highlightJsUrl);
+
+  // Need to use full URL including protocol/domain
+  // for use in a worker
+  if (hljsUrl.startsWith("/")) {
+    hljsUrl = window.location.protocol + "//" + window.location.host + hljsUrl;
+  }
+
+  return hljsUrl;
+}
+
+// To be used in qunit tests. Running highlight in a worker means that the
+// normal system which waits for ember rendering in tests doesn't work.
+// This promise will resolve once all pending highlights are done
+export function waitForHighlighting() {
+  if (!isTesting()) {
+    throw "This function should only be called in a test environment";
+  }
+  const promises = Object.values(_pendingResolution).map(r => r.promise);
+  return new Promise(resolve => {
+    Promise.all(promises).then(() => next(resolve));
+  });
+}

app/assets/javascripts/discourse/app/pre-initializers/discourse-bootstrap.js

@@ -9,6 +9,7 @@ import {
 } from "discourse-common/config/environment";
 import { setupURL, setupS3CDN } from "discourse-common/lib/get-url";
 import deprecated from "discourse-common/lib/deprecated";
+import { setupHighlightJs } from "discourse/lib/highlight-syntax";
 
 export default {
   name: "discourse-bootstrap",
@@ -81,7 +82,11 @@ export default {
       Session.currentProp("safe_mode", setupData.safeMode);
     }
 
-    app.HighlightJSPath = setupData.highlightJsPath;
+    setupHighlightJs({
+      highlightJsUrl: setupData.highlightJsUrl,
+      highlightJsWorkerUrl: setupData.highlightJsWorkerUrl
+    });
+
     app.SvgSpritePath = setupData.svgSpritePath;
 
     if (app.Environment === "development") {

app/assets/javascripts/highlightjs-worker.js

@@ -0,0 +1,47 @@
+// Standalone worker for highlightjs syntax generation
+
+// The highlightjs path changes based on site settings,
+// so we wait for Discourse to pass the path into the worker
+const loadHighlightJs = path => {
+  self.importScripts(path);
+};
+
+const highlight = ({ id, text, language }) => {
+  if (!self.hljs) {
+    throw "HighlightJS is not loaded";
+  }
+
+  const result = language
+    ? self.hljs.highlight(language, text, true).value
+    : self.hljs.highlightAuto(text).value;
+
+  postMessage({
+    type: "highlightResult",
+    id: id,
+    result: result
+  });
+};
+
+const registerLanguage = ({ name, definition }) => {
+  if (!self.hljs) {
+    throw "HighlightJS is not loaded";
+  }
+  self.hljs.registerLanguage(name, () => {
+    return definition;
+  });
+};
+
+onmessage = event => {
+  const data = event.data;
+  const messageType = data.type;
+
+  if (messageType === "loadHighlightJs") {
+    loadHighlightJs(data.path);
+  } else if (messageType === "registerLanguage") {
+    registerLanguage(data);
+  } else if (messageType === "highlight") {
+    highlight(data);
+  } else {
+    throw `Unknown message type: ${messageType}`;
+  }
+};

app/helpers/application_helper.rb

@@ -469,7 +469,8 @@ module ApplicationHelper
       default_locale: SiteSetting.default_locale,
       asset_version: Discourse.assets_digest,
       disable_custom_css: loading_admin?,
-      highlight_js_path: HighlightJs.path,
+      highlight_js_url: HighlightJs.path,
+      highlight_js_worker_url: script_asset_path('highlightjs-worker'),
       svg_sprite_path: SvgSprite.path(theme_ids),
       enable_js_error_reporting: GlobalSetting.enable_js_error_reporting,
     }

config/application.rb

@@ -171,6 +171,7 @@ module Discourse
       confirm-new-email/bootstrap.js
       onpopstate-handler.js
       embed-application.js
+      highlightjs-worker.js
     }
 
     # Precompile all available locales

test/javascripts/components/highlighted-code-test.js

@@ -1,4 +1,8 @@
 import componentTest from "helpers/component-test";
+import {
+  waitForHighlighting,
+  setupHighlightJs
+} from "discourse/lib/highlight-syntax";
 
 const LONG_CODE_BLOCK = "puts a\n".repeat(15000);
 
@@ -8,12 +12,15 @@ componentTest("highlighting code", {
   template: "{{highlighted-code lang='ruby' code=code}}",
 
   beforeEach() {
-    Discourse.HighlightJSPath =
-      "assets/highlightjs/highlight-test-bundle.min.js";
-    this.set("code", "def test; end");
+    setupHighlightJs({
+      highlightJsUrl: "/assets/highlightjs/highlight-test-bundle.min.js",
+      highlightJsWorkerUrl: "/assets/highlightjs-worker.js"
+    });
   },
 
   async test(assert) {
+    this.set("code", "def test; end");
+    await waitForHighlighting();
     assert.equal(
       find("code.ruby.hljs .hljs-function .hljs-keyword")
         .text()
@@ -23,16 +30,19 @@ componentTest("highlighting code", {
   }
 });
 
-componentTest("large code blocks are not highlighted", {
+componentTest("highlighting code limit", {
   template: "{{highlighted-code lang='ruby' code=code}}",
 
   beforeEach() {
-    Discourse.HighlightJSPath =
-      "assets/highlightjs/highlight-test-bundle.min.js";
-    this.set("code", LONG_CODE_BLOCK);
+    setupHighlightJs({
+      highlightJsUrl: "/assets/highlightjs/highlight-test-bundle.min.js",
+      highlightJsWorkerUrl: "/assets/highlightjs-worker.js"
+    });
   },
 
   async test(assert) {
+    this.set("code", LONG_CODE_BLOCK);
+    await waitForHighlighting();
     assert.equal(
       find("code")
         .text()