Add TLS resolver connection reusing and pooling

Also, fix caching issues and add more tests

resolver/clients.go

@@ -1,6 +1,7 @@
 package resolver
 
 import (
+	"context"
 	"crypto/tls"
 	"net"
 	"sync"
@@ -9,6 +10,12 @@ import (
 	"github.com/miekg/dns"
 )
 
+const (
+	defaultClientTTL         = 5 * time.Minute
+	defaultRequestTimeout    = 5 * time.Second
+	connectionEOLGracePeriod = 10 * time.Second
+)
+
 var (
 	localAddrFactory func(network string) net.Addr
 )
@@ -27,21 +34,72 @@ func getLocalAddr(network string) net.Addr {
 	return nil
 }
 
-type clientManager struct {
-	dnsClient *dns.Client
-	factory   func() *dns.Client
+type dnsClientManager struct {
+	lock sync.Mutex
 
-	lock         sync.Mutex
-	refreshAfter time.Time
-	ttl          time.Duration // force refresh of connection to reduce traceability
+	// set by creator
+	serverAddress string
+	ttl           time.Duration // force refresh of connection to reduce traceability
+	factory       func() *dns.Client
+
+	// internal
+	pool []*dnsClient
 }
 
-func newDNSClientManager(_ *Resolver) *clientManager {
-	return &clientManager{
-		ttl: 0, // new client for every request, as we need to randomize the port
+type dnsClient struct {
+	mgr *dnsClientManager
+
+	inUse     bool
+	useUntil  time.Time
+	dead      bool
+	inPool    bool
+	poolIndex int
+
+	client *dns.Client
+	conn   *dns.Conn
+}
+
+// conn returns the *dns.Conn and if it's new. This function may only be called between clientManager.getDNSClient() and dnsClient.done().
+func (dc *dnsClient) getConn() (c *dns.Conn, new bool, err error) {
+	if dc.conn == nil {
+		dc.conn, err = dc.client.Dial(dc.mgr.serverAddress)
+		if err != nil {
+			return nil, false, err
+		}
+		return dc.conn, true, nil
+	}
+	return dc.conn, false, nil
+}
+
+func (dc *dnsClient) done() {
+	dc.mgr.lock.Lock()
+	defer dc.mgr.lock.Unlock()
+
+	dc.inUse = false
+}
+
+func (dc *dnsClient) destroy() {
+	dc.mgr.lock.Lock()
+	dc.inUse = true // block from being used
+	dc.dead = true  // abort cleaning
+	if dc.inPool {
+		dc.inPool = false
+		dc.mgr.pool[dc.poolIndex] = nil
+	}
+	dc.mgr.lock.Unlock()
+
+	if dc.conn != nil {
+		_ = dc.conn.Close()
+	}
+}
+
+func newDNSClientManager(resolver *Resolver) *dnsClientManager {
+	return &dnsClientManager{
+		serverAddress: resolver.ServerAddress,
+		ttl:           0, // new client for every request, as we need to randomize the port
 		factory: func() *dns.Client {
 			return &dns.Client{
-				Timeout: 5 * time.Second,
+				Timeout: defaultRequestTimeout,
 				Dialer: &net.Dialer{
 					LocalAddr: getLocalAddr("udp"),
 				},
@@ -50,25 +108,27 @@ func newDNSClientManager(_ *Resolver) *clientManager {
 	}
 }
 
-func newTCPClientManager(_ *Resolver) *clientManager {
-	return &clientManager{
-		ttl: 0, // TODO: build a custom client that can reuse connections to some degree (performance / privacy tradeoff)
+func newTCPClientManager(resolver *Resolver) *dnsClientManager {
+	return &dnsClientManager{
+		serverAddress: resolver.ServerAddress,
+		ttl:           defaultClientTTL,
 		factory: func() *dns.Client {
 			return &dns.Client{
 				Net:     "tcp",
-				Timeout: 5 * time.Second,
+				Timeout: defaultRequestTimeout,
 				Dialer: &net.Dialer{
 					LocalAddr: getLocalAddr("tcp"),
-					KeepAlive: 15 * time.Second,
+					KeepAlive: defaultClientTTL,
 				},
 			}
 		},
 	}
 }
 
-func newTLSClientManager(resolver *Resolver) *clientManager {
-	return &clientManager{
-		ttl: 0, // TODO: build a custom client that can reuse connections to some degree (performance / privacy tradeoff)
+func newTLSClientManager(resolver *Resolver) *dnsClientManager {
+	return &dnsClientManager{
+		serverAddress: resolver.ServerAddress,
+		ttl:           defaultClientTTL,
 		factory: func() *dns.Client {
 			return &dns.Client{
 				Net: "tcp-tls",
@@ -77,24 +137,90 @@ func newTLSClientManager(resolver *Resolver) *clientManager {
 					ServerName: resolver.VerifyDomain,
 					// TODO: use portbase rng
 				},
-				Timeout: 5 * time.Second,
+				Timeout: defaultRequestTimeout,
 				Dialer: &net.Dialer{
 					LocalAddr: getLocalAddr("tcp"),
-					KeepAlive: 15 * time.Second,
+					KeepAlive: defaultClientTTL,
 				},
 			}
 		},
 	}
 }
 
-func (cm *clientManager) getDNSClient() *dns.Client {
+func (cm *dnsClientManager) getDNSClient() *dnsClient {
 	cm.lock.Lock()
 	defer cm.lock.Unlock()
 
-	if cm.dnsClient == nil || cm.ttl == 0 || time.Now().After(cm.refreshAfter) {
-		cm.dnsClient = cm.factory()
-		cm.refreshAfter = time.Now().Add(cm.ttl)
+	// return new immediately if a new client should be used for every request
+	if cm.ttl == 0 {
+		return &dnsClient{
+			mgr:    cm,
+			client: cm.factory(),
+		}
 	}
 
-	return cm.dnsClient
+	// get first unused from pool
+	now := time.Now().UTC()
+	for _, dc := range cm.pool {
+		if dc != nil && !dc.inUse && now.Before(dc.useUntil) {
+			dc.inUse = true
+			return dc
+		}
+	}
+
+	// no available in pool, create new
+	newClient := &dnsClient{
+		mgr:      cm,
+		inUse:    true,
+		useUntil: now.Add(cm.ttl),
+		inPool:   true,
+		client:   cm.factory(),
+	}
+	newClient.startCleaner()
+
+	// find free spot in pool
+	for poolIndex, dc := range cm.pool {
+		if dc == nil {
+			cm.pool[poolIndex] = newClient
+			newClient.poolIndex = poolIndex
+			return newClient
+		}
+	}
+
+	// append to pool
+	cm.pool = append(cm.pool, newClient)
+	newClient.poolIndex = len(cm.pool) - 1
+	// TODO: shrink pool again?
+
+	return newClient
+}
+
+// startCleaner waits for EOL of the client and then removes it from the pool.
+func (dc *dnsClient) startCleaner() {
+	// While a single worker to clean all connections may be slightly more performant, this approach focuses on least as possible locking and is simpler, thus less error prone.
+	module.StartWorker("dns client cleanup", func(ctx context.Context) error {
+		select {
+		case <-time.After(dc.mgr.ttl + time.Second):
+			dc.mgr.lock.Lock()
+			cleanNow := dc.dead || !dc.inUse
+			dc.mgr.lock.Unlock()
+
+			if cleanNow {
+				dc.destroy()
+				return nil
+			}
+		case <-ctx.Done():
+			// give a short time before kill for graceful request completion
+			time.Sleep(100 * time.Millisecond)
+		}
+
+		// wait for grace period to end, then kill
+		select {
+		case <-time.After(connectionEOLGracePeriod):
+		case <-ctx.Done():
+		}
+
+		dc.destroy()
+		return nil
+	})
 }