feat: DNS interceptor storage (#4589)

Co-authored-by: o.drapeza <[email protected]>

Author: SuperOleg39

docs/docs/api/Dispatcher.md

@@ -1043,6 +1043,7 @@ The `dns` interceptor enables you to cache DNS lookups for a given duration, per
   - The function should return a single record from the records array.
   - By default a simplified version of Round Robin is used.
   - The `records` property can be mutated to store the state of the balancing algorithm.
+- `storage: DNSStorage` - Custom storage for resolved DNS records
 
 > The `Dispatcher#options` also gets extended with the options `dns.affinity`, `dns.dualStack`, `dns.lookup` and `dns.pick` which can be used to configure the interceptor at a request-per-request basis.
 
@@ -1057,6 +1058,15 @@ It represents a map of DNS IP addresses records for a single origin.
 - `4.ips` - (`DNSInterceptorRecord[] | null`) The IPv4 addresses.
 - `6.ips` - (`DNSInterceptorRecord[] | null`) The IPv6 addresses.
 
+**DNSStorage**
+It represents a storage object for resolved DNS records.
+- `size` - (`number`) current size of the storage.
+- `get` - (`(origin: string) => DNSInterceptorOriginRecords | null`) method to get the records for a given origin.
+- `set` - (`(origin: string, records: DNSInterceptorOriginRecords | null, options: { ttl: number }) => void`) method to set the records for a given origin.
+- `delete` - (`(origin: string) => void`) method to delete records for a given origin.
+- `full` - (`() => boolean`) method to check if the storage is full, if returns `true`, DNS lookup will be skipped in this interceptor and new records will not be stored.
+```ts
+
 **Example - Basic DNS Interceptor**
 
 ```js
@@ -1073,6 +1083,45 @@ const response = await client.request({
 })
 ```
 
+**Example - DNS Interceptor and LRU cache as a storage**
+
+```js
+const { Client, interceptors } = require("undici");
+const QuickLRU = require("quick-lru");
+const { dns } = interceptors;
+
+const lru = new QuickLRU({ maxSize: 100 });
+
+const lruAdapter = {
+  get size() {
+    return lru.size;
+  },
+  get(origin) {
+    return lru.get(origin);
+  },
+  set(origin, records, { ttl }) {
+    lru.set(origin, records, { maxAge: ttl });
+  },
+  delete(origin) {
+    lru.delete(origin);
+  },
+  full() {
+    // For LRU cache, we can always store new records,
+    // old records will be evicted automatically
+    return false;
+  }
+}
+
+const client = new Agent().compose([
+  dns({ storage: lruAdapter })
+])
+
+const response = await client.request({
+  origin: `http://localhost:3030`,
+  ...requestOpts
+})
+```
+
 ##### `responseError`
 
 The `responseError` interceptor throws an error for responses with status code errors (>= 400).

lib/interceptor/dns.js

@@ -5,14 +5,44 @@ const DecoratorHandler = require('../handler/decorator-handler')
 const { InvalidArgumentError, InformationalError } = require('../core/errors')
 const maxInt = Math.pow(2, 31) - 1
 
+class DNSStorage {
+  #maxItems = 0
+  #records = new Map()
+
+  constructor (opts) {
+    this.#maxItems = opts.maxItems
+  }
+
+  get size () {
+    return this.#records.size
+  }
+
+  get (hostname) {
+    return this.#records.get(hostname) ?? null
+  }
+
+  set (hostname, records) {
+    this.#records.set(hostname, records)
+  }
+
+  delete (hostname) {
+    this.#records.delete(hostname)
+  }
+
+  // Delegate to storage decide can we do more lookups or not
+  full () {
+    return this.size >= this.#maxItems
+  }
+}
+
 class DNSInstance {
   #maxTTL = 0
   #maxItems = 0
-  #records = new Map()
   dualStack = true
   affinity = null
   lookup = null
   pick = null
+  storage = null
 
   constructor (opts) {
     this.#maxTTL = opts.maxTTL
@@ -21,17 +51,14 @@ class DNSInstance {
     this.affinity = opts.affinity
     this.lookup = opts.lookup ?? this.#defaultLookup
     this.pick = opts.pick ?? this.#defaultPick
-  }
-
-  get full () {
-    return this.#records.size === this.#maxItems
+    this.storage = opts.storage ?? new DNSStorage(opts)
   }
 
   runLookup (origin, opts, cb) {
-    const ips = this.#records.get(origin.hostname)
+    const ips = this.storage.get(origin.hostname)
 
     // If full, we just return the origin
-    if (ips == null && this.full) {
+    if (ips == null && this.storage.full()) {
       cb(null, origin)
       return
     }
@@ -55,7 +82,7 @@ class DNSInstance {
         }
 
         this.setRecords(origin, addresses)
-        const records = this.#records.get(origin.hostname)
+        const records = this.storage.get(origin.hostname)
 
         const ip = this.pick(
           origin,
@@ -89,7 +116,7 @@ class DNSInstance {
 
       // If no IPs we lookup - deleting old records
       if (ip == null) {
-        this.#records.delete(origin.hostname)
+        this.storage.delete(origin.hostname)
         this.runLookup(origin, opts, cb)
         return
       }
@@ -193,7 +220,7 @@ class DNSInstance {
   }
 
   pickFamily (origin, ipFamily) {
-    const records = this.#records.get(origin.hostname)?.records
+    const records = this.storage.get(origin.hostname)?.records
     if (!records) {
       return null
     }
@@ -227,11 +254,13 @@ class DNSInstance {
   setRecords (origin, addresses) {
     const timestamp = Date.now()
     const records = { records: { 4: null, 6: null } }
+    let minTTL = this.#maxTTL
     for (const record of addresses) {
       record.timestamp = timestamp
       if (typeof record.ttl === 'number') {
         // The record TTL is expected to be in ms
         record.ttl = Math.min(record.ttl, this.#maxTTL)
+        minTTL = Math.min(minTTL, record.ttl)
       } else {
         record.ttl = this.#maxTTL
       }
@@ -242,11 +271,12 @@ class DNSInstance {
       records.records[record.family] = familyRecords
     }
 
-    this.#records.set(origin.hostname, records)
+    // We provide a default TTL if external storage will be used without TTL per record-level support
+    this.storage.set(origin.hostname, records, { ttl: minTTL })
   }
 
   deleteRecords (origin) {
-    this.#records.delete(origin.hostname)
+    this.storage.delete(origin.hostname)
   }
 
   getHandler (meta, opts) {
@@ -372,6 +402,17 @@ module.exports = interceptorOpts => {
     throw new InvalidArgumentError('Invalid pick. Must be a function')
   }
 
+  if (
+    interceptorOpts?.storage != null &&
+    (typeof interceptorOpts?.storage?.get !== 'function' ||
+      typeof interceptorOpts?.storage?.set !== 'function' ||
+      typeof interceptorOpts?.storage?.full !== 'function' ||
+      typeof interceptorOpts?.storage?.delete !== 'function'
+    )
+  ) {
+    throw new InvalidArgumentError('Invalid storage. Must be a object with methods: { get, set, full, delete }')
+  }
+
   const dualStack = interceptorOpts?.dualStack ?? true
   let affinity
   if (dualStack) {
@@ -386,7 +427,8 @@ module.exports = interceptorOpts => {
     pick: interceptorOpts?.pick ?? null,
     dualStack,
     affinity,
-    maxItems: interceptorOpts?.maxItems ?? Infinity
+    maxItems: interceptorOpts?.maxItems ?? Infinity,
+    storage: interceptorOpts?.storage
   }
 
   const instance = new DNSInstance(opts)

test/interceptors/dns.js

@@ -15,7 +15,7 @@ const { interceptors, Agent } = require('../..')
 const { dns } = interceptors
 
 test('Should validate options', t => {
-  t = tspl(t, { plan: 10 })
+  t = tspl(t, { plan: 11 })
 
   t.throws(() => dns({ dualStack: 'true' }), { code: 'UND_ERR_INVALID_ARG' })
   t.throws(() => dns({ dualStack: 0 }), { code: 'UND_ERR_INVALID_ARG' })
@@ -27,6 +27,7 @@ test('Should validate options', t => {
   t.throws(() => dns({ maxItems: -1 }), { code: 'UND_ERR_INVALID_ARG' })
   t.throws(() => dns({ lookup: {} }), { code: 'UND_ERR_INVALID_ARG' })
   t.throws(() => dns({ pick: [] }), { code: 'UND_ERR_INVALID_ARG' })
+  t.throws(() => dns({ storage: new Map() }), { code: 'UND_ERR_INVALID_ARG' })
 })
 
 test('Should automatically resolve IPs (dual stack)', async t => {
@@ -1727,6 +1728,130 @@ test('Should handle max cached items', async t => {
   t.equal(await response3.body.text(), 'hello world! (x2)')
 })
 
+test('Should support external storage', async t => {
+  t = tspl(t, { plan: 9 })
+
+  let counter = 0
+  const server1 = createServer({ joinDuplicateHeaders: true })
+  const server2 = createServer({ joinDuplicateHeaders: true })
+  const requestOptions = {
+    method: 'GET',
+    path: '/',
+    headers: {
+      'content-type': 'application/json'
+    }
+  }
+
+  server1.on('request', (req, res) => {
+    res.writeHead(200, { 'content-type': 'text/plain' })
+    res.end('hello world!')
+  })
+
+  server1.listen(0)
+
+  server2.on('request', (req, res) => {
+    res.writeHead(200, { 'content-type': 'text/plain' })
+    res.end('hello world! (x2)')
+  })
+  server2.listen(0)
+
+  await Promise.all([once(server1, 'listening'), once(server2, 'listening')])
+
+  const cache = new Map()
+  const storage = {
+    get (origin) {
+      return cache.get(origin)
+    },
+    set (origin, records) {
+      cache.set(origin, records)
+    },
+    delete (origin) {
+      cache.delete(origin)
+    },
+    // simulate internal DNSStorage behaviour with `maxItems: 1` parameter
+    full () {
+      return cache.size === 1
+    }
+  }
+
+  const client = new Agent().compose([
+    dispatch => {
+      return (opts, handler) => {
+        ++counter
+        const url = new URL(opts.origin)
+
+        switch (counter) {
+          case 1:
+            t.equal(isIP(url.hostname), 4)
+            break
+
+          case 2:
+            // [::1] -> ::1
+            t.equal(isIP(url.hostname.slice(1, 4)), 6)
+            break
+
+          case 3:
+            t.equal(url.hostname, 'developer.mozilla.org')
+            // Rewrite origin to avoid reaching internet
+            opts.origin = `http://127.0.0.1:${server2.address().port}`
+            break
+          default:
+            t.fails('should not reach this point')
+        }
+
+        return dispatch(opts, handler)
+      }
+    },
+    dns({
+      storage,
+      lookup: (_origin, _opts, cb) => {
+        cb(null, [
+          {
+            address: '::1',
+            family: 6
+          },
+          {
+            address: '127.0.0.1',
+            family: 4
+          }
+        ])
+      }
+    })
+  ])
+
+  after(async () => {
+    await client.close()
+    server1.close()
+    server2.close()
+
+    await Promise.all([once(server1, 'close'), once(server2, 'close')])
+  })
+
+  const response = await client.request({
+    ...requestOptions,
+    origin: `http://localhost:${server1.address().port}`
+  })
+
+  t.equal(response.statusCode, 200)
+  t.equal(await response.body.text(), 'hello world!')
+
+  const response2 = await client.request({
+    ...requestOptions,
+    origin: `http://localhost:${server1.address().port}`
+  })
+
+  t.equal(response2.statusCode, 200)
+  t.equal(await response2.body.text(), 'hello world!')
+
+  const response3 = await client.request({
+    ...requestOptions,
+    origin: 'https://developer.mozilla.org'
+  })
+
+  t.equal(response3.statusCode, 200)
+  t.equal(await response3.body.text(), 'hello world! (x2)')
+})
+
 test('retry once with dual-stack', async t => {
   t = tspl(t, { plan: 2 })