Merge pull request bitly#1 from hailocab/sethosts

Allow hosts for an existing hostpool to be changed

Author: Boyan Dimitrov

epsilon_greedy.go

@@ -71,6 +71,16 @@ func (p *epsilonGreedyHostPool) SetEpsilon(newEpsilon float32) {
 	p.epsilon = newEpsilon
 }
 
+func (p *epsilonGreedyHostPool) SetHosts(hosts []string) {
+	p.Lock()
+	defer p.Unlock()
+	p.standardHostPool.setHosts(hosts)
+	for _, h := range p.hostList {
+		h.epsilonCounts = make([]int64, epsilonBuckets)
+		h.epsilonValues = make([]int64, epsilonBuckets)
+	}
+}
+
 func (p *epsilonGreedyHostPool) epsilonGreedyDecay() {
 	durationPerBucket := p.decayDuration / epsilonBuckets
 	ticker := time.Tick(durationPerBucket)

hostpool.go

@@ -45,6 +45,7 @@ type HostPool interface {
 
 	ResetAll()
 	Hosts() []string
+	SetHosts([]string)
 }
 
 type standardHostPool struct {
@@ -146,6 +147,26 @@ func (p *standardHostPool) ResetAll() {
 	p.doResetAll()
 }
 
+func (p *standardHostPool) SetHosts(hosts []string) {
+	p.Lock()
+	defer p.Unlock()
+	p.setHosts(hosts)
+}
+
+func (p *standardHostPool) setHosts(hosts []string) {
+	p.hosts = make(map[string]*hostEntry, len(hosts))
+	p.hostList = make([]*hostEntry, len(hosts))
+
+	for i, h := range hosts {
+		e := &hostEntry{
+			host:       h,
+			retryDelay: p.initialRetryDelay,
+		}
+		p.hosts[h] = e
+		p.hostList[i] = e
+	}
+}
+
 // this actually performs the logic to reset,
 // and should only be called when the lock has
 // already been acquired