func TestRoundRobinDistribution(t *testing.T) {
var (
ctx = context.Background()
counts = []int{0, 0, 0}
endpoints = []endpoint.Endpoint{
func(context.Context, interface{}) (interface{}, error) { counts[0]++; return struct{}{}, nil },
func(context.Context, interface{}) (interface{}, error) { counts[1]++; return struct{}{}, nil },
func(context.Context, interface{}) (interface{}, error) { counts[2]++; return struct{}{}, nil },
}
)
lb := loadbalancer.NewRoundRobin(static.NewPublisher(endpoints))
for i, want := range [][]int{
{1, 0, 0},
{1, 1, 0},
{1, 1, 1},
{2, 1, 1},
{2, 2, 1},
{2, 2, 2},
{3, 2, 2},
} {
e, err := lb.Endpoint()
if err != nil {
t.Fatal(err)
}
e(ctx, struct{}{})
if have := counts; !reflect.DeepEqual(want, have) {
t.Fatalf("%d: want %v, have %v", i, want, have)
}
}
}
func TestRetryMaxTotalFail(t *testing.T) {
var (
endpoints = []endpoint.Endpoint{} // no endpoints
p = fixed.NewPublisher(endpoints)
lb = loadbalancer.NewRoundRobin(p)
retry = loadbalancer.Retry(999, time.Second, lb) // lots of retries
ctx = context.Background()
)
if _, err := retry(ctx, struct{}{}); err == nil {
t.Errorf("expected error, got none") // should fail
}
}
func TestRetryMaxSuccess(t *testing.T) {
var (
endpoints = []endpoint.Endpoint{
func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error one") },
func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error two") },
func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil /* OK */ },
}
retries = len(endpoints) // exactly enough retries
p = fixed.NewPublisher(endpoints)
lb = loadbalancer.NewRoundRobin(p)
ctx = context.Background()
)
if _, err := loadbalancer.Retry(retries, time.Second, lb)(ctx, struct{}{}); err != nil {
t.Error(err)
}
}
func TestRetryMaxPartialFail(t *testing.T) {
var (
endpoints = []endpoint.Endpoint{
func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error one") },
func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error two") },
func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil /* OK */ },
}
retries = len(endpoints) - 1 // not quite enough retries
p = fixed.NewPublisher(endpoints)
lb = loadbalancer.NewRoundRobin(p)
ctx = context.Background()
)
if _, err := loadbalancer.Retry(retries, time.Second, lb)(ctx, struct{}{}); err == nil {
t.Errorf("expected error, got none")
}
}
func proxyingMiddleware(proxyList string, ctx context.Context, logger log.Logger) ServiceMiddleware {
if proxyList == "" {
_ = logger.Log("proxy_to", "none")
return func(next StringService) StringService { return next }
}
proxies := split(proxyList)
_ = logger.Log("proxy_to", fmt.Sprint(proxies))
return func(next StringService) StringService {
var (
qps = 100 // max to each instance
publisher = static.NewPublisher(proxies, factory(ctx, qps), logger)
lb = loadbalancer.NewRoundRobin(publisher)
maxAttempts = 3
maxTime = 100 * time.Millisecond
endpoint = loadbalancer.Retry(maxAttempts, maxTime, lb)
)
return proxymw{ctx, endpoint, next}
}
}
func TestRetryTimeout(t *testing.T) {
var (
step = make(chan struct{})
e = func(context.Context, interface{}) (interface{}, error) { <-step; return struct{}{}, nil }
timeout = time.Millisecond
retry = loadbalancer.Retry(999, timeout, loadbalancer.NewRoundRobin(fixed.NewPublisher([]endpoint.Endpoint{e})))
errs = make(chan error, 1)
invoke = func() { _, err := retry(context.Background(), struct{}{}); errs <- err }
)
go func() { step <- struct{}{} }() // queue up a flush of the endpoint
invoke() // invoke the endpoint and trigger the flush
if err := <-errs; err != nil { // that should succeed
t.Error(err)
}
go func() { time.Sleep(10 * timeout); step <- struct{}{} }() // a delayed flush
invoke() // invoke the endpoint
if err := <-errs; err != context.DeadlineExceeded { // that should not succeed
t.Errorf("wanted %v, got none", context.DeadlineExceeded)
}
}
func main() {
var (
httpAddr = flag.String("http.addr", ":8000", "Address for HTTP (JSON) server")
consulAddr = flag.String("consul.addr", "", "Consul agent address")
retryMax = flag.Int("retry.max", 3, "per-request retries to different instances")
retryTimeout = flag.Duration("retry.timeout", 500*time.Millisecond, "per-request timeout, including retries")
)
flag.Parse()
// Log domain
logger := log.NewLogfmtLogger(os.Stderr)
logger = log.NewContext(logger).With("ts", log.DefaultTimestampUTC).With("caller", log.DefaultCaller)
stdlog.SetFlags(0) // flags are handled by Go kit's logger
stdlog.SetOutput(log.NewStdlibAdapter(logger)) // redirect anything using stdlib log to us
// Service discovery domain. In this example we use Consul.
consulConfig := api.DefaultConfig()
if len(*consulAddr) > 0 {
consulConfig.Address = *consulAddr
}
consulClient, err := api.NewClient(consulConfig)
if err != nil {
logger.Log("err", err)
os.Exit(1)
}
discoveryClient := consul.NewClient(consulClient)
// Context domain.
ctx := context.Background()
// Set up our routes.
//
// Each Consul service name maps to multiple instances of that service. We
// connect to each instance according to its pre-determined transport: in this
// case, we choose to access addsvc via its gRPC client, and stringsvc over
// plain transport/http (it has no client package).
//
// Each service instance implements multiple methods, and we want to map each
// method to a unique path on the API gateway. So, we define that path and its
// corresponding factory function, which takes an instance string and returns an
// endpoint.Endpoint for the specific method.
//
// Finally, we mount that path + endpoint handler into the router.
r := mux.NewRouter()
for consulName, methods := range map[string][]struct {
path string
factory loadbalancer.Factory
}{
"addsvc": {
{path: "/api/addsvc/concat", factory: grpc.MakeConcatEndpointFactory(opentracing.GlobalTracer(), nil)},
{path: "/api/addsvc/sum", factory: grpc.MakeSumEndpointFactory(opentracing.GlobalTracer(), nil)},
},
"stringsvc": {
{path: "/api/stringsvc/uppercase", factory: httpFactory(ctx, "GET", "uppercase/")},
{path: "/api/stringsvc/concat", factory: httpFactory(ctx, "GET", "concat/")},
},
} {
for _, method := range methods {
publisher, err := consul.NewPublisher(discoveryClient, method.factory, logger, consulName)
if err != nil {
logger.Log("service", consulName, "path", method.path, "err", err)
continue
}
lb := loadbalancer.NewRoundRobin(publisher)
e := loadbalancer.Retry(*retryMax, *retryTimeout, lb)
h := makeHandler(ctx, e, logger)
r.HandleFunc(method.path, h)
}
}
// Mechanical stuff.
errc := make(chan error)
go func() {
errc <- interrupt()
}()
go func() {
logger.Log("transport", "http", "addr", *httpAddr)
errc <- http.ListenAndServe(*httpAddr, r)
}()
logger.Log("err", <-errc)
}