Commit 8357da3f3de51c1e663589dae756616bb42009aa - golang-github-go-kit-kit

+0

-44

~~loadbalancer/README.md~~ less more

0		# package loadbalancer
1
2		`package loadbalancer` provides a client-side load balancer abstraction.
3
4		A publisher is responsible for emitting the most recent set of endpoints for a
5		single logical service. Publishers exist for static endpoints, and endpoints
6		discovered via periodic DNS SRV lookups on a single logical name. Consul and
7		etcd publishers are planned.
8
9		Different load balancing strategies are implemented on top of publishers. Go
10		kit currently provides random and round-robin semantics. Smarter behaviors,
11		e.g. load balancing based on underlying endpoint priority/weight, is planned.
12
13		## Rationale
14
15		TODO
16
17		## Usage
18
19		In your client, define a publisher, wrap it with a balancing strategy, and pass
20		it to a retry strategy, which returns an endpoint. Use that endpoint to make
21		requests, or wrap it with other value-add middleware.
22
23		```go
24		func main() {
25		var (
26		fooPublisher = loadbalancer.NewDNSSRVPublisher("foo.mynet.local", 5*time.Second, makeEndpoint)
27		fooBalancer = loadbalancer.RoundRobin(fooPublisher)
28		fooEndpoint = loadbalancer.Retry(3, time.Second, fooBalancer)
29		)
30		http.HandleFunc("/", handle(fooEndpoint))
31		log.Fatal(http.ListenAndServe(":8080", nil))
32		}
33
34		func makeEndpoint(hostport string) endpoint.Endpoint {
35		// Convert a host:port to a endpoint via your defined transport.
36		}
37
38		func handle(foo endpoint.Endpoint) http.HandlerFunc {
39		return func(w http.ResponseWriter, r *http.Request) {
40		// foo is usable as a load-balanced remote endpoint.
41		}
42		}
43		```

+122

-0

loadbalancer/dnssrv/publisher.go less more

	0	package dnssrv
	1
	2	import (
	3	"crypto/md5"
	4	"fmt"
	5	"net"
	6	"sort"
	7	"time"
	8
	9	"github.com/go-kit/kit/endpoint"
	10	"github.com/go-kit/kit/loadbalancer"
	11	"github.com/go-kit/kit/log"
	12	)
	13
	14	// Publisher yields endpoints taken from the named DNS SRV record. The name is
	15	// resolved on a fixed schedule. Priorities and weights are ignored.
	16	type Publisher struct {
	17	name string
	18	ttl time.Duration
	19	factory loadbalancer.Factory
	20	logger log.Logger
	21	endpoints chan []endpoint.Endpoint
	22	quit chan struct{}
	23	}
	24
	25	// NewPublisher returns a DNS SRV publisher. The name is resolved
	26	// synchronously as part of construction; if that resolution fails, the
	27	// constructor will return an error. The factory is used to convert a
	28	// host:port to a usable endpoint. The logger is used to report DNS and
	29	// factory errors.
	30	func NewPublisher(name string, ttl time.Duration, f loadbalancer.Factory, logger log.Logger) (*Publisher, error) {
	31	logger = log.NewContext(logger).With("component", "DNS SRV Publisher")
	32	addrs, md5, err := resolve(name)
	33	if err != nil {
	34	return nil, err
	35	}
	36	p := &Publisher{
	37	name: name,
	38	ttl: ttl,
	39	factory: f,
	40	logger: logger,
	41	endpoints: make(chan []endpoint.Endpoint),
	42	quit: make(chan struct{}),
	43	}
	44	go p.loop(lift(addrs, f, logger), md5)
	45	return p, nil
	46	}
	47
	48	// Stop terminates the publisher.
	49	func (p *Publisher) Stop() {
	50	close(p.quit)
	51	}
	52
	53	func (p *Publisher) loop(endpoints []endpoint.Endpoint, md5 string) {
	54	t := newTicker(p.ttl)
	55	defer t.Stop()
	56	for {
	57	select {
	58	case p.endpoints <- endpoints:
	59
	60	case <-t.C:
	61	// TODO should we do this out-of-band?
	62	addrs, newmd5, err := resolve(p.name)
	63	if err != nil {
	64	p.logger.Log("name", p.name, "err", err)
	65	continue // don't replace good endpoints with bad ones
	66	}
	67	if newmd5 == md5 {
	68	continue // no change
	69	}
	70	endpoints = lift(addrs, p.factory, p.logger)
	71	md5 = newmd5
	72
	73	case <-p.quit:
	74	return
	75	}
	76	}
	77	}
	78
	79	// Endpoints implements the Publisher interface.
	80	func (p *Publisher) Endpoints() ([]endpoint.Endpoint, error) {
	81	return <-p.endpoints, nil
	82	}
	83
	84	var (
	85	lookupSRV = net.LookupSRV
	86	newTicker = time.NewTicker
	87	)
	88
	89	func resolve(name string) (addrs []*net.SRV, md5sum string, err error) {
	90	_, addrs, err = lookupSRV("", "", name)
	91	if err != nil {
	92	return addrs, "", err
	93	}
	94	hostports := make([]string, len(addrs))
	95	for i, addr := range addrs {
	96	hostports[i] = fmt.Sprintf("%s:%d", addr.Target, addr.Port)
	97	}
	98	sort.Sort(sort.StringSlice(hostports))
	99	h := md5.New()
	100	for _, hostport := range hostports {
	101	fmt.Fprintf(h, hostport)
	102	}
	103	return addrs, fmt.Sprintf("%x", h.Sum(nil)), nil
	104	}
	105
	106	func lift(addrs []*net.SRV, f loadbalancer.Factory, logger log.Logger) []endpoint.Endpoint {
	107	endpoints := make([]endpoint.Endpoint, 0, len(addrs))
	108	for _, addr := range addrs {
	109	endpoint, err := f(addr2instance(addr))
	110	if err != nil {
	111	logger.Log("instance", addr2instance(addr), "err", err)
	112	continue
	113	}
	114	endpoints = append(endpoints, endpoint)
	115	}
	116	return endpoints
	117	}
	118
	119	func addr2instance(addr *net.SRV) string {
	120	return net.JoinHostPort(addr.Target, fmt.Sprint(addr.Port))
	121	}

+145

-0

loadbalancer/dnssrv/publisher_internal_test.go less more

	0	package dnssrv
	1
	2	import (
	3	"errors"
	4	"net"
	5	"sync/atomic"
	6	"testing"
	7	"time"
	8
	9	"golang.org/x/net/context"
	10
	11	"github.com/go-kit/kit/endpoint"
	12	"github.com/go-kit/kit/log"
	13	)
	14
	15	func TestPublisher(t *testing.T) {
	16	var (
	17	target = "my-target"
	18	port = uint16(1234)
	19	addr = &net.SRV{Target: target, Port: port}
	20	addrs = []*net.SRV{addr}
	21	name = "my-name"
	22	ttl = time.Second
	23	logger = log.NewNopLogger()
	24	e = func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
	25	)
	26
	27	oldLookup := lookupSRV
	28	defer func() { lookupSRV = oldLookup }()
	29	lookupSRV = mockLookupSRV(addrs, nil, nil)
	30
	31	factory := func(instance string) (endpoint.Endpoint, error) {
	32	if want, have := addr2instance(addr), instance; want != have {
	33	t.Errorf("want %q, have %q", want, have)
	34	}
	35	return e, nil
	36	}
	37
	38	p, err := NewPublisher(name, ttl, factory, logger)
	39	if err != nil {
	40	t.Fatal(err)
	41	}
	42	defer p.Stop()
	43
	44	if _, err := p.Endpoints(); err != nil {
	45	t.Fatal(err)
	46	}
	47	}
	48
	49	func TestBadLookup(t *testing.T) {
	50	oldLookup := lookupSRV
	51	defer func() { lookupSRV = oldLookup }()
	52	lookupSRV = mockLookupSRV([]*net.SRV{}, errors.New("kaboom"), nil)
	53
	54	var (
	55	name = "some-name"
	56	ttl = time.Second
	57	factory = func(string) (endpoint.Endpoint, error) { return nil, errors.New("unreachable") }
	58	logger = log.NewNopLogger()
	59	)
	60
	61	if _, err := NewPublisher(name, ttl, factory, logger); err == nil {
	62	t.Fatal("wanted error, got none")
	63	}
	64	}
	65
	66	func TestBadFactory(t *testing.T) {
	67	var (
	68	addr = &net.SRV{Target: "foo", Port: 1234}
	69	addrs = []*net.SRV{addr}
	70	name = "some-name"
	71	ttl = time.Second
	72	factory = func(string) (endpoint.Endpoint, error) { return nil, errors.New("kaboom") }
	73	logger = log.NewNopLogger()
	74	)
	75
	76	oldLookup := lookupSRV
	77	defer func() { lookupSRV = oldLookup }()
	78	lookupSRV = mockLookupSRV(addrs, nil, nil)
	79
	80	p, err := NewPublisher(name, ttl, factory, logger)
	81	if err != nil {
	82	t.Fatal(err)
	83	}
	84	defer p.Stop()
	85
	86	endpoints, err := p.Endpoints()
	87	if err != nil {
	88	t.Fatal(err)
	89	}
	90	if want, have := 0, len(endpoints); want != have {
	91	t.Errorf("want %q, have %q", want, have)
	92	}
	93	}
	94
	95	func TestRefreshWithChange(t *testing.T) {
	96	t.Skip("TODO")
	97	}
	98
	99	func TestRefreshNoChange(t *testing.T) {
	100	var (
	101	tick = make(chan time.Time)
	102	target = "my-target"
	103	port = uint16(5678)
	104	addr = &net.SRV{Target: target, Port: port}
	105	addrs = []*net.SRV{addr}
	106	name = "my-name"
	107	ttl = time.Second
	108	factory = func(string) (endpoint.Endpoint, error) { return nil, errors.New("kaboom") }
	109	logger = log.NewNopLogger()
	110	)
	111
	112	oldTicker := newTicker
	113	defer func() { newTicker = oldTicker }()
	114	newTicker = func(time.Duration) *time.Ticker { return &time.Ticker{C: tick} }
	115
	116	var resolves uint64
	117	oldLookup := lookupSRV
	118	defer func() { lookupSRV = oldLookup }()
	119	lookupSRV = mockLookupSRV(addrs, nil, &resolves)
	120
	121	p, err := NewPublisher(name, ttl, factory, logger)
	122	if err != nil {
	123	t.Fatal(err)
	124	}
	125	defer p.Stop()
	126
	127	tick <- time.Now()
	128	if want, have := uint64(2), resolves; want != have {
	129	t.Errorf("want %d, have %d", want, have)
	130	}
	131	}
	132
	133	func TestRefreshResolveError(t *testing.T) {
	134	t.Skip("TODO")
	135	}
	136
	137	func mockLookupSRV(addrs []net.SRV, err error, count uint64) func(service, proto, name string) (string, []*net.SRV, error) {
	138	return func(service, proto, name string) (string, []*net.SRV, error) {
	139	if count != nil {
	140	atomic.AddUint64(count, 1)
	141	}
	142	return "", addrs, err
	143	}
	144	}

+10

-0

loadbalancer/factory.go less more

	0	package loadbalancer
	1
	2	import "github.com/go-kit/kit/endpoint"
	3
	4	// Factory is a function that converts an instance string, e.g. a host:port,
	5	// to a usable endpoint. Factories are used by load balancers to lift
	6	// instances returned by Publishers into endpoints. Users are expected to
	7	// provide their own factory functions that assume specific transports, or can
	8	// deduce transports by parsing the instance string.
	9	type Factory func(instance string) (endpoint.Endpoint, error)

+0

-17

~~loadbalancer/load_balancer.go~~ less more

0		package loadbalancer
1
2		import (
3		"errors"
4
5		"github.com/go-kit/kit/endpoint"
6		)
7
8		// LoadBalancer yields endpoints one-by-one.
9		type LoadBalancer interface {
10		Count() int
11		Get() (endpoint.Endpoint, error)
12		}
13
14		// ErrNoEndpointsAvailable is given by a load balancer or strategy when no
15		// endpoints are available to be returned.
16		var ErrNoEndpointsAvailable = errors.New("no endpoints available")

+8

-0

loadbalancer/loadbalancer.go less more

	0	package loadbalancer
	1
	2	import "errors"
	3
	4	// ErrNoEndpoints is returned when a load balancer (or one of its components)
	5	// has no endpoints to return. In a request lifecycle, this is usually a fatal
	6	// error.
	7	var ErrNoEndpoints = errors.New("no endpoints available")

+0

-110

~~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		}

+0

-77

~~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		}

+0

-10

~~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		}

+0

-51

~~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		}

+0

-30

~~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		}

+10

-0

loadbalancer/publisher.go less more

	0	package loadbalancer
	1
	2	import "github.com/go-kit/kit/endpoint"
	3
	4	// Publisher describes something that provides a set of identical endpoints.
	5	// Different publisher implementations exist for different kinds of service
	6	// discovery systems.
	7	type Publisher interface {
	8	Endpoints() ([]endpoint.Endpoint, error)
	9	}

+34

-0

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 is a completely stateless load balancer that chooses a random
	9	// endpoint to return each time.
	10	type Random struct {
	11	p Publisher
	12	r *rand.Rand
	13	}
	14
	15	// NewRandom returns a new Random load balancer.
	16	func NewRandom(p Publisher, seed int64) *Random {
	17	return &Random{
	18	p: p,
	19	r: rand.New(rand.NewSource(seed)),
	20	}
	21	}
	22
	23	// Endpoint implements the LoadBalancer interface.
	24	func (r *Random) Endpoint() (endpoint.Endpoint, error) {
	25	endpoints, err := r.p.Endpoints()
	26	if err != nil {
	27	return nil, err
	28	}
	29	if len(endpoints) <= 0 {
	30	return nil, ErrNoEndpoints
	31	}
	32	return endpoints[r.r.Intn(len(endpoints))], nil
	33	}

+58

-0

loadbalancer/random_test.go less more

	0	package loadbalancer_test
	1
	2	import (
	3	"math"
	4	"testing"
	5
	6	"golang.org/x/net/context"
	7
	8	"github.com/go-kit/kit/endpoint"
	9	"github.com/go-kit/kit/loadbalancer"
	10	"github.com/go-kit/kit/loadbalancer/static"
	11	)
	12
	13	func TestRandomDistribution(t *testing.T) {
	14	var (
	15	n = 3
	16	endpoints = make([]endpoint.Endpoint, n)
	17	counts = make([]int, n)
	18	seed = int64(123)
	19	ctx = context.Background()
	20	iterations = 100000
	21	want = iterations / n
	22	tolerance = want / 100 // 1%
	23	)
	24
	25	for i := 0; i < n; i++ {
	26	i0 := i
	27	endpoints[i] = func(context.Context, interface{}) (interface{}, error) { counts[i0]++; return struct{}{}, nil }
	28	}
	29
	30	lb := loadbalancer.NewRandom(static.Publisher(endpoints), seed)
	31
	32	for i := 0; i < iterations; i++ {
	33	e, err := lb.Endpoint()
	34	if err != nil {
	35	t.Fatal(err)
	36	}
	37	e(ctx, struct{}{})
	38	}
	39
	40	for i, have := range counts {
	41	if math.Abs(float64(want-have)) > float64(tolerance) {
	42	t.Errorf("%d: want %d, have %d", i, want, have)
	43	}
	44	}
	45	}
	46
	47	func TestRandomBadPublisher(t *testing.T) {
	48	t.Skip("TODO")
	49	}
	50
	51	func TestRandomNoEndpoints(t *testing.T) {
	52	lb := loadbalancer.NewRandom(static.Publisher([]endpoint.Endpoint{}), 123)
	53	_, have := lb.Endpoint()
	54	if want := loadbalancer.ErrNoEndpoints; want != have {
	55	t.Errorf("want %q, have %q", want, have)
	56	}
	57	}

+0

-52

~~loadbalancer/retry.go~~ less more

0		package loadbalancer
1
2		import (
3		"fmt"
4		"strings"
5		"time"
6
7		"golang.org/x/net/context"
8
9		"github.com/go-kit/kit/endpoint"
10		)
11
12		// Retry yields an endpoint that takes endpoints from the load balancer.
13		// Invocations that return errors will be retried until they succeed, up to
14		// max times, or until the timeout is elapsed, whichever comes first.
15		func Retry(max int, timeout time.Duration, lb LoadBalancer) endpoint.Endpoint {
16		return func(ctx context.Context, request interface{}) (interface{}, error) {
17		var (
18		newctx, cancel = context.WithTimeout(ctx, timeout)
19		responses = make(chan interface{}, 1)
20		errs = make(chan error, 1)
21		a = []string{}
22		)
23		defer cancel()
24		for i := 1; i <= max; i++ {
25		go func() {
26		e, err := lb.Get()
27		if err != nil {
28		errs <- err
29		return
30		}
31		response, err := e(newctx, request)
32		if err != nil {
33		errs <- err
34		return
35		}
36		responses <- response
37		}()
38
39		select {
40		case <-newctx.Done():
41		return nil, newctx.Err()
42		case response := <-responses:
43		return response, nil
44		case err := <-errs:
45		a = append(a, err.Error())
46		continue
47		}
48		}
49		return nil, fmt.Errorf("retry attempts exceeded (%s)", strings.Join(a, "; "))
50		}
51		}

+0

-67

~~loadbalancer/retry_test.go~~ less more

0		package loadbalancer_test
1
2		import (
3		"errors"
4		"testing"
5		"time"
6
7		"golang.org/x/net/context"
8
9		"github.com/go-kit/kit/endpoint"
10		"github.com/go-kit/kit/loadbalancer"
11		"github.com/go-kit/kit/loadbalancer/publisher/static"
12		"github.com/go-kit/kit/loadbalancer/strategy"
13		)
14
15		func TestRetryMax(t *testing.T) {
16		var (
17		endpoints = []endpoint.Endpoint{}
18		p = static.NewPublisher(endpoints)
19		lb = strategy.RoundRobin(p)
20		)
21
22		if _, err := loadbalancer.Retry(999, time.Second, lb)(context.Background(), struct{}{}); err == nil {
23		t.Errorf("expected error, got none")
24		}
25
26		endpoints = []endpoint.Endpoint{
27		func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error one") },
28		func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error two") },
29		func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil /* OK */ },
30		}
31		p.Replace(endpoints)
32		time.Sleep(10 * time.Millisecond) //assertLoadBalancerNotEmpty(t, lb) // TODO
33
34		if _, err := loadbalancer.Retry(len(endpoints)-1, time.Second, lb)(context.Background(), struct{}{}); err == nil {
35		t.Errorf("expected error, got none")
36		}
37
38		if _, err := loadbalancer.Retry(len(endpoints), time.Second, lb)(context.Background(), struct{}{}); err != nil {
39		t.Error(err)
40		}
41		}
42
43		func TestRetryTimeout(t *testing.T) {
44		var (
45		step = make(chan struct{})
46		e = func(context.Context, interface{}) (interface{}, error) { <-step; return struct{}{}, nil }
47		timeout = time.Millisecond
48		retry = loadbalancer.Retry(999, timeout, strategy.RoundRobin(static.NewPublisher([]endpoint.Endpoint{e})))
49		errs = make(chan error)
50		invoke = func() { _, err := retry(context.Background(), struct{}{}); errs <- err }
51		)
52
53		go invoke() // invoke the endpoint
54		step <- struct{}{} // tell the endpoint to return
55		if err := <-errs; err != nil { // that should succeed
56		t.Error(err)
57		}
58
59		go invoke() // invoke the endpoint
60		time.Sleep(2 * timeout) // wait
61		time.Sleep(2 * timeout) // wait again (CI servers!!)
62		step <- struct{}{} // tell the endpoint to return
63		if err := <-errs; err != context.DeadlineExceeded { // that should not succeed
64		t.Errorf("wanted error, got none")
65		}
66		}

+41

-0

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 is a simple load balancer that returns each of the published
	9	// endpoints in sequence.
	10	type RoundRobin struct {
	11	p Publisher
	12	counter uint64
	13	}
	14
	15	// NewRoundRobin returns a new RoundRobin load balancer.
	16	func NewRoundRobin(p Publisher) *RoundRobin {
	17	return &RoundRobin{
	18	p: p,
	19	counter: 0,
	20	}
	21	}
	22
	23	// Endpoint implements the LoadBalancer interface.
	24	func (rr *RoundRobin) Endpoint() (endpoint.Endpoint, error) {
	25	endpoints, err := rr.p.Endpoints()
	26	if err != nil {
	27	return nil, err
	28	}
	29	if len(endpoints) <= 0 {
	30	return nil, ErrNoEndpoints
	31	}
	32	var old uint64
	33	for {
	34	old = atomic.LoadUint64(&rr.counter)
	35	if atomic.CompareAndSwapUint64(&rr.counter, old, old+1) {
	36	break
	37	}
	38	}
	39	return endpoints[old%uint64(len(endpoints))], nil
	40	}

+49

-0

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	"github.com/go-kit/kit/loadbalancer/static"
	9	"golang.org/x/net/context"
	10	)
	11
	12	func TestRoundRobinDistribution(t *testing.T) {
	13	var (
	14	ctx = context.Background()
	15	counts = []int{0, 0, 0}
	16	endpoints = []endpoint.Endpoint{
	17	func(context.Context, interface{}) (interface{}, error) { counts[0]++; return struct{}{}, nil },
	18	func(context.Context, interface{}) (interface{}, error) { counts[1]++; return struct{}{}, nil },
	19	func(context.Context, interface{}) (interface{}, error) { counts[2]++; return struct{}{}, nil },
	20	}
	21	)
	22
	23	lb := loadbalancer.NewRoundRobin(static.Publisher(endpoints))
	24
	25	for i, want := range [][]int{
	26	{1, 0, 0},
	27	{1, 1, 0},
	28	{1, 1, 1},
	29	{2, 1, 1},
	30	{2, 2, 1},
	31	{2, 2, 2},
	32	{3, 2, 2},
	33	} {
	34	e, err := lb.Endpoint()
	35	if err != nil {
	36	t.Fatal(err)
	37	}
	38	e(ctx, struct{}{})
	39	if have := counts; !reflect.DeepEqual(want, have) {
	40	t.Fatalf("%d: want %v, have %v", i, want, have)
	41	}
	42
	43	}
	44	}
	45
	46	func TestRoundRobinBadPublisher(t *testing.T) {
	47	t.Skip("TODO")
	48	}

+9

-0

loadbalancer/static/publisher.go less more

	0	package static
	1
	2	import "github.com/go-kit/kit/endpoint"
	3
	4	// Publisher yields the same set of static endpoints.
	5	type Publisher []endpoint.Endpoint
	6
	7	// Endpoints implements the Publisher interface.
	8	func (p Publisher) Endpoints() ([]endpoint.Endpoint, error) { return p, nil }

+27

-0

loadbalancer/static/publisher_test.go less more

	0	package static_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/static"
	10	)
	11
	12	func TestStatic(t *testing.T) {
	13	var (
	14	e1 = func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
	15	e2 = func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
	16	endpoints = []endpoint.Endpoint{e1, e2}
	17	)
	18	p := static.Publisher(endpoints)
	19	have, err := p.Endpoints()
	20	if err != nil {
	21	t.Fatal(err)
	22	}
	23	if want := endpoints; !reflect.DeepEqual(want, have) {
	24	t.Fatalf("want %#+v, have %#+v", want, have)
	25	}
	26	}

+0

-50

~~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		}

+0

-34

~~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		}

+0

-26

~~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		}

+0

-43

~~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		}

+0

-36

~~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		}

+0

-46

~~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		}

+0

-9

~~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		}

+0

-35

~~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		}