plugin/forward/grpc: Revert forward/grpc policy dedup (#3919)

* revert de-dup

Signed-off-by: Chris O'Haver <cohaver@infoblox.com>

* unit test

Signed-off-by: Chris O'Haver <cohaver@infoblox.com>

* use roundrobin policy in test

Signed-off-by: Chris O'Haver <cohaver@infoblox.com>

4 files changed, 96 insertions, 15 deletions

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)
 		}