Codebase list golang-github-go-kit-kit / 059ec16
Retry wrapper Peter Bourgon 8 years ago
2 changed file(s) with 86 addition(s) and 0 deletion(s). Raw diff Collapse all Expand all
0 package loadbalancer
1
2 import (
3 "fmt"
4 "strings"
5
6 "golang.org/x/net/context"
7
8 "github.com/go-kit/kit/endpoint"
9 )
10
11 // Retry yields an endpoint that invokes the load balancer up to max times.
12 func Retry(max int, lb LoadBalancer) endpoint.Endpoint {
13 return func(ctx context.Context, request interface{}) (interface{}, error) {
14 errs := []string{}
15 for i := 1; i <= max; i++ {
16 e, err := lb.Get()
17 if err != nil {
18 errs = append(errs, err.Error())
19 return nil, fmt.Errorf("%s", strings.Join(errs, "; ")) // fatal
20 }
21 response, err := e(ctx, request)
22 if err != nil {
23 errs = append(errs, err.Error())
24 continue // try again
25 }
26 return response, err
27 }
28 if len(errs) <= 0 {
29 panic("impossible state in retry load balancer")
30 }
31 return nil, fmt.Errorf("%s", strings.Join(errs, "; "))
32 }
33 }
0 package loadbalancer_test
1
2 import (
3 "errors"
4 "runtime"
5
6 "github.com/go-kit/kit/endpoint"
7 "github.com/go-kit/kit/loadbalancer"
8 "golang.org/x/net/context"
9
10 "testing"
11 )
12
13 func TestRetry(t *testing.T) {
14 var (
15 endpoints = []endpoint.Endpoint{}
16 p = loadbalancer.NewStaticPublisher(endpoints)
17 lb = loadbalancer.RoundRobin(p)
18 )
19
20 {
21 max := 999
22 e := loadbalancer.Retry(max, lb)
23 if _, err := e(context.Background(), struct{}{}); err == nil {
24 t.Errorf("expected error, got none")
25 }
26 }
27
28 endpoints = []endpoint.Endpoint{
29 func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error one") },
30 func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error two") },
31 func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil /* OK */ },
32 }
33 p.Replace(endpoints)
34 runtime.Gosched()
35
36 {
37 max := len(endpoints) - 1
38 e := loadbalancer.Retry(max, lb)
39 if _, err := e(context.Background(), struct{}{}); err == nil {
40 t.Errorf("expected error, got none")
41 }
42 }
43
44 {
45 max := len(endpoints)
46 e := loadbalancer.Retry(max, lb)
47 if _, err := e(context.Background(), struct{}{}); err != nil {
48 t.Error(err)
49 }
50 }
51 }