diff --git a/plugin/forward/forward.go b/plugin/forward/forward.go index 20a3fcf38..59e27b216 100644 --- a/plugin/forward/forward.go +++ b/plugin/forward/forward.go @@ -14,7 +14,6 @@ import ( "github.com/coredns/coredns/plugin" "github.com/coredns/coredns/plugin/debug" clog "github.com/coredns/coredns/plugin/pkg/log" - "github.com/coredns/coredns/plugin/pkg/policy" "github.com/coredns/coredns/request" "github.com/miekg/dns" @@ -29,7 +28,7 @@ type Forward struct { concurrent int64 // atomic counters need to be first in struct for proper alignment proxies []*Proxy - p policy.Policy + p Policy hcInterval time.Duration from string @@ -52,7 +51,7 @@ type Forward struct { // New returns a new Forward. func New() *Forward { - f := &Forward{maxfails: 2, tlsConfig: new(tls.Config), expire: defaultExpire, p: new(policy.Random), from: ".", hcInterval: hcInterval, opts: options{forceTCP: false, preferUDP: false, hcRecursionDesired: true}} + f := &Forward{maxfails: 2, tlsConfig: new(tls.Config), expire: defaultExpire, p: new(random), from: ".", hcInterval: hcInterval, opts: options{forceTCP: false, preferUDP: false, hcRecursionDesired: true}} return f } @@ -109,8 +108,8 @@ func (f *Forward) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg } // All upstream proxies are dead, assume healthcheck is completely broken and randomly // select an upstream to connect to. - r := new(policy.Random) - proxy = r.List(f.proxies)[0].([]*Proxy)[0] + r := new(random) + proxy = r.List(f.proxies)[0] HealthcheckBrokenCount.Add(1) } @@ -206,12 +205,7 @@ func (f *Forward) ForceTCP() bool { return f.opts.forceTCP } func (f *Forward) PreferUDP() bool { return f.opts.preferUDP } // List returns a set of proxies to be used for this client depending on the policy in f. -func (f *Forward) List() []*Proxy { - if len(f.p.List(f.proxies)) == 1 { - return f.p.List(f.proxies)[0].([]*Proxy) - } - return nil -} +func (f *Forward) List() []*Proxy {return f.p.List(f.proxies)} var ( // ErrNoHealthy means no healthy proxies left. diff --git a/plugin/forward/forward_test.go b/plugin/forward/forward_test.go new file mode 100644 index 000000000..b0ef47ba9 --- /dev/null +++ b/plugin/forward/forward_test.go @@ -0,0 +1,24 @@ +package forward + +import ( + "testing" +) + +func TestList(t *testing.T) { + f := Forward{ + proxies: []*Proxy{{addr: "1.1.1.1:53"}, {addr: "2.2.2.2:53"}, {addr: "3.3.3.3:53"}}, + p: &roundRobin{}, + } + + expect := []*Proxy{{addr: "2.2.2.2:53"}, {addr: "1.1.1.1:53"}, {addr: "3.3.3.3:53"}} + got := f.List() + + if len(got) != len(expect) { + t.Fatalf("Expected: %v results, got: %v", len(expect), len(got)) + } + for i, p := range got { + if p.addr != expect[i].addr { + t.Fatalf("Expected proxy %v to be '%v', got: '%v'", i, expect[i].addr, p.addr) + } + } +} diff --git a/plugin/forward/policy.go b/plugin/forward/policy.go new file mode 100644 index 000000000..2066e1316 --- /dev/null +++ b/plugin/forward/policy.go @@ -0,0 +1,64 @@ +package forward + +import ( + "math/rand" + "sync/atomic" +) + +// Policy defines a policy we use for selecting upstreams. +type Policy interface { + List([]*Proxy) []*Proxy + String() string +} + +// random is a policy that implements random upstream selection. +type random struct{} + +func (r *random) String() string { return "random" } + +func (r *random) List(p []*Proxy) []*Proxy { + switch len(p) { + case 1: + return p + case 2: + if rand.Int()%2 == 0 { + return []*Proxy{p[1], p[0]} // swap + } + return p + } + + perms := rand.Perm(len(p)) + rnd := make([]*Proxy, len(p)) + + for i, p1 := range perms { + rnd[i] = p[p1] + } + return rnd +} + +// roundRobin is a policy that selects hosts based on round robin ordering. +type roundRobin struct { + robin uint32 +} + +func (r *roundRobin) String() string { return "round_robin" } + +func (r *roundRobin) List(p []*Proxy) []*Proxy { + poolLen := uint32(len(p)) + i := atomic.AddUint32(&r.robin, 1) % poolLen + + robin := []*Proxy{p[i]} + robin = append(robin, p[:i]...) + robin = append(robin, p[i+1:]...) + + return robin +} + +// sequential is a policy that selects hosts based on sequential ordering. +type sequential struct{} + +func (r *sequential) String() string { return "sequential" } + +func (r *sequential) List(p []*Proxy) []*Proxy { + return p +} diff --git a/plugin/forward/setup.go b/plugin/forward/setup.go index b5eeab912..9d873a411 100644 --- a/plugin/forward/setup.go +++ b/plugin/forward/setup.go @@ -10,7 +10,6 @@ import ( "github.com/coredns/coredns/plugin" "github.com/coredns/coredns/plugin/metrics" "github.com/coredns/coredns/plugin/pkg/parse" - "github.com/coredns/coredns/plugin/pkg/policy" pkgtls "github.com/coredns/coredns/plugin/pkg/tls" "github.com/coredns/coredns/plugin/pkg/transport" @@ -221,11 +220,11 @@ func parseBlock(c *caddy.Controller, f *Forward) error { } switch x := c.Val(); x { case "random": - f.p = &policy.Random{} + f.p = &random{} case "round_robin": - f.p = &policy.RoundRobin{} + f.p = &roundRobin{} case "sequential": - f.p = &policy.Sequential{} + f.p = &sequential{} default: return c.Errf("unknown policy '%s'", x) } diff --git a/plugin/grpc/grpc.go b/plugin/grpc/grpc.go index 878b6c3f4..3dda225df 100644 --- a/plugin/grpc/grpc.go +++ b/plugin/grpc/grpc.go @@ -7,7 +7,6 @@ import ( "github.com/coredns/coredns/plugin" "github.com/coredns/coredns/plugin/debug" - "github.com/coredns/coredns/plugin/pkg/policy" "github.com/coredns/coredns/request" "github.com/miekg/dns" @@ -18,7 +17,7 @@ import ( // It has a list of proxies each representing one upstream proxy. type GRPC struct { proxies []*Proxy - p policy.Policy + p Policy from string ignored []string @@ -94,7 +93,7 @@ func (g *GRPC) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) ( // NewGRPC returns a new GRPC. func newGRPC() *GRPC { g := &GRPC{ - p: new(policy.Random), + p: new(random), } return g } @@ -127,11 +126,6 @@ func (g *GRPC) isAllowedDomain(name string) bool { } // List returns a set of proxies to be used for this client depending on the policy in p. -func (g *GRPC) list() []*Proxy { - if len(g.p.List(g.proxies)) == 1 { - return g.p.List(g.proxies)[0].([]*Proxy) - } - return nil -} +func (g *GRPC) list() []*Proxy { return g.p.List(g.proxies) } const defaultTimeout = 5 * time.Second diff --git a/plugin/grpc/policy.go b/plugin/grpc/policy.go new file mode 100644 index 000000000..66351d822 --- /dev/null +++ b/plugin/grpc/policy.go @@ -0,0 +1,64 @@ +package grpc + +import ( + "math/rand" + "sync/atomic" +) + +// Policy defines a policy we use for selecting upstreams. +type Policy interface { + List([]*Proxy) []*Proxy + String() string +} + +// random is a policy that implements random upstream selection. +type random struct{} + +func (r *random) String() string { return "random" } + +func (r *random) List(p []*Proxy) []*Proxy { + switch len(p) { + case 1: + return p + case 2: + if rand.Int()%2 == 0 { + return []*Proxy{p[1], p[0]} // swap + } + return p + } + + perms := rand.Perm(len(p)) + rnd := make([]*Proxy, len(p)) + + for i, p1 := range perms { + rnd[i] = p[p1] + } + return rnd +} + +// roundRobin is a policy that selects hosts based on round robin ordering. +type roundRobin struct { + robin uint32 +} + +func (r *roundRobin) String() string { return "round_robin" } + +func (r *roundRobin) List(p []*Proxy) []*Proxy { + poolLen := uint32(len(p)) + i := atomic.AddUint32(&r.robin, 1) % poolLen + + robin := []*Proxy{p[i]} + robin = append(robin, p[:i]...) + robin = append(robin, p[i+1:]...) + + return robin +} + +// sequential is a policy that selects hosts based on sequential ordering. +type sequential struct{} + +func (r *sequential) String() string { return "sequential" } + +func (r *sequential) List(p []*Proxy) []*Proxy { + return p +} diff --git a/plugin/grpc/setup.go b/plugin/grpc/setup.go index c178932c4..a234efb37 100644 --- a/plugin/grpc/setup.go +++ b/plugin/grpc/setup.go @@ -8,7 +8,6 @@ import ( "github.com/coredns/coredns/plugin" "github.com/coredns/coredns/plugin/metrics" "github.com/coredns/coredns/plugin/pkg/parse" - "github.com/coredns/coredns/plugin/pkg/policy" pkgtls "github.com/coredns/coredns/plugin/pkg/tls" "github.com/caddyserver/caddy" @@ -133,11 +132,11 @@ func parseBlock(c *caddy.Controller, g *GRPC) error { } switch x := c.Val(); x { case "random": - g.p = &policy.Random{} + g.p = &random{} case "round_robin": - g.p = &policy.RoundRobin{} + g.p = &roundRobin{} case "sequential": - g.p = &policy.Sequential{} + g.p = &sequential{} default: return c.Errf("unknown policy '%s'", x) } diff --git a/plugin/pkg/policy/policy.go b/plugin/pkg/policy/policy.go deleted file mode 100644 index c5e9de6f4..000000000 --- a/plugin/pkg/policy/policy.go +++ /dev/null @@ -1,76 +0,0 @@ -package policy - -import ( - "math/rand" - "sync/atomic" -) - -// Policy defines a policy we use for selecting upstreams. -type Policy interface { - List(policy ...interface{}) []interface{} - String() string -} - -// Random is a policy that implements random upstream selection. -type Random struct{} - -var _ Policy = &Random{} - -// String returns the name of policy Random -func (r *Random) String() string { return "random" } - -// List returns a set of proxies to be used for this client depending on Random policy. -func (r *Random) List(p ...interface{}) []interface{} { - switch len(p) { - case 1: - return p - case 2: - if rand.Int()%2 == 0 { - return []interface{}{p[1], p[0]} // swap - } - return p - } - - perms := rand.Perm(len(p)) - rnd := make([]interface{}, len(p)) - - for i, p1 := range perms { - rnd[i] = p[p1] - } - return rnd -} - -// RoundRobin is a policy that selects hosts based on round robin ordering. -type RoundRobin struct { - robin uint32 -} - -var _ Policy = &RoundRobin{} - -// String returns the name of policy RoundRobin -func (r *RoundRobin) String() string { return "round_robin" } - -// List returns a set of proxies to be used for this client depending on RoundRobin policy. -func (r *RoundRobin) List(p ...interface{}) []interface{} { - poolLen := uint32(len(p)) - i := atomic.AddUint32(&r.robin, 1) % poolLen - - robin := []interface{}{p[i]} - robin = append(robin, p[:i]...) - robin = append(robin, p[i+1:]...) - - return robin -} - -// Sequential is a policy that selects hosts based on sequential ordering. -type Sequential struct{} - -var _ Policy = &Sequential{} - -// String returns the name of policy Sequential -func (r *Sequential) String() string { return "sequential" } - -// List returns a set of proxies without filter. -func (r *Sequential) List(p ...interface{}) []interface{} { - return p -}