From cc72e1337e3db8afdca86fc3c32f86fa737d14a1 Mon Sep 17 00:00:00 2001
From: fengmk2 <m@fengmk2.com>
Date: Tue, 27 Sep 2016 01:29:39 +0800
Subject: [PATCH] feat: support lookup with timeout

https://github.com/nodejs/node/issues/7231
---
 .travis.yml               |  4 +--
 README.md                 |  4 ++-
 index.js                  | 29 +++++++++++++++++--
 test/failover-dns.test.js | 60 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 92 insertions(+), 5 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 09e26b4..2408fd7 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,8 +1,8 @@
 sudo: false
 language: node_js
 node_js:
-  - 6
-  - 4
+  - '6'
+  - '4'
 script:
   - npm run ci
 after_script:
diff --git a/README.md b/README.md
index 0e73d42..91914a5 100644
--- a/README.md
+++ b/README.md
@@ -23,6 +23,8 @@ failover-dns
 
 Use local cache dns query result when dns query fail.
 
+- Support dns lookup with `options.timeout`.
+
 ## Installation
 
 ```bash
@@ -37,7 +39,7 @@ const dns = require('failover-dns');
 // must listen `error` event to logging by yourself
 dns.on('error', err => console.error(err));
 
-dns.lookup('cnpmjs.org', { family: 4 }, (err, ip, family) => {
+dns.lookup('cnpmjs.org', { family: 4, timeout: 2000 }, (err, ip, family) => {
   if (err) throw err;
   console.log(ip, family);
 });
diff --git a/index.js b/index.js
index 8891f41..4728bc5 100644
--- a/index.js
+++ b/index.js
@@ -14,13 +14,15 @@ exports.lookup = function lookup(hostname, options, callback) {
     callback = options;
     options = {};
   }
+  options = options || {};
+
+  // don't failover on `options.all = true`
   if (options.all) {
-    // don't failover on options.all
     return dns.lookup(hostname, options, callback);
   }
 
   const cacheKey = `${hostname}_${options.family}_${options.hints}`;
-  dns.lookup(hostname, options, (err, ip, family) => {
+  exports._lookupWithTimeout(hostname, options, (err, ip, family) => {
     if (err) {
       const address = DNS_LOOKUP_CACHE[cacheKey];
       if (address) {
@@ -49,3 +51,26 @@ exports.lookup = function lookup(hostname, options, callback) {
     callback(null, ip, family);
   });
 };
+
+exports._lookupWithTimeout = function lookupWithTimeout(hostname, options, callback) {
+  if (!options.timeout) {
+    return dns.lookup(hostname, options, callback);
+  }
+  let timer = setTimeout(() => {
+    timer = null;
+    const cb = callback;
+    callback = null;
+    const err = new Error(`getaddrinfo TIMEOUT ${hostname}`);
+    err.options = options;
+    err.code = 'TIMEOUT';
+    cb(err);
+  }, options.timeout);
+
+  dns.lookup(hostname, options, (err, ip, family) => {
+    timer && clearTimeout(timer);
+    if (!callback) {
+      return;
+    }
+    callback(err, ip, family);
+  });
+};
diff --git a/test/failover-dns.test.js b/test/failover-dns.test.js
index dca31c4..c835e63 100644
--- a/test/failover-dns.test.js
+++ b/test/failover-dns.test.js
@@ -18,6 +18,26 @@ describe('test/failover-dns.test.js', () => {
     });
   });
 
+  it('should lookup with options=null from dns server', done => {
+    dns.lookup('a.alipayobjects.com', null, (err, ip, family) => {
+      console.log(err, ip, family);
+      assert(err == null);
+      assert((/^\d+\.\d+\.\d+\.\d+$/).test(ip));
+      assert(family === 4);
+      done();
+    });
+  });
+
+  it('should lookup with options.timeout', done => {
+    dns.lookup('as.alipayobjects.com', { timeout: 5000 }, (err, ip, family) => {
+      console.log(err, ip, family);
+      assert(err == null);
+      assert((/^\d+\.\d+\.\d+\.\d+$/).test(ip));
+      assert(family === 4);
+      done();
+    });
+  });
+
   it('should lookup with options.family from dns server', done => {
     dns.lookup('cnpmjs.org', { family: 4 }, (err, ip, family) => {
       console.log(err, ip, family);
@@ -66,6 +86,46 @@ describe('test/failover-dns.test.js', () => {
     });
   });
 
+  it('should mock timeout then return address from local cache', done => {
+    done = pedding(2, done);
+    mm(require('dns'), 'lookup', (hostname, options, callback) => {
+      setTimeout(() => {
+        callback(null, '127.0.0.1', 4);
+      }, 600);
+    });
+
+    dns.once('error', err => {
+      console.log(err);
+      assert(err.message === 'getaddrinfo TIMEOUT a.alipayobjects.com');
+      done();
+    });
+    dns.lookup('a.alipayobjects.com', { family: 4, timeout: 500 }, (err, ip, family) => {
+      console.log('lookup result', err, ip, family);
+      assert(err == null);
+      assert((/^\d+\.\d+\.\d+\.\d+$/).test(ip));
+      assert(family === 4);
+      // wait for dns.lookup() done and ignore callback
+      setTimeout(done, 200);
+    });
+  });
+
+  it('should mock timeout and local cache missing', done => {
+    mm(require('dns'), 'lookup', (hostname, options, callback) => {
+      setTimeout(() => {
+        callback(null, '127.0.0.1', 4);
+      }, 600);
+    });
+
+    dns.lookup('foo.cnpmjs.org', { family: 4, timeout: 500 }, (err, ip, family) => {
+      console.log('lookup result', err, ip, family);
+      assert(err);
+      assert(err.message === 'getaddrinfo TIMEOUT foo.cnpmjs.org');
+      assert(err.code === 'TIMEOUT');
+      // wait for dns.lookup() done and ignore callback
+      setTimeout(done, 200);
+    });
+  });
+
   it('should lookup with options.all from dns server', done => {
     dns.lookup('a.alipayobjects.com', { all: true }, (err, ips) => {
       console.log(err, ips);