Commit 9bdd9bde authored by Matt Holt's avatar Matt Holt

Merge pull request #94 from acmacalister/feature/custom-policies

added custom policy support
parents 130301e3 dd946f8a
...@@ -13,6 +13,12 @@ type Policy interface { ...@@ -13,6 +13,12 @@ type Policy interface {
Select(pool HostPool) *UpstreamHost 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. // Random is a policy that selects up hosts from a pool at random.
type Random struct{} type Random struct{}
......
...@@ -4,6 +4,12 @@ import ( ...@@ -4,6 +4,12 @@ import (
"testing" "testing"
) )
type customPolicy struct{}
func (r *customPolicy) Select(pool HostPool) *UpstreamHost {
return pool[0]
}
func testPool() HostPool { func testPool() HostPool {
pool := []*UpstreamHost{ pool := []*UpstreamHost{
&UpstreamHost{ &UpstreamHost{
...@@ -55,3 +61,12 @@ func TestLeastConnPolicy(t *testing.T) { ...@@ -55,3 +61,12 @@ func TestLeastConnPolicy(t *testing.T) {
t.Error("Expected least connection host to be first or second host.") 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.")
}
}
...@@ -12,6 +12,8 @@ import ( ...@@ -12,6 +12,8 @@ import (
"github.com/mholt/caddy/config/parse" "github.com/mholt/caddy/config/parse"
) )
var supportedPolicies map[string]func() Policy = make(map[string]func() Policy)
type staticUpstream struct { type staticUpstream struct {
from string from string
Hosts HostPool Hosts HostPool
...@@ -53,14 +55,10 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) { ...@@ -53,14 +55,10 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) {
if !c.NextArg() { if !c.NextArg() {
return upstreams, c.ArgErr() return upstreams, c.ArgErr()
} }
switch c.Val() {
case "random": if policyCreateFunc, ok := supportedPolicies[c.Val()]; ok {
upstream.Policy = &Random{} upstream.Policy = policyCreateFunc()
case "round_robin": } else {
upstream.Policy = &RoundRobin{}
case "least_conn":
upstream.Policy = &LeastConn{}
default:
return upstreams, c.ArgErr() return upstreams, c.ArgErr()
} }
case "fail_timeout": case "fail_timeout":
...@@ -147,6 +145,11 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) { ...@@ -147,6 +145,11 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) {
return upstreams, nil return upstreams, nil
} }
// RegisterPolicy adds a custom policy to the proxy.
func RegisterPolicy(name string, policy func() Policy) {
supportedPolicies[name] = policy
}
func (u *staticUpstream) From() string { func (u *staticUpstream) From() string {
return u.from return u.from
} }
......
...@@ -41,3 +41,13 @@ func TestSelect(t *testing.T) { ...@@ -41,3 +41,13 @@ func TestSelect(t *testing.T) {
t.Error("Expected select to not return nil") t.Error("Expected select to not return nil")
} }
} }
func TestRegisterPolicy(t *testing.T) {
name := "custom"
customPolicy := &customPolicy{}
RegisterPolicy(name, func() Policy { return customPolicy })
if _, ok := supportedPolicies[name]; !ok {
t.Error("Expected supportedPolicies to have a custom policy.")
}
}
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment