/
dnspool.go
112 lines (96 loc) · 3.01 KB
/
dnspool.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
/*
Package dnspool creates a goroutine pool for DNS query to limit the number of
OS threads that will be spawned by the go runtime.
Default number of goroutines for DNS lookup is set to 8 now. You can increase
goroutine number by calling SetGoroutineNumber. Note the maximum
number of goroutines is 32. (Please tell me if this number is not enough.)
This package will no longer be needed when the go runtime provides a way to
limit OS threads creation.
*/
package dnspool
import (
"fmt"
"net"
)
type lookupRequest struct {
host string
addrs []string // lookup result
err error
done chan byte // channel notify that lookup is done
}
// Resolver provides LookupHost method which can be used non-concurrently.
type Resolver lookupRequest
// Maximum concurrent lookup request
const maxLookupReq = 32
var requestChan = make(chan *lookupRequest, maxLookupReq)
// Default goroutine pool size. Making this larger than maxLookupReq is
// useless as the channel can hold only that many lookup request.
var nGoroutine = 8
func init() {
for i := 0; i < nGoroutine; i++ {
go lookup()
}
}
// SetGoroutineNumber sets the number of goroutines used to do DNS query.
// If n <= current goroutine number or is larger than maximum concurrent DNS
// request, this function will do nothing.
func SetGoroutineNumber(n int) {
if n <= nGoroutine {
fmt.Printf("SetGoroutineNumber: %d <= current goroutine number %d, do nothing\n", n, nGoroutine)
return
}
if n > maxLookupReq {
fmt.Printf("SetGoroutineNumber: %d > maximum goroutines number %d, do nothing\n", n, maxLookupReq)
return
}
for ; nGoroutine < n; nGoroutine++ {
go lookup()
}
}
func lookup() {
for {
req := <-requestChan
req.addrs, req.err = net.LookupHost(req.host)
req.done <- 1
}
}
func NewResolver() *Resolver {
return &Resolver{done: make(chan byte)}
}
// LookupHost has the same usage as net.LookupHost. Not thread safe.
func (r *Resolver) LookupHost(host string) (addrs []string, err error) {
r.host = host
requestChan <- (*lookupRequest)(r)
<-r.done
return r.addrs, r.err
}
// LookupHost has the same usage as net.LookupHost. If you are going to call
// this repeatedly in the same goroutine, it's better to create a new Resolver
// to avoid some performance overhead.
func LookupHost(host string) (addrs []string, err error) {
return NewResolver().LookupHost(host)
}
// Dial has the same usage as as net.Dial. This function will use LookupHost
// to resolve host address, then try net.Dial on each returned ip address till
// one succeeds or all fail.
func Dial(hostPort string) (c net.Conn, err error) {
var addrs []string
var host, port string
if host, port, err = net.SplitHostPort(hostPort); err != nil {
return
}
// No need to call LookupHost if host part is IP address
if ip := net.ParseIP(host); ip != nil {
return net.Dial("tcp", hostPort)
}
if addrs, err = LookupHost(host); err != nil {
return
}
for _, ip := range addrs {
ipHost := net.JoinHostPort(ip, port)
if c, err = net.Dial("tcp", ipHost); err == nil {
return
}
}
return nil, err
}