Fix domain sharding for ESM output

packages/core/integration-tests/test/domain-sharding.js

@@ -2,7 +2,8 @@
 import assert from 'assert';
 import path from 'path';
 
-import {fsFixture, overlayFS, bundle} from '@atlaspack/test-utils';
+import {fsFixture, overlayFS, bundle, run} from '@atlaspack/test-utils';
+import {domainShardingKey} from '@atlaspack/domain-sharding';
 
 const maxShards = 8;
 
@@ -120,5 +121,70 @@ describe('domain-sharding', () => {
         'Expected generated code for shardUrl was not found',
       );
     });
+
+    it('for ESM loaded bundle manifest', async () => {
+      const maxShards = 8;
+      await fsFixture(overlayFS)`
+        package.json:
+          {
+            "name": "bundle-sharding-test",
+            "@atlaspack/runtime-js": {
+              "domainSharding": {
+                "maxShards": ${maxShards}
+              }
+            }
+          }
+
+        src/index.js:
+          async function fn() {
+            const a = await import('./a.js');
+            const b = await import('./b.js');
+            console.log('a', a, b);
+          }
+          fn();
+
+        src/a.js:
+          export const a = async () => {
+            const b = await import('./b');
+            return b + 'A';
+          }
+        src/b.js:
+          export const b = 'B';
+
+        yarn.lock:
+      `;
+
+      const bundleGraph = await bundle('src/index.js', {
+        inputFS: overlayFS,
+        mode: 'production',
+        defaultTargetOptions: {
+          shouldOptimize: false,
+          outputFormat: 'esmodule',
+        },
+      });
+
+      const mainBundle = bundleGraph
+        .getBundles()
+        .find((b) => b.name === 'index.js');
+
+      if (!mainBundle) {
+        return assert(mainBundle);
+      }
+
+      // $FlowFixMe - Flow doesn't seem to know about doesNotReject
+      await assert.doesNotReject(
+        run(bundleGraph, {[domainShardingKey]: true}),
+        'Expected bundle to be able to execute',
+      );
+
+      const code = await overlayFS.readFile(mainBundle.filePath, 'utf-8');
+
+      const expectedCode = `var$load = (parcelRequire("cfSxn"))(${maxShards});`
+        .replaceAll('$', '\\$')
+        .replaceAll('(', '\\(')
+        .replaceAll(')', '\\)');
+
+      assert.match(code, new RegExp(expectedCode));
+    });
   });
 });

packages/core/test-utils/src/utils.js

@@ -1048,7 +1048,11 @@ export async function runESM(
 
   function load(inputSpecifier, referrer, code = null) {
     // ESM can request bundles with an absolute URL. Normalize this to the baseDir.
-    let specifier = inputSpecifier.replace('http://localhost', baseDir);
+    // Any digits after the - can be ignored, for domain sharding tests
+    let specifier = inputSpecifier.replace(
+      /http:\/\/localhost(-\d+)?/,
+      baseDir,
+    );
 
     if (path.isAbsolute(specifier) || specifier.startsWith('.')) {
       let extname = path.extname(specifier);

packages/runtimes/js/src/JSRuntime.js

@@ -625,7 +625,11 @@ function getLoaderRuntime({
   let code = [];
 
   if (needsEsmLoadPrelude) {
-    code.push(`let load = require('./helpers/browser/esm-js-loader');`);
+    let preludeLoad = shardingConfig
+      ? `let load = require('./helpers/browser/esm-js-loader-shards')(${shardingConfig.maxShards});`
+      : `let load = require('./helpers/browser/esm-js-loader');`;
+
+    code.push(preludeLoad);
   }
 
   code.push(`module.exports = ${loaderCode};`);

packages/runtimes/js/src/helpers/browser/esm-js-loader-shards.js

@@ -0,0 +1,11 @@
+let load = (maxShards) => (id) => {
+  // eslint-disable-next-line no-undef
+  return __parcel__import__(
+    require('@atlaspack/domain-sharding').shardUrl(
+      require('../bundle-manifest').resolve(id),
+      maxShards,
+    ),
+  );
+};
+
+module.exports = load;