toolkit/http/loadbalancer.go

package http

import (
	"context"
	"io"
	"net/url"
	"strings"
	"time"

	"github.com/go-kit/kit/endpoint"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/sd"
	"github.com/go-kit/kit/sd/lb"

	consulsd "github.com/go-kit/kit/sd/consul"
)

// dec DecodeRequestFunc, enc EncodeResponseFunc

// LoadBalancer endpoint
type LoadBalancer struct {
	Ctx    context.Context
	Inst   *consulsd.Instancer
	Route  string
	Dec    DecodeResponseFunc
	Logger log.Logger
}

var (
	// retry max
	retryMax = 3
	// retry timeout
	retryTimeout = 500 * time.Millisecond
)

func factory(ctx context.Context, method, router string, dec DecodeResponseFunc) sd.Factory {
	return func(instance string) (endpoint.Endpoint, io.Closer, error) {
		if !strings.HasPrefix(instance, "http") {
			instance = "http://" + instance
		}
		tgt, err := url.Parse(instance)
		if err != nil {
			return nil, nil, err
		}

		return ClientRequestEndpoint(ctx, tgt, method, router, dec), nil, nil
	}
}

// FactoryLoadBalancer factory load balance
func FactoryLoadBalancer(ctx context.Context, instancer *consulsd.Instancer, method, route string, dec DecodeResponseFunc, logger log.Logger) endpoint.Endpoint {
	endpointer := sd.NewEndpointer(
		instancer,
		factory(ctx, method, route, dec),
		logger)
	balancer := lb.NewRoundRobin(endpointer)

	return lb.Retry(retryMax, retryTimeout, balancer)
}

// GetLoadBalancer add GET router to load balance
func GetLoadBalancer(loadbalancer LoadBalancer) endpoint.Endpoint {
	endpointer := sd.NewEndpointer(loadbalancer.Inst, factory(loadbalancer.Ctx, `GET`, loadbalancer.Route, loadbalancer.Dec), loadbalancer.Logger)
	balancer := lb.NewRoundRobin(endpointer)

	return lb.Retry(retryMax, retryTimeout, balancer)
}

// PostLoadBalancer add POST router to load balance
func PostLoadBalancer(loadbalancer LoadBalancer) endpoint.Endpoint {
	endpointer := sd.NewEndpointer(loadbalancer.Inst, factory(loadbalancer.Ctx, `POST`, loadbalancer.Route, loadbalancer.Dec), loadbalancer.Logger)
	balancer := lb.NewRoundRobin(endpointer)

	return lb.Retry(retryMax, retryTimeout, balancer)
}

// PutLoadBalancer add PUT router to load balance
func PutLoadBalancer(loadbalancer LoadBalancer) endpoint.Endpoint {
	endpointer := sd.NewEndpointer(loadbalancer.Inst, factory(loadbalancer.Ctx, `PUT`, loadbalancer.Route, loadbalancer.Dec), loadbalancer.Logger)
	balancer := lb.NewRoundRobin(endpointer)

	return lb.Retry(retryMax, retryTimeout, balancer)
}

// HeadLoadBalancer add HEAD router to load balance
func HeadLoadBalancer(loadbalancer LoadBalancer) endpoint.Endpoint {
	endpointer := sd.NewEndpointer(loadbalancer.Inst, factory(loadbalancer.Ctx, `HEAD`, loadbalancer.Route, loadbalancer.Dec), loadbalancer.Logger)
	balancer := lb.NewRoundRobin(endpointer)

	return lb.Retry(retryMax, retryTimeout, balancer)
}

// DeleteLoadBalancer add DELETE router to load balance
func DeleteLoadBalancer(loadbalancer LoadBalancer) endpoint.Endpoint {
	endpointer := sd.NewEndpointer(loadbalancer.Inst, factory(loadbalancer.Ctx, `DELETE`, loadbalancer.Route, loadbalancer.Dec), loadbalancer.Logger)
	balancer := lb.NewRoundRobin(endpointer)

	return lb.Retry(retryMax, retryTimeout, balancer)
}

// OptionsLoadBalancer add OPTIONS router to load balance
func OptionsLoadBalancer(loadbalancer LoadBalancer) endpoint.Endpoint {
	endpointer := sd.NewEndpointer(loadbalancer.Inst, factory(loadbalancer.Ctx, `OPTIONS`, loadbalancer.Route, loadbalancer.Dec), loadbalancer.Logger)
	balancer := lb.NewRoundRobin(endpointer)

	return lb.Retry(retryMax, retryTimeout, balancer)
}

// PatchLoadBalancer add PATCH router to load balance
func PatchLoadBalancer(loadbalancer LoadBalancer) endpoint.Endpoint {
	endpointer := sd.NewEndpointer(loadbalancer.Inst, factory(loadbalancer.Ctx, `PATCH`, loadbalancer.Route, loadbalancer.Dec), loadbalancer.Logger)
	balancer := lb.NewRoundRobin(endpointer)

	return lb.Retry(retryMax, retryTimeout, balancer)
}