Added ability to prevent unhealthy nodes from being returned

This could be useful if you would prefer to fall back to another process
if there are no healthy nodes available or if you wanted to "fail fast".

This PR updates the Get method to optionally return `nil` if
`ReturnUnhealthy` has been called with false.

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