From 6b173b517049d23da228830540a74768d8ef2de4 Mon Sep 17 00:00:00 2001 From: Austin Date: Thu, 28 May 2015 15:56:11 -0700 Subject: [PATCH 1/3] added custom policy support --- middleware/proxy/policy.go | 16 ++++++++++++++++ middleware/proxy/policy_test.go | 19 +++++++++++++++++++ middleware/proxy/upstream.go | 23 +++++++++++++++-------- middleware/proxy/upstream_test.go | 9 +++++++++ 4 files changed, 59 insertions(+), 8 deletions(-) diff --git a/middleware/proxy/policy.go b/middleware/proxy/policy.go index 27e02a177..f90a1e239 100644 --- a/middleware/proxy/policy.go +++ b/middleware/proxy/policy.go @@ -11,6 +11,7 @@ type HostPool []*UpstreamHost // Policy decides how a host will be selected from a pool. type Policy interface { Select(pool HostPool) *UpstreamHost + Name() string } // Random is a policy that selects up hosts from a pool at random. @@ -39,6 +40,11 @@ func (r *Random) Select(pool HostPool) *UpstreamHost { return randHost } +// Name returns the name of the policy. +func (r *Random) Name() string { + return "random" +} + // LeastConn is a policy that selects the host with the least connections. type LeastConn struct{} @@ -74,6 +80,11 @@ func (r *LeastConn) Select(pool HostPool) *UpstreamHost { return bestHost } +// Name returns the name of the policy. +func (r *LeastConn) Name() string { + return "least_conn" +} + // RoundRobin is a policy that selects hosts based on round robin ordering. type RoundRobin struct { Robin uint32 @@ -93,3 +104,8 @@ func (r *RoundRobin) Select(pool HostPool) *UpstreamHost { } return host } + +// Name returns the name of the policy. +func (r *RoundRobin) Name() string { + return "round_robin" +} diff --git a/middleware/proxy/policy_test.go b/middleware/proxy/policy_test.go index 11269a4f2..4272e332a 100644 --- a/middleware/proxy/policy_test.go +++ b/middleware/proxy/policy_test.go @@ -4,6 +4,16 @@ import ( "testing" ) +type customPolicy struct{} + +func (r *customPolicy) Select(pool HostPool) *UpstreamHost { + return pool[0] +} + +func (r *customPolicy) Name() string { + return "custom" +} + func testPool() HostPool { pool := []*UpstreamHost{ &UpstreamHost{ @@ -55,3 +65,12 @@ func TestLeastConnPolicy(t *testing.T) { t.Error("Expected least connection host to be first or second host.") } } + +func TestCustomPolicy(t *testing.T) { + pool := testPool() + customPolicy := &customPolicy{} + h := customPolicy.Select(pool) + if h != pool[0] { + t.Error("Expected custom policy host to be the first host.") + } +} diff --git a/middleware/proxy/upstream.go b/middleware/proxy/upstream.go index c724da78e..58325f1d4 100644 --- a/middleware/proxy/upstream.go +++ b/middleware/proxy/upstream.go @@ -12,6 +12,8 @@ import ( "github.com/mholt/caddy/config/parse" ) +var supportedPolicies map[string]Policy = make(map[string]Policy) + type staticUpstream struct { from string Hosts HostPool @@ -30,6 +32,10 @@ type staticUpstream struct { func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) { var upstreams []Upstream + RegisterPolicy(&Random{}) + RegisterPolicy(&LeastConn{}) + RegisterPolicy(&RoundRobin{}) + for c.Next() { upstream := &staticUpstream{ from: "", @@ -53,16 +59,12 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) { if !c.NextArg() { return upstreams, c.ArgErr() } - switch c.Val() { - case "random": - upstream.Policy = &Random{} - case "round_robin": - upstream.Policy = &RoundRobin{} - case "least_conn": - upstream.Policy = &LeastConn{} - default: + + policy, ok := supportedPolicies[c.Val()] + if !ok { return upstreams, c.ArgErr() } + upstream.Policy = policy case "fail_timeout": if !c.NextArg() { return upstreams, c.ArgErr() @@ -147,6 +149,11 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) { return upstreams, nil } +// RegisterPolicy adds a custom policy to the proxy. +func RegisterPolicy(policy Policy) { + supportedPolicies[policy.Name()] = policy +} + func (u *staticUpstream) From() string { return u.from } diff --git a/middleware/proxy/upstream_test.go b/middleware/proxy/upstream_test.go index 6be3f6cea..1d1cc317d 100644 --- a/middleware/proxy/upstream_test.go +++ b/middleware/proxy/upstream_test.go @@ -41,3 +41,12 @@ func TestSelect(t *testing.T) { t.Error("Expected select to not return nil") } } + +func TestRegisterPolicy(t *testing.T) { + customPolicy := &customPolicy{} + RegisterPolicy(customPolicy) + if _, ok := supportedPolicies[customPolicy.Name()]; !ok { + t.Error("Expected supportedPolicies to have a custom policy.") + } + +} From 593aec9ab1e00500200ac194f7d73e5024019bde Mon Sep 17 00:00:00 2001 From: Austin Date: Thu, 28 May 2015 16:53:54 -0700 Subject: [PATCH 2/3] changes per comment --- middleware/proxy/policy.go | 16 ---------------- middleware/proxy/policy_test.go | 4 ---- middleware/proxy/upstream.go | 22 ++++++++++++---------- middleware/proxy/upstream_test.go | 5 +++-- 4 files changed, 15 insertions(+), 32 deletions(-) diff --git a/middleware/proxy/policy.go b/middleware/proxy/policy.go index f90a1e239..27e02a177 100644 --- a/middleware/proxy/policy.go +++ b/middleware/proxy/policy.go @@ -11,7 +11,6 @@ type HostPool []*UpstreamHost // Policy decides how a host will be selected from a pool. type Policy interface { Select(pool HostPool) *UpstreamHost - Name() string } // Random is a policy that selects up hosts from a pool at random. @@ -40,11 +39,6 @@ func (r *Random) Select(pool HostPool) *UpstreamHost { return randHost } -// Name returns the name of the policy. -func (r *Random) Name() string { - return "random" -} - // LeastConn is a policy that selects the host with the least connections. type LeastConn struct{} @@ -80,11 +74,6 @@ func (r *LeastConn) Select(pool HostPool) *UpstreamHost { return bestHost } -// Name returns the name of the policy. -func (r *LeastConn) Name() string { - return "least_conn" -} - // RoundRobin is a policy that selects hosts based on round robin ordering. type RoundRobin struct { Robin uint32 @@ -104,8 +93,3 @@ func (r *RoundRobin) Select(pool HostPool) *UpstreamHost { } return host } - -// Name returns the name of the policy. -func (r *RoundRobin) Name() string { - return "round_robin" -} diff --git a/middleware/proxy/policy_test.go b/middleware/proxy/policy_test.go index 4272e332a..d4c752256 100644 --- a/middleware/proxy/policy_test.go +++ b/middleware/proxy/policy_test.go @@ -10,10 +10,6 @@ func (r *customPolicy) Select(pool HostPool) *UpstreamHost { return pool[0] } -func (r *customPolicy) Name() string { - return "custom" -} - func testPool() HostPool { pool := []*UpstreamHost{ &UpstreamHost{ diff --git a/middleware/proxy/upstream.go b/middleware/proxy/upstream.go index 58325f1d4..03bba0e6f 100644 --- a/middleware/proxy/upstream.go +++ b/middleware/proxy/upstream.go @@ -12,7 +12,7 @@ import ( "github.com/mholt/caddy/config/parse" ) -var supportedPolicies map[string]Policy = make(map[string]Policy) +var supportedPolicies map[string]func() Policy = make(map[string]func() Policy) type staticUpstream struct { from string @@ -27,15 +27,17 @@ type staticUpstream struct { } } +func init() { + RegisterPolicy("random", func() Policy { return &Random{} }) + RegisterPolicy("least_conn", func() Policy { return &LeastConn{} }) + RegisterPolicy("round_robin", func() Policy { return &RoundRobin{} }) +} + // NewStaticUpstreams parses the configuration input and sets up // static upstreams for the proxy middleware. func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) { var upstreams []Upstream - RegisterPolicy(&Random{}) - RegisterPolicy(&LeastConn{}) - RegisterPolicy(&RoundRobin{}) - for c.Next() { upstream := &staticUpstream{ from: "", @@ -60,11 +62,11 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) { return upstreams, c.ArgErr() } - policy, ok := supportedPolicies[c.Val()] - if !ok { + if policyCreateFunc, ok := supportedPolicies[c.Val()]; ok { + upstream.Policy = policyCreateFunc() + } else { return upstreams, c.ArgErr() } - upstream.Policy = policy case "fail_timeout": if !c.NextArg() { return upstreams, c.ArgErr() @@ -150,8 +152,8 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) { } // RegisterPolicy adds a custom policy to the proxy. -func RegisterPolicy(policy Policy) { - supportedPolicies[policy.Name()] = policy +func RegisterPolicy(name string, policy func() Policy) { + supportedPolicies[name] = policy } func (u *staticUpstream) From() string { diff --git a/middleware/proxy/upstream_test.go b/middleware/proxy/upstream_test.go index 1d1cc317d..f3df16136 100644 --- a/middleware/proxy/upstream_test.go +++ b/middleware/proxy/upstream_test.go @@ -43,9 +43,10 @@ func TestSelect(t *testing.T) { } func TestRegisterPolicy(t *testing.T) { + name := "custom" customPolicy := &customPolicy{} - RegisterPolicy(customPolicy) - if _, ok := supportedPolicies[customPolicy.Name()]; !ok { + RegisterPolicy(name, func() Policy { return customPolicy }) + if _, ok := supportedPolicies[name]; !ok { t.Error("Expected supportedPolicies to have a custom policy.") } From dd946f8ab56fb861bfa026a1aa911f502838a0ed Mon Sep 17 00:00:00 2001 From: Austin Date: Thu, 28 May 2015 18:16:23 -0700 Subject: [PATCH 3/3] moved init to policy.go --- middleware/proxy/policy.go | 6 ++++++ middleware/proxy/upstream.go | 6 ------ 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/middleware/proxy/policy.go b/middleware/proxy/policy.go index 27e02a177..a2522bcb1 100644 --- a/middleware/proxy/policy.go +++ b/middleware/proxy/policy.go @@ -13,6 +13,12 @@ type Policy interface { Select(pool HostPool) *UpstreamHost } +func init() { + RegisterPolicy("random", func() Policy { return &Random{} }) + RegisterPolicy("least_conn", func() Policy { return &LeastConn{} }) + RegisterPolicy("round_robin", func() Policy { return &RoundRobin{} }) +} + // Random is a policy that selects up hosts from a pool at random. type Random struct{} diff --git a/middleware/proxy/upstream.go b/middleware/proxy/upstream.go index 03bba0e6f..a657a088e 100644 --- a/middleware/proxy/upstream.go +++ b/middleware/proxy/upstream.go @@ -27,12 +27,6 @@ type staticUpstream struct { } } -func init() { - RegisterPolicy("random", func() Policy { return &Random{} }) - RegisterPolicy("least_conn", func() Policy { return &LeastConn{} }) - RegisterPolicy("round_robin", func() Policy { return &RoundRobin{} }) -} - // NewStaticUpstreams parses the configuration input and sets up // static upstreams for the proxy middleware. func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) {