From a32ee13fc93e79a2869e7e333e7a5a3d717987fa Mon Sep 17 00:00:00 2001
From: Dreamacro <305009791@qq.com>
Date: Mon, 31 Aug 2020 00:32:18 +0800
Subject: [PATCH] Feature: reuse dns resolver cache when hot reload

---
 common/cache/lrucache.go      | 17 +++++++++++++++++
 common/cache/lrucache_test.go | 18 ++++++++++++++++++
 dns/resolver.go               |  5 +++++
 hub/executor/executor.go      |  8 ++++++++
 4 files changed, 48 insertions(+)

diff --git a/common/cache/lrucache.go b/common/cache/lrucache.go
index 1b1e492..4246d84 100644
--- a/common/cache/lrucache.go
+++ b/common/cache/lrucache.go
@@ -146,6 +146,23 @@ func (c *LruCache) SetWithExpire(key interface{}, value interface{}, expires tim
 	c.maybeDeleteOldest()
 }
 
+// CloneTo clone and overwrite elements to another LruCache
+func (c *LruCache) CloneTo(n *LruCache) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	n.mu.Lock()
+	defer n.mu.Unlock()
+
+	n.lru = list.New()
+	n.cache = make(map[interface{}]*list.Element)
+
+	for e := c.lru.Front(); e != nil; e = e.Next() {
+		elm := e.Value.(*entry)
+		n.cache[elm.key] = n.lru.PushBack(elm)
+	}
+}
+
 func (c *LruCache) get(key interface{}) *entry {
 	c.mu.Lock()
 	defer c.mu.Unlock()
diff --git a/common/cache/lrucache_test.go b/common/cache/lrucache_test.go
index c3c629d..13675bf 100644
--- a/common/cache/lrucache_test.go
+++ b/common/cache/lrucache_test.go
@@ -164,3 +164,21 @@ func TestStale(t *testing.T) {
 	assert.Equal(t, tenSecBefore, expires)
 	assert.Equal(t, true, exist)
 }
+
+func TestCloneTo(t *testing.T) {
+	o := NewLRUCache(WithSize(10))
+	o.Set("1", 1)
+	o.Set("2", 2)
+
+	n := NewLRUCache(WithSize(2))
+	n.Set("3", 3)
+	n.Set("4", 4)
+
+	o.CloneTo(n)
+
+	assert.False(t, n.Exist("3"))
+	assert.True(t, n.Exist("1"))
+
+	n.Set("5", 5)
+	assert.False(t, n.Exist("1"))
+}
diff --git a/dns/resolver.go b/dns/resolver.go
index 04ea41b..21ebb00 100644
--- a/dns/resolver.go
+++ b/dns/resolver.go
@@ -186,6 +186,11 @@ func (r *Resolver) IsFakeIP(ip net.IP) bool {
 	return false
 }
 
+// PatchCache overwrite lruCache to the new resolver
+func (r *Resolver) PatchCache(n *Resolver) {
+	r.lruCache.CloneTo(n.lruCache)
+}
+
 func (r *Resolver) batchExchange(clients []dnsClient, m *D.Msg) (msg *D.Msg, err error) {
 	fast, ctx := picker.WithTimeout(context.Background(), time.Second*5)
 	for _, client := range clients {
diff --git a/hub/executor/executor.go b/hub/executor/executor.go
index 3bb29f5..f0f4c9b 100644
--- a/hub/executor/executor.go
+++ b/hub/executor/executor.go
@@ -120,6 +120,14 @@ func updateDNS(c *config.DNS) {
 		},
 		Default: c.DefaultNameserver,
 	})
+
+	// reuse cache of old resolver
+	if resolver.DefaultResolver != nil {
+		if o, ok := resolver.DefaultResolver.(*dns.Resolver); ok {
+			o.PatchCache(r)
+		}
+	}
+
 	resolver.DefaultResolver = r
 	tunnel.SetResolver(r)
 	if err := dns.ReCreateServer(c.Listen, r); err != nil {