pre-load source maps to minimize calls at exit

This is a workaround for whatever bug in node or V8 is causing
tapjs/tapjs#934

Best guess is that it has something to do with the FinalizationRegistry,
and Module.findSourceMap is triggering this somehow when it's called
many times in rapid succession at the end of the process.

package.json

@@ -9,6 +9,7 @@
   "types": "./dist/commonjs/index.d.ts",
   "tshy": {
     "main": true,
+    "selfLink": false,
     "exports": {
       ".": "./src/index.ts",
       "./loader": "./src/loader-legacy.mts",

src/find-source-map-safe.ts

@@ -3,19 +3,23 @@ import { pathToURL } from './path-to-url.js'
 
 const sourceMaps = new Map<string, SourceMap>()
 
-export const findSourceMapSafe = (s: string | URL) => {
+export const getSourceMaps = () => sourceMaps
+
+export const findSourceMapSafe = (
+  s: string | URL
+): false | undefined | SourceMap => {
   // Have to look up by URL, because the ?tapmock param will be in
   // the internal key used by node, as those are "different" modules.
   const mod = pathToURL(s)
   const c = sourceMaps.get(mod)
   if (c) return c
 
-  // this can throw in some cases in node 19
+  // this can throw in some cases, eg if the sourcemap file is missing
   try {
     const sm = findSourceMap(mod)
     if (sm) sourceMaps.set(mod, sm)
     return sm
-    /* c8 ignore start */
-  } catch {}
-  /* c8 ignore stop */
+  } catch {
+    return false
+  }
 }

src/hooks.mts

@@ -1,5 +1,6 @@
 // hooks used by loader-legacy.mjs and loader.mjs
 
+import { readFile } from 'node:fs/promises'
 import { parse } from 'node:path'
 import { fileURLToPath } from 'node:url'
 import type { MessagePort } from 'node:worker_threads'
@@ -11,11 +12,14 @@ import {
   reset as processInfoReset,
 } from './get-process-info.js'
 import { saveLineLengths } from './line-lengths.js'
+import { likelyHasSourceMap } from './lookup-sources.js'
 
 let getProcessInfo = _getProcessInfo
-
 let PORT: undefined | MessagePort = undefined
 
+const exclude = getExclude('_TAPJS_PROCESSINFO_EXCLUDE_', false)
+const smMagicComment = /\/[*\/]#\s+sourceMappingURL=[^\s]+/
+
 export const globalPreload = (context: { port?: MessagePort }) => {
   // this will be something like path/to/dist/esm/lib/esm.mjs
   // but we need path/to/dist/commonjs/cjs.js
@@ -29,13 +33,15 @@ if (typeof port !== 'undefined') {
   const require = createRequire(${JSON.stringify(base)})
   const { getProcessInfo } = require('./get-process-info.js')
   const { saveLineLengths } = require('./line-lengths.js')
+  const { likelyHasSourceMap } = require('./lookup-sources.js')
   // must be called eagerly here.
   // this does all the registration as well.
   const processInfo = getProcessInfo()
   port.onmessage = (e) => {
-    const { filename, content } = e.data
+    const { filename, content, url } = e.data
     processInfo.files.push(filename)
     saveLineLengths(filename, content)
+    if (url) likelyHasSourceMap(url)
   }
   port.unref()
 }
@@ -46,9 +52,11 @@ export const initialize = ({ port }: { port: MessagePort }) => {
   PORT = port
 }
 
-const exclude = getExclude('_TAPJS_PROCESSINFO_EXCLUDE_', false)
-
-const record = async (url: string, content?: string) => {
+const record = async (
+  url: string,
+  content?: string,
+  originSource?: string
+) => {
   const filename = url.startsWith('file://') ? fileURLToPath(url) : url
   if (exclude.test(filename)) {
     return
@@ -57,12 +65,29 @@ const record = async (url: string, content?: string) => {
     return
   }
 
+  let maybeSM = false
+  if (
+    originSource !== content ||
+    (content === undefined && url.startsWith('file://'))
+  ) {
+    // try to read the file, fall back to the content we have, or ''
+    // if any source maps anywhere, flag it as possibly having one
+    originSource ??=
+      (await readFile(filename, 'utf8').catch(() => content)) ?? ''
+  }
+  maybeSM = smMagicComment.test(originSource as string)
+
   if (PORT) {
-    PORT.postMessage({ filename, content })
+    PORT.postMessage({
+      filename,
+      content,
+      ...(maybeSM && { url }),
+    })
   } else {
     // call lazily so we don't double-register
     getProcessInfo().files.push(filename)
     saveLineLengths(filename, content)
+    if (maybeSM) likelyHasSourceMap(url)
   }
 }
 
@@ -90,9 +115,11 @@ export const load = async (
     }
   }
 
-  // we actually need the transpiled
+  // get line lengths from final source
+  // if origin source doesn't match, check for possible source map
+  const originSource = context.source
   const ret = await nextLoad(url, context)
-  await record(url, ret.source)
+  await record(url, ret.source, originSource)
   return ret
 }
 

src/import.mts

@@ -6,15 +6,17 @@ import { MessageChannel } from 'node:worker_threads'
 import { getProcessInfo } from './get-process-info.js'
 import { saveLineLengths } from './line-lengths.js'
 import { getImportMetaURL } from './get-import-meta-url.js'
+import {likelyHasSourceMap} from './lookup-sources.js'
 
 const { port1, port2 } = new MessageChannel()
 
 // must be called eagerly here.
 // this does all the registration as well.
 const processInfo = getProcessInfo()
-port1.on('message', ({ filename, content }) => {
+port1.on('message', ({ filename, content, url }) => {
   processInfo.files.push(filename)
   saveLineLengths(filename, content)
+  if (url) likelyHasSourceMap(url)
 })
 
 port1.unref()

src/line-lengths.ts

@@ -23,10 +23,8 @@ export const saveLineLengths = (
   const last = content.trimEnd().split('\n').pop()
   if (cache.has(filename) || !last?.startsWith(sourceMapComment)) return
   const ll = content
-    .replace(/[\n\u2028\u2029]$/, '')
     .split(/\n|\u2028|\u2029/)
     .map(l => l.length)
-  cache.set(filename, ll)
 
   cache.set(filename, ll)
 }

src/lookup-sources.ts

@@ -0,0 +1,52 @@
+// For some reason that is very mysterious as of the time of writing this,
+// node sporadically will in rare cases hang and fail to gracefully exit if
+// a sufficiently large number of findSourceMap calls are made during the
+// process exit event.
+//
+// However, we cannot look up source maps until *after* the module load
+// event is completely finished and the module is about to be executed,
+// because that is when the source map is added to node's cache.
+//
+// To work around this, every time a module is loaded, we attempt to
+// determine whether it likely has the magic sourceMappingURL comment.
+// If so, then we put it in a list, and at each message, attempt to load
+// the sources for all modules in the list. Then, on process exit, if
+// there's anything still pending that likely has a source map, we only
+// have to look up at most one module (ie, if the last module loaded had a
+// source map), which seems to not trigger the hang.
+
+import { findSourceMapSafe } from './find-source-map-safe.js'
+
+// the list of modules that likely have source maps
+const maybeSM = new Set<string>()
+export const sourcesCache = new Map<string, string[]>()
+
+export const loadPendingSourceMaps = () => {
+  for (const url of maybeSM) {
+    const sm = findSourceMapSafe(url)
+    if (sm === false) {
+      // can only happen if node found the SM comment, and tried to load it,
+      // but got an error creating the sourcemap, because it's invalid or
+      // the file is not present. No need to keep trying.
+      maybeSM.delete(url)
+    } else {
+      const sources = sm?.payload?.sources
+      if (sources) {
+        sourcesCache.set(url, sources)
+        maybeSM.delete(url)
+      }
+    }
+  }
+}
+
+export const lookupSources = (url: string) => getSources().get(url)
+
+export const getSources = () => {
+  if (maybeSM.size) loadPendingSourceMaps()
+  return sourcesCache
+}
+
+export const likelyHasSourceMap = (url: string) => {
+  if (!sourcesCache.has(url)) maybeSM.add(url)
+  loadPendingSourceMaps()
+}

src/register-coverage.ts

@@ -10,6 +10,7 @@ import { fileURLToPath } from 'node:url'
 import { findSourceMapSafe } from './find-source-map-safe.js'
 import { getExclude } from './get-exclude.js'
 import { getLineLengths } from './line-lengths.js'
+import { lookupSources } from './lookup-sources.js'
 import { ProcessInfoNodeData } from './process-info-node.js'
 
 export let SESSION: Session | undefined = undefined
@@ -39,16 +40,14 @@ const exclude = p.env._TAPJS_PROCESSINFO_COV_EXCLUDE_
 
 const fileCovered = (
   f: string,
-  s?: SourceMapPayload,
+  sources: string[] = [],
   files: string[] = []
 ) => {
   const testFiles = [f]
-  if (s) {
-    for (const src of s.sources || []) {
-      testFiles.push(
-        resolve(src.startsWith('file://') ? fileURLToPath(src) : src)
-      )
-    }
+  for (const src of sources || []) {
+    testFiles.push(
+      resolve(src.startsWith('file://') ? fileURLToPath(src) : src)
+    )
   }
 
   // never include coverage if the file is fully ignored.
@@ -134,14 +133,22 @@ export const coverageOnProcessEnd = (
       // see if it has a source map
       // need to look up via the url, not the file path, because mocks
       // attach a tapmock search param, which is in node's internal key.
-      const s = findSourceMapSafe(obj.url)
-      if (!fileCovered(f, s?.payload, processInfo.files)) {
+      const sources = lookupSources(obj.url)
+      if (!fileCovered(f, sources, processInfo.files)) {
         return false
       }
+      // Most of the time this will be cached at the time of recording, but
+      // if it's the last module loaded, or transpiled in-place by ts-node,
+      // the sourcemap won't be pre-loaded and will have to be looked up.
+      const s = findSourceMapSafe(obj.url)
       const { payload } = s || { payload: null }
       if (payload) {
         sourceMapCache[obj.url] = Object.assign(Object.create(null), {
-          lineLengths: getLineLengths(f),
+          /* c8 ignore start */
+          // node's SourceMap objects provide this as of 20.5.0
+          //@ts-ignore
+          lineLengths: s?.lineLengths || getLineLengths(f),
+          /* c8 ignore stop */
           data: payload,
         })
       }

src/register-process-end.ts

@@ -23,9 +23,7 @@ export const register = () => {
       // try to find the actual sources of the files we loaded
       // This can't be done up front, because the sourcemap isn't
       // present during the load phase, since it's in the contents.
-      for (const file of processInfo.files) {
-        setSources(file)
-      }
+      setSources(processInfo)
       processInfo.runtime = runtime[0] * 1e3 + runtime[1] / 1e6
       const globalsAdded = Object.keys(global).filter(k => !globals.has(k))
       if (globalsAdded.length) {

src/set-sources.ts

@@ -1,15 +1,19 @@
-import { findSourceMapSafe } from './find-source-map-safe.js'
-import { getProcessInfo } from './get-process-info.js'
+import { fileURLToPath } from 'url'
+import { ProcessInfoNodeData } from './index.js'
+import { getSources } from './lookup-sources.js'
 import { urlToPath } from './url-to-path.js'
 
+let sourcesCache: Map<string, string[]>
+
 // set the processInfo.sources for a given file, but don't clobber
 // if called multiple times, or create duplicate entries.
 // Should only be called *after* the file in question has been loaded.
-export const setSources = (file: string) => {
-  const sm = findSourceMapSafe(file)
-  if (!sm) return
-  const pi = getProcessInfo()
-  const s = pi.sources[file] || []
-  s.push(...sm.payload.sources.map(s => urlToPath(s)))
-  pi.sources[file] = [...new Set(s)]
+export const setSources = (pi: ProcessInfoNodeData) => {
+  sourcesCache ??= getSources()
+  for (const [url, sources] of sourcesCache.entries()) {
+    const file = fileURLToPath(url)
+    const s = pi.sources[file] || []
+    s.push(...sources.map(s => urlToPath(s)))
+    pi.sources[file] = [...new Set(s)]
+  }
 }

test/find-source-map-safe.cts

@@ -1,7 +1,19 @@
 import { findSourceMap } from 'module'
 import t from 'tap'
 import { findSourceMapSafe } from '../dist/commonjs/find-source-map-safe.js'
-t.strictSame(
-  findSourceMapSafe(__filename)?.payload,
-  findSourceMap(__filename)?.payload
-)
+const sm = findSourceMapSafe(__filename)
+t.ok(sm)
+t.strictSame(sm && sm?.payload, findSourceMap(__filename)?.payload)
+// run a second time to exercise cache path
+t.equal(sm, findSourceMapSafe(__filename))
+
+t.test('throwing source map lookup returns false', async t => {
+  await import('./fixtures/missing-sm.min.mjs')
+  t.throws(() => findSourceMap(
+    require.resolve('./fixtures/missing-sm.min.mjs')
+  ))
+  const missing = findSourceMapSafe(
+    require.resolve('./fixtures/missing-sm.min.mjs')
+  )
+  t.equal(missing, false)
+})

test/fixtures/missing-sm.min.d.mts

@@ -0,0 +1 @@
+export const d: 1

test/fixtures/missing-sm.min.mjs

@@ -0,0 +1,2 @@
+"use strict";export const d=1;
+//# sourceMappingURL=missing-sm.min.mjs.map

test/hooks.cts

@@ -13,6 +13,7 @@ const dir = t.testdir({
       console.log('otherload', url)
       return next(url + '?otherload', context, next)
     }
+//# sourceMappingURL=not.valid.map
   `,
   'otherresolve.mjs': `
     export const resolve = async (url, context, next) => {