Merge pull request #5 from hailocab/maybe-get

Added ability to prevent unhealthy nodes from being returned

epsilon_greedy.go

@@ -104,6 +104,10 @@ func (p *epsilonGreedyHostPool) Get() HostPoolResponse {
 	p.Lock()
 	defer p.Unlock()
 	host := p.getEpsilonGreedy()
+	if host == "" {
+		return nil
+	}
+
 	started := time.Now()
 	return &epsilonHostPoolResponse{
 		standardHostPoolResponse: standardHostPoolResponse{host: host, pool: p},
@@ -161,6 +165,7 @@ func (p *epsilonGreedyHostPool) getEpsilonGreedy() string {
 		if len(possibleHosts) != 0 {
 			log.Println("Failed to randomly choose a host, Dan loses")
 		}
+
 		return p.getRoundRobin()
 	}
 

hostpool.go

@@ -44,6 +44,10 @@ type HostPool interface {
 	markFailed(HostPoolResponse)
 
 	ResetAll()
+	// ReturnUnhealthy when called with true will prevent an unhealthy node from
+	// being returned and will instead return a nil HostPoolResponse. If using
+	// this feature then you should check the result of Get for nil
+	ReturnUnhealthy(v bool)
 	Hosts() []string
 	SetHosts([]string)
 }
@@ -52,6 +56,7 @@ type standardHostPool struct {
 	sync.RWMutex
 	hosts             map[string]*hostEntry
 	hostList          []*hostEntry
+	returnUnhealthy   bool
 	initialRetryDelay time.Duration
 	maxRetryInterval  time.Duration
 	nextHostIndex     int
@@ -68,6 +73,7 @@ const defaultDecayDuration = time.Duration(5) * time.Minute
 // Construct a basic HostPool using the hostnames provided
 func New(hosts []string) HostPool {
 	p := &standardHostPool{
+		returnUnhealthy:   true,
 		hosts:             make(map[string]*hostEntry, len(hosts)),
 		hostList:          make([]*hostEntry, len(hosts)),
 		initialRetryDelay: time.Duration(30) * time.Second,
@@ -113,6 +119,10 @@ func (p *standardHostPool) Get() HostPoolResponse {
 	p.Lock()
 	defer p.Unlock()
 	host := p.getRoundRobin()
+	if host == "" {
+		return nil
+	}
+
 	return &standardHostPoolResponse{host: host, pool: p}
 }
 
@@ -135,6 +145,11 @@ func (p *standardHostPool) getRoundRobin() string {
 		}
 	}
 
+	// all hosts are down and returnUnhealhy is false then return no host
+	if p.returnUnhealthy {
+		return ""
+	}
+
 	// all hosts are down. re-add them
 	p.doResetAll()
 	p.nextHostIndex = 0
@@ -153,6 +168,12 @@ func (p *standardHostPool) SetHosts(hosts []string) {
 	p.setHosts(hosts)
 }
 
+func (p *standardHostPool) ReturnUnhealthy(v bool) {
+	p.Lock()
+	defer p.Unlock()
+	p.returnUnhealthy = v
+}
+
 func (p *standardHostPool) setHosts(hosts []string) {
 	p.hosts = make(map[string]*hostEntry, len(hosts))
 	p.hostList = make([]*hostEntry, len(hosts))