198 lines
4.9 KiB
Go
198 lines
4.9 KiB
Go
|
/*
|
||
|
* MinIO Cloud Storage, (C) 2020 MinIO, Inc.
|
||
|
*
|
||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
* you may not use this file except in compliance with the License.
|
||
|
* You may obtain a copy of the License at
|
||
|
*
|
||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||
|
*
|
||
|
* Unless required by applicable law or agreed to in writing, software
|
||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
* See the License for the specific language governing permissions and
|
||
|
* limitations under the License.
|
||
|
*/
|
||
|
|
||
|
package http
|
||
|
|
||
|
import (
|
||
|
"context"
|
||
|
"log"
|
||
|
"math/rand"
|
||
|
"net"
|
||
|
"sync"
|
||
|
"time"
|
||
|
)
|
||
|
|
||
|
var randPerm = func(n int) []int {
|
||
|
return rand.Perm(n)
|
||
|
}
|
||
|
|
||
|
// DialContextWithDNSCache is a helper function which returns `net.DialContext` function.
|
||
|
// It randomly fetches an IP from the DNS cache and dials it by the given dial
|
||
|
// function. It dials one by one and returns first connected `net.Conn`.
|
||
|
// If it fails to dial all IPs from cache it returns first error. If no baseDialFunc
|
||
|
// is given, it sets default dial function.
|
||
|
//
|
||
|
// You can use returned dial function for `http.Transport.DialContext`.
|
||
|
//
|
||
|
// In this function, it uses functions from `rand` package. To make it really random,
|
||
|
// you MUST call `rand.Seed` and change the value from the default in your application
|
||
|
func DialContextWithDNSCache(cache *DNSCache, baseDialCtx DialContext) DialContext {
|
||
|
if baseDialCtx == nil {
|
||
|
// This is same as which `http.DefaultTransport` uses.
|
||
|
baseDialCtx = (&net.Dialer{
|
||
|
Timeout: 30 * time.Second,
|
||
|
KeepAlive: 30 * time.Second,
|
||
|
DualStack: true,
|
||
|
}).DialContext
|
||
|
}
|
||
|
return func(ctx context.Context, network, host string) (net.Conn, error) {
|
||
|
h, p, err := net.SplitHostPort(host)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
|
||
|
// Fetch DNS result from cache.
|
||
|
//
|
||
|
// ctxLookup is only used for canceling DNS Lookup.
|
||
|
ctxLookup, cancelF := context.WithTimeout(ctx, cache.lookupTimeout)
|
||
|
defer cancelF()
|
||
|
addrs, err := cache.Fetch(ctxLookup, h)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
|
||
|
var firstErr error
|
||
|
for _, randomIndex := range randPerm(len(addrs)) {
|
||
|
conn, err := baseDialCtx(ctx, "tcp", net.JoinHostPort(addrs[randomIndex], p))
|
||
|
if err == nil {
|
||
|
return conn, nil
|
||
|
}
|
||
|
if firstErr == nil {
|
||
|
firstErr = err
|
||
|
}
|
||
|
}
|
||
|
|
||
|
return nil, firstErr
|
||
|
}
|
||
|
}
|
||
|
|
||
|
const (
|
||
|
// cacheSize is initial size of addr and IP list cache map.
|
||
|
cacheSize = 64
|
||
|
)
|
||
|
|
||
|
// defaultFreq is default frequency a resolver refreshes DNS cache.
|
||
|
var (
|
||
|
defaultFreq = 3 * time.Second
|
||
|
defaultLookupTimeout = 10 * time.Second
|
||
|
)
|
||
|
|
||
|
// DNSCache is DNS cache resolver which cache DNS resolve results in memory.
|
||
|
type DNSCache struct {
|
||
|
sync.RWMutex
|
||
|
|
||
|
lookupHostFn func(ctx context.Context, host string) ([]string, error)
|
||
|
lookupTimeout time.Duration
|
||
|
|
||
|
cache map[string][]string
|
||
|
closer func()
|
||
|
}
|
||
|
|
||
|
// NewDNSCache initializes DNS cache resolver and starts auto refreshing
|
||
|
// in a new goroutine. To stop auto refreshing, call `Stop()` function.
|
||
|
// Once `Stop()` is called auto refreshing cannot be resumed.
|
||
|
func NewDNSCache(freq time.Duration, lookupTimeout time.Duration) *DNSCache {
|
||
|
if freq <= 0 {
|
||
|
freq = defaultFreq
|
||
|
}
|
||
|
|
||
|
if lookupTimeout <= 0 {
|
||
|
lookupTimeout = defaultLookupTimeout
|
||
|
}
|
||
|
|
||
|
ticker := time.NewTicker(freq)
|
||
|
ch := make(chan struct{})
|
||
|
closer := func() {
|
||
|
ticker.Stop()
|
||
|
close(ch)
|
||
|
}
|
||
|
|
||
|
r := &DNSCache{
|
||
|
lookupHostFn: net.DefaultResolver.LookupHost,
|
||
|
lookupTimeout: lookupTimeout,
|
||
|
cache: make(map[string][]string, cacheSize),
|
||
|
closer: closer,
|
||
|
}
|
||
|
|
||
|
go func() {
|
||
|
for {
|
||
|
select {
|
||
|
case <-ticker.C:
|
||
|
r.Refresh()
|
||
|
case <-ch:
|
||
|
return
|
||
|
}
|
||
|
}
|
||
|
}()
|
||
|
|
||
|
return r
|
||
|
}
|
||
|
|
||
|
// LookupHost lookups address list from DNS server, persist the results
|
||
|
// in-memory cache. `Fetch` is used to obtain the values for a given host.
|
||
|
func (r *DNSCache) LookupHost(ctx context.Context, host string) ([]string, error) {
|
||
|
addrs, err := r.lookupHostFn(ctx, host)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
|
||
|
r.Lock()
|
||
|
r.cache[host] = addrs
|
||
|
r.Unlock()
|
||
|
|
||
|
return addrs, nil
|
||
|
}
|
||
|
|
||
|
// Fetch fetches IP list from the cache. If IP list of the given addr is not in the cache,
|
||
|
// then it lookups from DNS server by `Lookup` function.
|
||
|
func (r *DNSCache) Fetch(ctx context.Context, host string) ([]string, error) {
|
||
|
r.RLock()
|
||
|
addrs, ok := r.cache[host]
|
||
|
r.RUnlock()
|
||
|
if ok {
|
||
|
return addrs, nil
|
||
|
}
|
||
|
return r.LookupHost(ctx, host)
|
||
|
}
|
||
|
|
||
|
// Refresh refreshes IP list cache, automatically.
|
||
|
func (r *DNSCache) Refresh() {
|
||
|
r.RLock()
|
||
|
hosts := make([]string, 0, len(r.cache))
|
||
|
for host := range r.cache {
|
||
|
hosts = append(hosts, host)
|
||
|
}
|
||
|
r.RUnlock()
|
||
|
|
||
|
for _, host := range hosts {
|
||
|
ctx, cancelF := context.WithTimeout(context.Background(), r.lookupTimeout)
|
||
|
if _, err := r.LookupHost(ctx, host); err != nil {
|
||
|
log.Println("failed to refresh DNS cache, resolver is unavailable", err)
|
||
|
}
|
||
|
cancelF()
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Stop stops auto refreshing.
|
||
|
func (r *DNSCache) Stop() {
|
||
|
r.Lock()
|
||
|
defer r.Unlock()
|
||
|
if r.closer != nil {
|
||
|
r.closer()
|
||
|
r.closer = nil
|
||
|
}
|
||
|
}
|