From 0a3b3d0b6da4a9fbca6e8eac346131e57a831632 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Thu, 22 Feb 2018 17:29:21 +0100
Subject: [PATCH] refactor UDPNameServer clean up task

---
 app/dns/nameserver.go | 33 ++++++++++++++-------------------
 1 file changed, 14 insertions(+), 19 deletions(-)

diff --git a/app/dns/nameserver.go b/app/dns/nameserver.go
index a689e785..f3e19e04 100644
--- a/app/dns/nameserver.go
+++ b/app/dns/nameserver.go
@@ -10,14 +10,10 @@ import (
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/dice"
 	"v2ray.com/core/common/net"
+	"v2ray.com/core/common/signal"
 	"v2ray.com/core/transport/internet/udp"
 )
 
-const (
-	CleanupInterval  = time.Second * 120
-	CleanupThreshold = 512
-)
-
 var (
 	multiQuestionDNS = map[net.Address]bool{
 		net.IPAddress([]byte{8, 8, 8, 8}): true,
@@ -42,10 +38,10 @@ type PendingRequest struct {
 
 type UDPNameServer struct {
 	sync.Mutex
-	address     net.Destination
-	requests    map[uint16]*PendingRequest
-	udpServer   *udp.Dispatcher
-	nextCleanup time.Time
+	address   net.Destination
+	requests  map[uint16]*PendingRequest
+	udpServer *udp.Dispatcher
+	cleanup   *signal.PeriodicTask
 }
 
 func NewUDPNameServer(address net.Destination, dispatcher core.Dispatcher) *UDPNameServer {
@@ -54,36 +50,35 @@ func NewUDPNameServer(address net.Destination, dispatcher core.Dispatcher) *UDPN
 		requests:  make(map[uint16]*PendingRequest),
 		udpServer: udp.NewDispatcher(dispatcher),
 	}
+	s.cleanup = &signal.PeriodicTask{
+		Interval: time.Minute,
+		Execute:  s.Cleanup,
+	}
+	s.cleanup.Start()
 	return s
 }
 
-func (s *UDPNameServer) Cleanup() {
-	expiredRequests := make([]uint16, 0, 16)
+func (s *UDPNameServer) Cleanup() error {
 	now := time.Now()
 	s.Lock()
 	for id, r := range s.requests {
 		if r.expire.Before(now) {
-			expiredRequests = append(expiredRequests, id)
 			close(r.response)
+			delete(s.requests, id)
 		}
 	}
-	for _, id := range expiredRequests {
-		delete(s.requests, id)
-	}
 	s.Unlock()
+	return nil
 }
 
 func (s *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
 	var id uint16
 	s.Lock()
-	if len(s.requests) > CleanupThreshold && s.nextCleanup.Before(time.Now()) {
-		s.nextCleanup = time.Now().Add(CleanupInterval)
-		go s.Cleanup()
-	}
 
 	for {
 		id = dice.RollUint16()
 		if _, found := s.requests[id]; found {
+			time.Sleep(time.Millisecond * 500)
 			continue
 		}
 		newError("add pending request id ", id).AtDebug().WriteToLog()