gobreaker.go

package breaker

import (
	"context"
	"sync"

	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/errors"
	"github.com/sony/gobreaker"
)

type BreakerMethod int

const (
	BreakService BreakerMethod = iota
	BreakServiceEndpoint
)

type clientWrapper struct {
	bs  gobreaker.Settings
	bm  BreakerMethod
	cbs map[string]*gobreaker.CircuitBreaker
	mu  sync.Mutex
	client.Client
}

func (c *clientWrapper) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...client.CallOption) error {
	var svc string

	switch c.bm {
	case BreakService:
		svc = req.Service()
	case BreakServiceEndpoint:
		svc = req.Service() + "." + req.Endpoint()
	}

	c.mu.Lock()
	cb, ok := c.cbs[svc]
	if !ok {
		cb = gobreaker.NewCircuitBreaker(c.bs)
		c.cbs[svc] = cb
	}
	c.mu.Unlock()

	_, err := cb.Execute(func() (interface{}, error) {
		cerr := c.Client.Call(ctx, req, rsp, opts...)
		return nil, cerr
	})

	if err != nil {
		err = errors.New(req.Service(), err.Error(), 503)
	}

	return err
}

// NewClientWrapper takes a gobreaker.Settings and returns a client Wrapper.
func NewClientWrapper(bs gobreaker.Settings, bm BreakerMethod) client.Wrapper {
	return func(c client.Client) client.Client {
		w := &clientWrapper{}
		w.bm = bm
		w.bs = bs
		w.cbs = make(map[string]*gobreaker.CircuitBreaker)
		w.Client = c
		return w
	}
}