Commit 78fd3911658fc953a89cc0e270dcd3aee3741001 - golang-github-go-kit-kit

+0

-47

~~loadbalancer/cache.go~~ less more

0		package loadbalancer
1
2		import "github.com/go-kit/kit/endpoint"
3
4		type cache struct {
5		req chan []endpoint.Endpoint
6		cnt chan int
7		quit chan struct{}
8		}
9
10		func newCache(p Publisher) *cache {
11		c := &cache{
12		req: make(chan []endpoint.Endpoint),
13		cnt: make(chan int),
14		quit: make(chan struct{}),
15		}
16		go c.loop(p)
17		return c
18		}
19
20		func (c *cache) loop(p Publisher) {
21		e := make(chan []endpoint.Endpoint, 1)
22		p.Subscribe(e)
23		defer p.Unsubscribe(e)
24		endpoints := <-e
25		for {
26		select {
27		case endpoints = <-e:
28		case c.cnt <- len(endpoints):
29		case c.req <- endpoints:
30		case <-c.quit:
31		return
32		}
33		}
34		}
35
36		func (c *cache) count() int {
37		return <-c.cnt
38		}
39
40		func (c *cache) get() []endpoint.Endpoint {
41		return <-c.req
42		}
43
44		func (c *cache) stop() {
45		close(c.quit)
46		}

+0

-33

~~loadbalancer/cache_internal_test.go~~ less more

0		package loadbalancer
1
2		import (
3		"runtime"
4		"testing"
5
6		"golang.org/x/net/context"
7
8		"github.com/go-kit/kit/endpoint"
9		)
10
11		func TestCache(t *testing.T) {
12		e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
13		endpoints := []endpoint.Endpoint{e}
14
15		p := NewStaticPublisher(endpoints)
16		defer p.Stop()
17
18		c := newCache(p)
19		defer c.stop()
20
21		for _, n := range []int{2, 10, 0} {
22		endpoints = make([]endpoint.Endpoint, n)
23		for i := 0; i < n; i++ {
24		endpoints[i] = e
25		}
26		p.Replace(endpoints)
27		runtime.Gosched()
28		if want, have := len(endpoints), len(c.get()); want != have {
29		t.Errorf("want %d, have %d", want, have)
30		}
31		}
32		}

+0

-105

~~loadbalancer/dns_srv_publisher.go~~ less more

0		package loadbalancer
1
2		import (
3		"crypto/md5"
4		"fmt"
5		"net"
6		"sort"
7		"time"
8
9		"github.com/go-kit/kit/endpoint"
10		)
11
12		type dnssrvPublisher struct {
13		subscribe chan chan<- []endpoint.Endpoint
14		unsubscribe chan chan<- []endpoint.Endpoint
15		quit chan struct{}
16		}
17
18		// NewDNSSRVPublisher returns a publisher that resolves the SRV name every ttl, and
19		func NewDNSSRVPublisher(name string, ttl time.Duration, makeEndpoint func(hostport string) endpoint.Endpoint) Publisher {
20		p := &dnssrvPublisher{
21		subscribe: make(chan chan<- []endpoint.Endpoint),
22		unsubscribe: make(chan chan<- []endpoint.Endpoint),
23		quit: make(chan struct{}),
24		}
25		go p.loop(name, ttl, makeEndpoint)
26		return p
27		}
28
29		func (p *dnssrvPublisher) Subscribe(c chan<- []endpoint.Endpoint) {
30		p.subscribe <- c
31		}
32
33		func (p *dnssrvPublisher) Unsubscribe(c chan<- []endpoint.Endpoint) {
34		p.unsubscribe <- c
35		}
36
37		func (p *dnssrvPublisher) Stop() {
38		close(p.quit)
39		}
40
41		var newTicker = time.NewTicker
42
43		func (p *dnssrvPublisher) loop(name string, ttl time.Duration, makeEndpoint func(hostport string) endpoint.Endpoint) {
44		var (
45		subscriptions = map[chan<- []endpoint.Endpoint]struct{}{}
46		addrs, md5, _ = resolve(name)
47		endpoints = convert(addrs, makeEndpoint)
48		ticker = newTicker(ttl)
49		)
50		defer ticker.Stop()
51		for {
52		select {
53		case <-ticker.C:
54		addrs, newmd5, err := resolve(name)
55		if err == nil && newmd5 != md5 {
56		endpoints = convert(addrs, makeEndpoint)
57		for c := range subscriptions {
58		c <- endpoints
59		}
60		md5 = newmd5
61		}
62
63		case c := <-p.subscribe:
64		subscriptions[c] = struct{}{}
65		c <- endpoints
66
67		case c := <-p.unsubscribe:
68		delete(subscriptions, c)
69
70		case <-p.quit:
71		return
72		}
73		}
74		}
75
76		// Allow mocking in tests.
77		var resolve = func(name string) (addrs []*net.SRV, md5sum string, err error) {
78		_, addrs, err = net.LookupSRV("", "", name)
79		if err != nil {
80		return addrs, "", err
81		}
82		hostports := make([]string, len(addrs))
83		for i, addr := range addrs {
84		hostports[i] = fmt.Sprintf("%s:%d", addr.Target, addr.Port)
85		}
86		sort.Sort(sort.StringSlice(hostports))
87		h := md5.New()
88		for _, hostport := range hostports {
89		fmt.Fprintf(h, hostport)
90		}
91		return addrs, fmt.Sprintf("%x", h.Sum(nil)), nil
92		}
93
94		func convert(addrs []*net.SRV, makeEndpoint func(hostport string) endpoint.Endpoint) []endpoint.Endpoint {
95		endpoints := make([]endpoint.Endpoint, len(addrs))
96		for i, addr := range addrs {
97		endpoints[i] = makeEndpoint(addr2hostport(addr))
98		}
99		return endpoints
100		}
101
102		func addr2hostport(addr *net.SRV) string {
103		return net.JoinHostPort(addr.Target, fmt.Sprintf("%d", addr.Port))
104		}

+0

-77

~~loadbalancer/dns_srv_publisher_internal_test.go~~ less more

0		package loadbalancer
1
2		import (
3		"fmt"
4		"net"
5		"testing"
6		"time"
7
8		"golang.org/x/net/context"
9
10		"github.com/go-kit/kit/endpoint"
11		)
12
13		func TestDNSSRVPublisher(t *testing.T) {
14		// Reset the vars when we're done
15		oldResolve := resolve
16		defer func() { resolve = oldResolve }()
17		oldNewTicker := newTicker
18		defer func() { newTicker = oldNewTicker }()
19
20		// Set up a fixture and swap the vars
21		a := []*net.SRV{
22		{Target: "foo", Port: 123},
23		{Target: "bar", Port: 456},
24		{Target: "baz", Port: 789},
25		}
26		ticker := make(chan time.Time)
27		resolve = func(string) ([]*net.SRV, string, error) { return a, fmt.Sprint(len(a)), nil }
28		newTicker = func(time.Duration) *time.Ticker { return &time.Ticker{C: ticker} }
29
30		// Construct endpoint
31		m := map[string]int{}
32		e := func(hostport string) endpoint.Endpoint {
33		return func(context.Context, interface{}) (interface{}, error) {
34		m[hostport]++
35		return struct{}{}, nil
36		}
37		}
38
39		// Build the publisher
40		var (
41		name = "irrelevant"
42		ttl = time.Second
43		makeEndpoint = func(hostport string) endpoint.Endpoint { return e(hostport) }
44		)
45		p := NewDNSSRVPublisher(name, ttl, makeEndpoint)
46		defer p.Stop()
47
48		// Subscribe
49		c := make(chan []endpoint.Endpoint, 1)
50		p.Subscribe(c)
51		defer p.Unsubscribe(c)
52
53		// Invoke all of the endpoints
54		for _, e := range <-c {
55		e(context.Background(), struct{}{})
56		}
57
58		// Make sure we invoked what we expected to
59		for _, addr := range a {
60		hostport := addr2hostport(addr)
61		if want, have := 1, m[hostport]; want != have {
62		t.Errorf("%q: want %d, have %d", name, want, have)
63		}
64		delete(m, hostport)
65		}
66		if want, have := 0, len(m); want != have {
67		t.Errorf("want %d, have %d", want, have)
68		}
69
70		// Reset the fixture, trigger the timer, count the endpoints
71		a = []*net.SRV{}
72		ticker <- time.Now()
73		if want, have := len(a), len(<-c); want != have {
74		t.Errorf("want %d, have %d", want, have)
75		}
76		}

+0

-41

~~loadbalancer/endpoint_cache.go~~ less more

0		package loadbalancer
1
2		import "github.com/go-kit/kit/endpoint"
3
4		type endpointCache struct {
5		requests chan []endpoint.Endpoint
6		quit chan struct{}
7		}
8
9		func newEndpointCache(p Publisher) *endpointCache {
10		c := &endpointCache{
11		requests: make(chan []endpoint.Endpoint),
12		quit: make(chan struct{}),
13		}
14		go c.loop(p)
15		return c
16		}
17
18		func (c *endpointCache) loop(p Publisher) {
19		updates := make(chan []endpoint.Endpoint, 1)
20		p.Subscribe(updates)
21		defer p.Unsubscribe(updates)
22		endpoints := <-updates
23
24		for {
25		select {
26		case endpoints = <-updates:
27		case c.requests <- endpoints:
28		case <-c.quit:
29		return
30		}
31		}
32		}
33
34		func (c *endpointCache) get() []endpoint.Endpoint {
35		return <-c.requests
36		}
37
38		func (c *endpointCache) stop() {
39		close(c.quit)
40		}

+0

-27

~~loadbalancer/endpoint_cache_internal_test.go~~ less more

0		package loadbalancer
1
2		import (
3		"testing"
4
5		"golang.org/x/net/context"
6
7		"github.com/go-kit/kit/endpoint"
8		)
9
10		func TestEndpointCache(t *testing.T) {
11		endpoints := []endpoint.Endpoint{
12		func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil },
13		func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil },
14		func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil },
15		}
16
17		p := NewStaticPublisher(endpoints)
18		defer p.Stop()
19
20		c := newEndpointCache(p)
21		defer c.stop()
22
23		if want, have := len(endpoints), len(c.get()); want != have {
24		t.Errorf("want %d, have %d", want, have)
25		}
26		}

+2

-2

loadbalancer/load_balancer.go less more

11	11	Get() (endpoint.Endpoint, error)
12	12	}
13	13
14		// ErrNoEndpointsAvailable is given by a load balancer when no endpoints are
15		// available to be returned.
	14	// ErrNoEndpointsAvailable is given by a load balancer or strategy when no
	15	// endpoints are available to be returned.
16	16	var ErrNoEndpointsAvailable = errors.New("no endpoints available")

+0

-46

~~loadbalancer/mock_publisher_test.go~~ less more

0		package loadbalancer_test
1
2		import (
3		"runtime"
4		"sync"
5
6		"github.com/go-kit/kit/endpoint"
7		)
8
9		type mockPublisher struct {
10		sync.Mutex
11		e []endpoint.Endpoint
12		s map[chan<- []endpoint.Endpoint]struct{}
13		}
14
15		func newMockPublisher(endpoints []endpoint.Endpoint) *mockPublisher {
16		return &mockPublisher{
17		e: endpoints,
18		s: map[chan<- []endpoint.Endpoint]struct{}{},
19		}
20		}
21
22		func (p *mockPublisher) replace(endpoints []endpoint.Endpoint) {
23		p.Lock()
24		defer p.Unlock()
25		p.e = endpoints
26		for s := range p.s {
27		s <- p.e
28		}
29		runtime.Gosched()
30		}
31
32		func (p *mockPublisher) Subscribe(c chan<- []endpoint.Endpoint) {
33		p.Lock()
34		defer p.Unlock()
35		p.s[c] = struct{}{}
36		c <- p.e
37		}
38
39		func (p *mockPublisher) Unsubscribe(c chan<- []endpoint.Endpoint) {
40		p.Lock()
41		defer p.Unlock()
42		delete(p.s, c)
43		}
44
45		func (p *mockPublisher) Stop() {}

+110

-0

loadbalancer/publisher/dns/srv_publisher.go less more

	0	package dns
	1
	2	import (
	3	"crypto/md5"
	4	"fmt"
	5	"net"
	6	"sort"
	7	"time"
	8
	9	"github.com/go-kit/kit/endpoint"
	10	)
	11
	12	// SRVPublisher implements Publisher.
	13	type SRVPublisher struct {
	14	subscribe chan chan<- []endpoint.Endpoint
	15	unsubscribe chan chan<- []endpoint.Endpoint
	16	quit chan struct{}
	17	}
	18
	19	// NewSRVPublisher returns a publisher that resolves the SRV name every ttl,
	20	// and yields endpoints constructed via the makeEndpoint factory.
	21	func NewSRVPublisher(name string, ttl time.Duration, makeEndpoint func(hostport string) endpoint.Endpoint) *SRVPublisher {
	22	p := &SRVPublisher{
	23	subscribe: make(chan chan<- []endpoint.Endpoint),
	24	unsubscribe: make(chan chan<- []endpoint.Endpoint),
	25	quit: make(chan struct{}),
	26	}
	27	go p.loop(name, ttl, makeEndpoint)
	28	return p
	29	}
	30
	31	// Subscribe implements Publisher.
	32	func (p *SRVPublisher) Subscribe(c chan<- []endpoint.Endpoint) {
	33	p.subscribe <- c
	34	}
	35
	36	// Unsubscribe implements Publisher.
	37	func (p *SRVPublisher) Unsubscribe(c chan<- []endpoint.Endpoint) {
	38	p.unsubscribe <- c
	39	}
	40
	41	// Stop implements Publisher.
	42	func (p *SRVPublisher) Stop() {
	43	close(p.quit)
	44	}
	45
	46	var newTicker = time.NewTicker
	47
	48	func (p *SRVPublisher) loop(name string, ttl time.Duration, makeEndpoint func(hostport string) endpoint.Endpoint) {
	49	var (
	50	subscriptions = map[chan<- []endpoint.Endpoint]struct{}{}
	51	addrs, md5, _ = resolve(name)
	52	endpoints = convert(addrs, makeEndpoint)
	53	ticker = newTicker(ttl)
	54	)
	55	defer ticker.Stop()
	56	for {
	57	select {
	58	case <-ticker.C:
	59	addrs, newmd5, err := resolve(name)
	60	if err == nil && newmd5 != md5 {
	61	endpoints = convert(addrs, makeEndpoint)
	62	for c := range subscriptions {
	63	c <- endpoints
	64	}
	65	md5 = newmd5
	66	}
	67
	68	case c := <-p.subscribe:
	69	subscriptions[c] = struct{}{}
	70	c <- endpoints
	71
	72	case c := <-p.unsubscribe:
	73	delete(subscriptions, c)
	74
	75	case <-p.quit:
	76	return
	77	}
	78	}
	79	}
	80
	81	// Allow mocking in tests.
	82	var resolve = func(name string) (addrs []*net.SRV, md5sum string, err error) {
	83	_, addrs, err = net.LookupSRV("", "", name)
	84	if err != nil {
	85	return addrs, "", err
	86	}
	87	hostports := make([]string, len(addrs))
	88	for i, addr := range addrs {
	89	hostports[i] = fmt.Sprintf("%s:%d", addr.Target, addr.Port)
	90	}
	91	sort.Sort(sort.StringSlice(hostports))
	92	h := md5.New()
	93	for _, hostport := range hostports {
	94	fmt.Fprintf(h, hostport)
	95	}
	96	return addrs, fmt.Sprintf("%x", h.Sum(nil)), nil
	97	}
	98
	99	func convert(addrs []*net.SRV, makeEndpoint func(hostport string) endpoint.Endpoint) []endpoint.Endpoint {
	100	endpoints := make([]endpoint.Endpoint, len(addrs))
	101	for i, addr := range addrs {
	102	endpoints[i] = makeEndpoint(addr2hostport(addr))
	103	}
	104	return endpoints
	105	}
	106
	107	func addr2hostport(addr *net.SRV) string {
	108	return net.JoinHostPort(addr.Target, fmt.Sprint(addr.Port))
	109	}

+77

-0

loadbalancer/publisher/dns/srv_publisher_internal_test.go less more

	0	package dns
	1
	2	import (
	3	"fmt"
	4	"net"
	5	"testing"
	6	"time"
	7
	8	"golang.org/x/net/context"
	9
	10	"github.com/go-kit/kit/endpoint"
	11	)
	12
	13	func TestDNSSRVPublisher(t *testing.T) {
	14	// Reset the vars when we're done
	15	oldResolve := resolve
	16	defer func() { resolve = oldResolve }()
	17	oldNewTicker := newTicker
	18	defer func() { newTicker = oldNewTicker }()
	19
	20	// Set up a fixture and swap the vars
	21	a := []*net.SRV{
	22	{Target: "foo", Port: 123},
	23	{Target: "bar", Port: 456},
	24	{Target: "baz", Port: 789},
	25	}
	26	ticker := make(chan time.Time)
	27	resolve = func(string) ([]*net.SRV, string, error) { return a, fmt.Sprint(len(a)), nil }
	28	newTicker = func(time.Duration) *time.Ticker { return &time.Ticker{C: ticker} }
	29
	30	// Construct endpoint
	31	m := map[string]int{}
	32	e := func(hostport string) endpoint.Endpoint {
	33	return func(context.Context, interface{}) (interface{}, error) {
	34	m[hostport]++
	35	return struct{}{}, nil
	36	}
	37	}
	38
	39	// Build the publisher
	40	var (
	41	name = "irrelevant"
	42	ttl = time.Second
	43	makeEndpoint = func(hostport string) endpoint.Endpoint { return e(hostport) }
	44	)
	45	p := NewSRVPublisher(name, ttl, makeEndpoint)
	46	defer p.Stop()
	47
	48	// Subscribe
	49	c := make(chan []endpoint.Endpoint, 1)
	50	p.Subscribe(c)
	51	defer p.Unsubscribe(c)
	52
	53	// Invoke all of the endpoints
	54	for _, e := range <-c {
	55	e(context.Background(), struct{}{})
	56	}
	57
	58	// Make sure we invoked what we expected to
	59	for _, addr := range a {
	60	hostport := addr2hostport(addr)
	61	if want, have := 1, m[hostport]; want != have {
	62	t.Errorf("%q: want %d, have %d", name, want, have)
	63	}
	64	delete(m, hostport)
	65	}
	66	if want, have := 0, len(m); want != have {
	67	t.Errorf("want %d, have %d", want, have)
	68	}
	69
	70	// Reset the fixture, trigger the timer, count the endpoints
	71	a = []*net.SRV{}
	72	ticker <- time.Now()
	73	if want, have := len(a), len(<-c); want != have {
	74	t.Errorf("want %d, have %d", want, have)
	75	}
	76	}

+10

-0

loadbalancer/publisher/publisher.go less more

	0	package publisher
	1
	2	import "github.com/go-kit/kit/endpoint"
	3
	4	// Publisher produces endpoints.
	5	type Publisher interface {
	6	Subscribe(chan<- []endpoint.Endpoint)
	7	Unsubscribe(chan<- []endpoint.Endpoint)
	8	Stop()
	9	}

+51

-0

loadbalancer/publisher/static/publisher.go less more

	0	package static
	1
	2	import (
	3	"sync"
	4
	5	"github.com/go-kit/kit/endpoint"
	6	)
	7
	8	// Publisher holds a static set of endpoints.
	9	type Publisher struct {
	10	mu sync.Mutex
	11	current []endpoint.Endpoint
	12	subscribers map[chan<- []endpoint.Endpoint]struct{}
	13	}
	14
	15	// NewPublisher returns a publisher that yields a static set of endpoints,
	16	// which can be completely replaced.
	17	func NewPublisher(endpoints []endpoint.Endpoint) *Publisher {
	18	return &Publisher{
	19	current: endpoints,
	20	subscribers: map[chan<- []endpoint.Endpoint]struct{}{},
	21	}
	22	}
	23
	24	// Subscribe implements Publisher.
	25	func (p *Publisher) Subscribe(c chan<- []endpoint.Endpoint) {
	26	p.mu.Lock()
	27	defer p.mu.Unlock()
	28	p.subscribers[c] = struct{}{}
	29	c <- p.current
	30	}
	31
	32	// Unsubscribe implements Publisher.
	33	func (p *Publisher) Unsubscribe(c chan<- []endpoint.Endpoint) {
	34	p.mu.Lock()
	35	defer p.mu.Unlock()
	36	delete(p.subscribers, c)
	37	}
	38
	39	// Stop implements Publisher, but is a no-op.
	40	func (p *Publisher) Stop() {}
	41
	42	// Replace replaces the endpoints and notifies all subscribers.
	43	func (p *Publisher) Replace(endpoints []endpoint.Endpoint) {
	44	p.mu.Lock()
	45	defer p.mu.Unlock()
	46	p.current = endpoints
	47	for c := range p.subscribers {
	48	c <- p.current
	49	}
	50	}

+30

-0

loadbalancer/publisher/static/publisher_test.go less more

	0	package static_test
	1
	2	import (
	3	"testing"
	4
	5	"golang.org/x/net/context"
	6
	7	"github.com/go-kit/kit/endpoint"
	8	"github.com/go-kit/kit/loadbalancer/publisher/static"
	9	)
	10
	11	func TestStaticPublisher(t *testing.T) {
	12	endpoints := []endpoint.Endpoint{
	13	func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil },
	14	}
	15	p := static.NewPublisher(endpoints)
	16	defer p.Stop()
	17
	18	c := make(chan []endpoint.Endpoint, 1)
	19	p.Subscribe(c)
	20	if want, have := len(endpoints), len(<-c); want != have {
	21	t.Errorf("want %d, have %d", want, have)
	22	}
	23
	24	endpoints = []endpoint.Endpoint{}
	25	p.Replace(endpoints)
	26	if want, have := len(endpoints), len(<-c); want != have {
	27	t.Errorf("want %d, have %d", want, have)
	28	}
	29	}

+0

-10

~~loadbalancer/publisher.go~~ less more

0		package loadbalancer
1
2		import "github.com/go-kit/kit/endpoint"
3
4		// Publisher produces endpoints.
5		type Publisher interface {
6		Subscribe(chan<- []endpoint.Endpoint)
7		Unsubscribe(chan<- []endpoint.Endpoint)
8		Stop()
9		}

+0

-24

~~loadbalancer/random.go~~ less more

0		package loadbalancer
1
2		import (
3		"math/rand"
4
5		"github.com/go-kit/kit/endpoint"
6		)
7
8		// Random returns a load balancer that yields random endpoints.
9		func Random(p Publisher) LoadBalancer {
10		return random{newCache(p)}
11		}
12
13		type random struct{ *cache }
14
15		func (r random) Count() int { return r.cache.count() }
16
17		func (r random) Get() (endpoint.Endpoint, error) {
18		endpoints := r.cache.get()
19		if len(endpoints) <= 0 {
20		return nil, ErrNoEndpointsAvailable
21		}
22		return endpoints[rand.Intn(len(endpoints))], nil
23		}

+0

-42

~~loadbalancer/random_test.go~~ less more

0		package loadbalancer_test
1
2		import (
3		"math"
4		"testing"
5
6		"github.com/go-kit/kit/endpoint"
7		"github.com/go-kit/kit/loadbalancer"
8		"golang.org/x/net/context"
9		)
10
11		func TestRandom(t *testing.T) {
12		p := loadbalancer.NewStaticPublisher([]endpoint.Endpoint{})
13		defer p.Stop()
14
15		lb := loadbalancer.Random(p)
16		if _, err := lb.Get(); err == nil {
17		t.Error("want error, got none")
18		}
19
20		counts := []int{0, 0, 0}
21		p.Replace([]endpoint.Endpoint{
22		func(context.Context, interface{}) (interface{}, error) { counts[0]++; return struct{}{}, nil },
23		func(context.Context, interface{}) (interface{}, error) { counts[1]++; return struct{}{}, nil },
24		func(context.Context, interface{}) (interface{}, error) { counts[2]++; return struct{}{}, nil },
25		})
26		assertLoadBalancerNotEmpty(t, lb)
27
28		n := 10000
29		for i := 0; i < n; i++ {
30		e, _ := lb.Get()
31		e(context.Background(), struct{}{})
32		}
33
34		want := float64(n) / float64(len(counts))
35		tolerance := (want / 100.0) * 5 // 5%
36		for _, have := range counts {
37		if math.Abs(want-float64(have)) > tolerance {
38		t.Errorf("want %.0f, have %d", want, have)
39		}
40		}
41		}

+9

-7

loadbalancer/retry_test.go less more

1	1
2	2	import (
3	3	"errors"
	4	"testing"
4	5	"time"
	6
	7	"golang.org/x/net/context"
5	8
6	9	"github.com/go-kit/kit/endpoint"
7	10	"github.com/go-kit/kit/loadbalancer"
8		"golang.org/x/net/context"
9
10		"testing"
	11	"github.com/go-kit/kit/loadbalancer/publisher/static"
	12	"github.com/go-kit/kit/loadbalancer/strategy"
11	13	)
12	14
13	15	func TestRetryMax(t *testing.T) {
14	16	var (
15	17	endpoints = []endpoint.Endpoint{}
16		p = loadbalancer.NewStaticPublisher(endpoints)
17		lb = loadbalancer.RoundRobin(p)
	18	p = static.NewPublisher(endpoints)
	19	lb = strategy.RoundRobin(p)
18	20	)
19	21
20	22	if _, err := loadbalancer.Retry(999, time.Second, lb)(context.Background(), struct{}{}); err == nil {

27	29	func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil /* OK */ },
28	30	}
29	31	p.Replace(endpoints)
30		assertLoadBalancerNotEmpty(t, lb)
	32	time.Sleep(10 * time.Millisecond) //assertLoadBalancerNotEmpty(t, lb) // TODO
31	33
32	34	if _, err := loadbalancer.Retry(len(endpoints)-1, time.Second, lb)(context.Background(), struct{}{}); err == nil {
33	35	t.Errorf("expected error, got none")

43	45	step = make(chan struct{})
44	46	e = func(context.Context, interface{}) (interface{}, error) { <-step; return struct{}{}, nil }
45	47	timeout = time.Millisecond
46		retry = loadbalancer.Retry(999, timeout, loadbalancer.RoundRobin(loadbalancer.NewStaticPublisher([]endpoint.Endpoint{e})))
	48	retry = loadbalancer.Retry(999, timeout, strategy.RoundRobin(static.NewPublisher([]endpoint.Endpoint{e})))
47	49	errs = make(chan error)
48	50	invoke = func() { _, err := retry(context.Background(), struct{}{}); errs <- err }
49	51	)

+0

-34

~~loadbalancer/round_robin.go~~ less more

0		package loadbalancer
1
2		import (
3		"sync/atomic"
4
5		"github.com/go-kit/kit/endpoint"
6		)
7
8		// RoundRobin returns a load balancer that yields endpoints in sequence.
9		func RoundRobin(p Publisher) LoadBalancer {
10		return &roundRobin{newCache(p), 0}
11		}
12
13		type roundRobin struct {
14		*cache
15		uint64
16		}
17
18		func (r *roundRobin) Count() int { return r.cache.count() }
19
20		func (r *roundRobin) Get() (endpoint.Endpoint, error) {
21		endpoints := r.cache.get()
22		if len(endpoints) <= 0 {
23		return nil, ErrNoEndpointsAvailable
24		}
25		var old uint64
26		for {
27		old = atomic.LoadUint64(&r.uint64)
28		if atomic.CompareAndSwapUint64(&r.uint64, old, old+1) {
29		break
30		}
31		}
32		return endpoints[old%uint64(len(endpoints))], nil
33		}

+0

-44

~~loadbalancer/round_robin_test.go~~ less more

0		package loadbalancer_test
1
2		import (
3		"reflect"
4		"testing"
5
6		"github.com/go-kit/kit/endpoint"
7		"github.com/go-kit/kit/loadbalancer"
8		"golang.org/x/net/context"
9		)
10
11		func TestRoundRobin(t *testing.T) {
12		p := loadbalancer.NewStaticPublisher([]endpoint.Endpoint{})
13		defer p.Stop()
14
15		lb := loadbalancer.RoundRobin(p)
16		if _, err := lb.Get(); err == nil {
17		t.Error("want error, got none")
18		}
19
20		counts := []int{0, 0, 0}
21		p.Replace([]endpoint.Endpoint{
22		func(context.Context, interface{}) (interface{}, error) { counts[0]++; return struct{}{}, nil },
23		func(context.Context, interface{}) (interface{}, error) { counts[1]++; return struct{}{}, nil },
24		func(context.Context, interface{}) (interface{}, error) { counts[2]++; return struct{}{}, nil },
25		})
26		assertLoadBalancerNotEmpty(t, lb)
27
28		for i, want := range [][]int{
29		{1, 0, 0},
30		{1, 1, 0},
31		{1, 1, 1},
32		{2, 1, 1},
33		{2, 2, 1},
34		{2, 2, 2},
35		{3, 2, 2},
36		} {
37		e, _ := lb.Get()
38		e(context.Background(), struct{}{})
39		if have := counts; !reflect.DeepEqual(want, have) {
40		t.Errorf("%d: want %v, have %v", i+1, want, have)
41		}
42		}
43		}

+0

-51

~~loadbalancer/static_publisher.go~~ less more

0		package loadbalancer
1
2		import (
3		"sync"
4
5		"github.com/go-kit/kit/endpoint"
6		)
7
8		// NewStaticPublisher returns a publisher that yields a static set of
9		// endpoints, which can be completely replaced.
10		func NewStaticPublisher(endpoints []endpoint.Endpoint) *StaticPublisher {
11		return &StaticPublisher{
12		current: endpoints,
13		subscribers: map[chan<- []endpoint.Endpoint]struct{}{},
14		}
15		}
16
17		// StaticPublisher holds a static set of endpoints.
18		type StaticPublisher struct {
19		mu sync.Mutex
20		current []endpoint.Endpoint
21		subscribers map[chan<- []endpoint.Endpoint]struct{}
22		}
23
24		// Subscribe implements Publisher.
25		func (p *StaticPublisher) Subscribe(c chan<- []endpoint.Endpoint) {
26		p.mu.Lock()
27		defer p.mu.Unlock()
28		p.subscribers[c] = struct{}{}
29		c <- p.current
30		}
31
32		// Unsubscribe implements Publisher.
33		func (p *StaticPublisher) Unsubscribe(c chan<- []endpoint.Endpoint) {
34		p.mu.Lock()
35		defer p.mu.Unlock()
36		delete(p.subscribers, c)
37		}
38
39		// Stop implements Publisher, but is a no-op.
40		func (p *StaticPublisher) Stop() {}
41
42		// Replace replaces the endpoints and notifies all subscribers.
43		func (p *StaticPublisher) Replace(endpoints []endpoint.Endpoint) {
44		p.mu.Lock()
45		defer p.mu.Unlock()
46		p.current = endpoints
47		for c := range p.subscribers {
48		c <- p.current
49		}
50		}

+0

-30

~~loadbalancer/static_publisher_test.go~~ less more

0		package loadbalancer_test
1
2		import (
3		"testing"
4
5		"golang.org/x/net/context"
6
7		"github.com/go-kit/kit/endpoint"
8		"github.com/go-kit/kit/loadbalancer"
9		)
10
11		func TestStaticPublisher(t *testing.T) {
12		endpoints := []endpoint.Endpoint{
13		func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil },
14		}
15		p := loadbalancer.NewStaticPublisher(endpoints)
16		defer p.Stop()
17
18		c := make(chan []endpoint.Endpoint, 1)
19		p.Subscribe(c)
20		if want, have := len(endpoints), len(<-c); want != have {
21		t.Errorf("want %d, have %d", want, have)
22		}
23
24		endpoints = []endpoint.Endpoint{}
25		p.Replace(endpoints)
26		if want, have := len(endpoints), len(<-c); want != have {
27		t.Errorf("want %d, have %d", want, have)
28		}
29		}

+0

-17

~~loadbalancer/strategies.go~~ less more

0		package loadbalancer
1
2		import (
3		"errors"
4
5		"github.com/go-kit/kit/endpoint"
6		)
7
8		// Strategy yields endpoints to consumers according to some algorithm.
9		type Strategy interface {
10		Next() (endpoint.Endpoint, error)
11		Stop()
12		}
13
14		// ErrNoEndpoints is returned by a strategy when there are no endpoints
15		// available.
16		var ErrNoEndpoints = errors.New("no endpoints available")

+50

-0

loadbalancer/strategy/cache.go less more

	0	package strategy
	1
	2	import (
	3	"github.com/go-kit/kit/endpoint"
	4	"github.com/go-kit/kit/loadbalancer/publisher"
	5	)
	6
	7	type cache struct {
	8	req chan []endpoint.Endpoint
	9	cnt chan int
	10	quit chan struct{}
	11	}
	12
	13	func newCache(p publisher.Publisher) *cache {
	14	c := &cache{
	15	req: make(chan []endpoint.Endpoint),
	16	cnt: make(chan int),
	17	quit: make(chan struct{}),
	18	}
	19	go c.loop(p)
	20	return c
	21	}
	22
	23	func (c *cache) loop(p publisher.Publisher) {
	24	e := make(chan []endpoint.Endpoint, 1)
	25	p.Subscribe(e)
	26	defer p.Unsubscribe(e)
	27	endpoints := <-e
	28	for {
	29	select {
	30	case endpoints = <-e:
	31	case c.cnt <- len(endpoints):
	32	case c.req <- endpoints:
	33	case <-c.quit:
	34	return
	35	}
	36	}
	37	}
	38
	39	func (c *cache) count() int {
	40	return <-c.cnt
	41	}
	42
	43	func (c *cache) get() []endpoint.Endpoint {
	44	return <-c.req
	45	}
	46
	47	func (c *cache) stop() {
	48	close(c.quit)
	49	}

+34

-0

loadbalancer/strategy/cache_internal_test.go less more

	0	package strategy
	1
	2	import (
	3	"runtime"
	4	"testing"
	5
	6	"golang.org/x/net/context"
	7
	8	"github.com/go-kit/kit/endpoint"
	9	"github.com/go-kit/kit/loadbalancer/publisher/static"
	10	)
	11
	12	func TestCache(t *testing.T) {
	13	e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
	14	endpoints := []endpoint.Endpoint{e}
	15
	16	p := static.NewPublisher(endpoints)
	17	defer p.Stop()
	18
	19	c := newCache(p)
	20	defer c.stop()
	21
	22	for _, n := range []int{2, 10, 0} {
	23	endpoints = make([]endpoint.Endpoint, n)
	24	for i := 0; i < n; i++ {
	25	endpoints[i] = e
	26	}
	27	p.Replace(endpoints)
	28	runtime.Gosched()
	29	if want, have := len(endpoints), len(c.get()); want != have {
	30	t.Errorf("want %d, have %d", want, have)
	31	}
	32	}
	33	}

+26

-0

loadbalancer/strategy/random.go less more

	0	package strategy
	1
	2	import (
	3	"math/rand"
	4
	5	"github.com/go-kit/kit/endpoint"
	6	"github.com/go-kit/kit/loadbalancer"
	7	"github.com/go-kit/kit/loadbalancer/publisher"
	8	)
	9
	10	// Random returns a load balancer that yields random endpoints.
	11	func Random(p publisher.Publisher) loadbalancer.LoadBalancer {
	12	return random{newCache(p)}
	13	}
	14
	15	type random struct{ *cache }
	16
	17	func (r random) Count() int { return r.cache.count() }
	18
	19	func (r random) Get() (endpoint.Endpoint, error) {
	20	endpoints := r.cache.get()
	21	if len(endpoints) <= 0 {
	22	return nil, loadbalancer.ErrNoEndpointsAvailable
	23	}
	24	return endpoints[rand.Intn(len(endpoints))], nil
	25	}

+43

-0

loadbalancer/strategy/random_test.go less more

	0	package strategy_test
	1
	2	import (
	3	"math"
	4	"testing"
	5
	6	"github.com/go-kit/kit/endpoint"
	7	"github.com/go-kit/kit/loadbalancer/publisher/static"
	8	"github.com/go-kit/kit/loadbalancer/strategy"
	9	"golang.org/x/net/context"
	10	)
	11
	12	func TestRandom(t *testing.T) {
	13	p := static.NewPublisher([]endpoint.Endpoint{})
	14	defer p.Stop()
	15
	16	lb := strategy.Random(p)
	17	if _, err := lb.Get(); err == nil {
	18	t.Error("want error, got none")
	19	}
	20
	21	counts := []int{0, 0, 0}
	22	p.Replace([]endpoint.Endpoint{
	23	func(context.Context, interface{}) (interface{}, error) { counts[0]++; return struct{}{}, nil },
	24	func(context.Context, interface{}) (interface{}, error) { counts[1]++; return struct{}{}, nil },
	25	func(context.Context, interface{}) (interface{}, error) { counts[2]++; return struct{}{}, nil },
	26	})
	27	assertLoadBalancerNotEmpty(t, lb)
	28
	29	n := 10000
	30	for i := 0; i < n; i++ {
	31	e, _ := lb.Get()
	32	e(context.Background(), struct{}{})
	33	}
	34
	35	want := float64(n) / float64(len(counts))
	36	tolerance := (want / 100.0) * 5 // 5%
	37	for _, have := range counts {
	38	if math.Abs(want-float64(have)) > tolerance {
	39	t.Errorf("want %.0f, have %d", want, have)
	40	}
	41	}
	42	}

+36

-0

loadbalancer/strategy/round_robin.go less more

	0	package strategy
	1
	2	import (
	3	"sync/atomic"
	4
	5	"github.com/go-kit/kit/endpoint"
	6	"github.com/go-kit/kit/loadbalancer"
	7	"github.com/go-kit/kit/loadbalancer/publisher"
	8	)
	9
	10	// RoundRobin returns a load balancer that yields endpoints in sequence.
	11	func RoundRobin(p publisher.Publisher) loadbalancer.LoadBalancer {
	12	return &roundRobin{newCache(p), 0}
	13	}
	14
	15	type roundRobin struct {
	16	*cache
	17	uint64
	18	}
	19
	20	func (r *roundRobin) Count() int { return r.cache.count() }
	21
	22	func (r *roundRobin) Get() (endpoint.Endpoint, error) {
	23	endpoints := r.cache.get()
	24	if len(endpoints) <= 0 {
	25	return nil, loadbalancer.ErrNoEndpointsAvailable
	26	}
	27	var old uint64
	28	for {
	29	old = atomic.LoadUint64(&r.uint64)
	30	if atomic.CompareAndSwapUint64(&r.uint64, old, old+1) {
	31	break
	32	}
	33	}
	34	return endpoints[old%uint64(len(endpoints))], nil
	35	}

+46

-0

loadbalancer/strategy/round_robin_test.go less more

	0	package strategy_test
	1
	2	import (
	3	"reflect"
	4	"testing"
	5
	6	"golang.org/x/net/context"
	7
	8	"github.com/go-kit/kit/endpoint"
	9	"github.com/go-kit/kit/loadbalancer/publisher/static"
	10	"github.com/go-kit/kit/loadbalancer/strategy"
	11	)
	12
	13	func TestRoundRobin(t *testing.T) {
	14	p := static.NewPublisher([]endpoint.Endpoint{})
	15	defer p.Stop()
	16
	17	lb := strategy.RoundRobin(p)
	18	if _, err := lb.Get(); err == nil {
	19	t.Error("want error, got none")
	20	}
	21
	22	counts := []int{0, 0, 0}
	23	p.Replace([]endpoint.Endpoint{
	24	func(context.Context, interface{}) (interface{}, error) { counts[0]++; return struct{}{}, nil },
	25	func(context.Context, interface{}) (interface{}, error) { counts[1]++; return struct{}{}, nil },
	26	func(context.Context, interface{}) (interface{}, error) { counts[2]++; return struct{}{}, nil },
	27	})
	28	assertLoadBalancerNotEmpty(t, lb)
	29
	30	for i, want := range [][]int{
	31	{1, 0, 0},
	32	{1, 1, 0},
	33	{1, 1, 1},
	34	{2, 1, 1},
	35	{2, 2, 1},
	36	{2, 2, 2},
	37	{3, 2, 2},
	38	} {
	39	e, _ := lb.Get()
	40	e(context.Background(), struct{}{})
	41	if have := counts; !reflect.DeepEqual(want, have) {
	42	t.Errorf("%d: want %v, have %v", i+1, want, have)
	43	}
	44	}
	45	}

+9

-0

loadbalancer/strategy/strategy.go less more

	0	package strategy
	1
	2	import "github.com/go-kit/kit/endpoint"
	3
	4	// Strategy yields endpoints to consumers according to some algorithm.
	5	type Strategy interface {
	6	Next() (endpoint.Endpoint, error)
	7	Stop()
	8	}

+35

-0

loadbalancer/strategy/util_test.go less more

	0	package strategy_test
	1
	2	import (
	3	"fmt"
	4	"testing"
	5	"time"
	6
	7	"github.com/go-kit/kit/loadbalancer"
	8	)
	9
	10	func assertLoadBalancerNotEmpty(t *testing.T, lb loadbalancer.LoadBalancer) {
	11	if err := within(10*time.Millisecond, func() bool {
	12	return lb.Count() > 0
	13	}); err != nil {
	14	t.Fatal("Publisher never updated endpoints")
	15	}
	16	}
	17
	18	func within(d time.Duration, f func() bool) error {
	19	var (
	20	deadline = time.After(d)
	21	ticker = time.NewTicker(d / 10)
	22	)
	23	defer ticker.Stop()
	24	for {
	25	select {
	26	case <-ticker.C:
	27	if f() {
	28	return nil
	29	}
	30	case <-deadline:
	31	return fmt.Errorf("deadline exceeded")
	32	}
	33	}
	34	}

+0

-35

~~loadbalancer/util_test.go~~ less more

0		package loadbalancer_test
1
2		import (
3		"fmt"
4		"testing"
5		"time"
6
7		"github.com/go-kit/kit/loadbalancer"
8		)
9
10		func assertLoadBalancerNotEmpty(t *testing.T, lb loadbalancer.LoadBalancer) {
11		if err := within(10*time.Millisecond, func() bool {
12		return lb.Count() > 0
13		}); err != nil {
14		t.Fatal("Publisher never updated endpoints")
15		}
16		}
17
18		func within(d time.Duration, f func() bool) error {
19		var (
20		deadline = time.After(d)
21		ticker = time.NewTicker(d / 10)
22		)
23		defer ticker.Stop()
24		for {
25		select {
26		case <-ticker.C:
27		if f() {
28		return nil
29		}
30		case <-deadline:
31		return fmt.Errorf("deadline exceeded")
32		}
33		}
34		}