feat: provide configurable warmup depth

antongolub · Dec 7, 2022 · 7c70468 · 7c70468
1 parent a34950e
commit 7c70468
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 24 deletions.
diff --git a/README.md b/README.md
@@ -246,12 +246,13 @@ type TFirewallConfigEntry = {
 type TFirewallConfig = Record<string, TFirewallConfigEntry>
 
 type TConfig = {
-  agent?: TAgentConfig
-  cache?: TCacheConfig | TCacheImpl
-  extend?: string
+  server: TServerConfig
   firewall: TFirewallConfig
+  extend?: string
+  agent?: TAgentConfig
   log?: { level?: TLogeLevel }
-  server: TServerConfig
+  cache?: TCacheConfig | TCacheImpl
+  warmup?: boolean | number
 }
 
 type TValidationContext = {
@@ -328,7 +329,7 @@ export function getMetricsDigest(): Record<string, any>
     "limit": 1000000          // Optional. Max cache size in bytes. Defaults to Infinity
   },
   "workerConcurrency": 2,     // Optional. os.cpus().length by default. NB: workers are used to perform `heavy` tasks like zip
-  "warmup": true,             // Optional. Lets the prefetcher guess the next packages to load. Defaults to true
+  "warmup": true,             // Optional. Lets the prefetcher guess the next packages to load. Defaults to true (infinity). If set to a number, limits the fetching depth.
   "firewall": {
     "/foo": {                 // Context path
       "registry": "https://registry.npmmirror.com",  // Remote registry

diff --git a/src/main/js/config.js b/src/main/js/config.js
@@ -21,7 +21,7 @@ const populate = (config) => {
   // assert.ok(config.firewall, 'cfg: firewall')
 
   const workerConcurrency = getConcurrencyLimit(config.workerConcurrency)
-  const warmup = !!(config.warmup ?? true)
+  const warmup = config.warmup === true || config.warmup === undefined ? Number.POSITIVE_INFINITY : config.warmup | 0
   const agent = config.agent
   const cache = config.cache
     ? {

diff --git a/src/main/js/firewall/middleware.js b/src/main/js/firewall/middleware.js
@@ -8,18 +8,20 @@ import {checkTarball} from './tarball.js'
 import {logger} from '../logger.js'
 import {getConfig} from '../config.js'
 
-const warmupPipeline = (pipeline, opts) => pipeline.forEach(([plugin, _opts]) => {
-  if (getConfig().warmup === false || isNoCache()) return
+const warmupPipeline = (pipeline, opts, warmup = getConfig().warmup) => {
+  if (warmup <= 0 || isNoCache()) return
 
-  try {
-    plugin.warmup?.({...opts, ..._opts })
-  } catch (e) {
-    logger.error(`Error in plugin ${plugin.name} warmup`, e)
-  }
-})
+  pipeline.forEach(([plugin, _opts]) => {
+    try {
+      plugin.warmup?.({...opts, ..._opts })
+    } catch (e) {
+      logger.error(`Error in plugin ${plugin.name} warmup`, e)
+    }
+  })
+}
 
-const warmupDepPackuments = (name, deps, boundContext, rules) => {
-  if (getConfig().warmup === false || isNoCache()) return
+const warmupDepPackuments = (name, deps, boundContext, rules, warmup = getConfig().warmup) => {
+  if (warmup <= 0 || isNoCache()) return
 
   const {registry, authorization, entrypoint, pipeline} = boundContext
   logger.debug(`warmup ${name} deps`, deps)
@@ -30,10 +32,10 @@ const warmupDepPackuments = (name, deps, boundContext, rules) => {
     }
     const org = name.charAt(0) === '@' ? name.slice(0, (name.indexOf('/') + 1 || name.indexOf('%') + 1) - 1) : null
     try {
-      warmupPipeline(pipeline, {name, registry, org})
+      warmupPipeline(pipeline, {name, registry, org}, warmup--)
 
       const {deps: _deps} = await getPackument({ boundContext: {registry, authorization, entrypoint, name, org, pipeline}, rules })
-      warmupDepPackuments(name, _deps, boundContext, rules)
+      warmupDepPackuments(name, _deps, boundContext, rules, warmup--)
     } catch (e) {
       logger.warn('warmup error', e.message, e.stack)
     }

diff --git a/src/main/js/index.d.ts b/src/main/js/index.d.ts
@@ -75,12 +75,13 @@ type TFirewallConfigEntry = {
 type TFirewallConfig = Record<string, TFirewallConfigEntry>
 
 type TConfig = {
-  agent?: TAgentConfig
-  cache?: TCacheConfig | TCacheImpl
-  extend?: string
+  server: TServerConfig
   firewall: TFirewallConfig
+  extend?: string
+  agent?: TAgentConfig
   log?: { level?: TLogeLevel }
-  server: TServerConfig
+  cache?: TCacheConfig | TCacheImpl
+  warmup?: boolean | number
 }
 
 type TValidationContext = {

diff --git a/src/test/js/config.js b/src/test/js/config.js
@@ -127,14 +127,14 @@ test('processes `workerConcurrency` && `warmup` opts', () => {
     firewall: {},
   }), {
     workerConcurrency: 1,
-    warmup: false,
+    warmup: 0,
   })
 
   objectContaining(loadConfig({
     server: {port: 3000},
     firewall: {},
   }), {
     workerConcurrency: os.cpus().length,
-    warmup: true,
+    warmup: Number.POSITIVE_INFINITY,
   })
 })