You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
65 lines
1.8 KiB
65 lines
1.8 KiB
2 years ago
|
package circuitbreaker
|
||
|
|
||
|
import (
|
||
|
"context"
|
||
|
|
||
|
"github.com/go-kratos/aegis/circuitbreaker"
|
||
|
"github.com/go-kratos/aegis/circuitbreaker/sre"
|
||
|
"github.com/go-kratos/kratos/v2/errors"
|
||
|
"github.com/go-kratos/kratos/v2/internal/group"
|
||
|
"github.com/go-kratos/kratos/v2/middleware"
|
||
|
"github.com/go-kratos/kratos/v2/transport"
|
||
|
)
|
||
|
|
||
|
// ErrNotAllowed is request failed due to circuit breaker triggered.
|
||
|
var ErrNotAllowed = errors.New(503, "CIRCUITBREAKER", "request failed due to circuit breaker triggered")
|
||
|
|
||
|
// Option is circuit breaker option.
|
||
|
type Option func(*options)
|
||
|
|
||
|
// WithGroup with circuit breaker group.
|
||
|
// NOTE: implements generics circuitbreaker.CircuitBreaker
|
||
|
func WithGroup(g *group.Group) Option {
|
||
|
return func(o *options) {
|
||
|
o.group = g
|
||
|
}
|
||
|
}
|
||
|
|
||
|
type options struct {
|
||
|
group *group.Group
|
||
|
}
|
||
|
|
||
|
// Client circuitbreaker middleware will return errBreakerTriggered when the circuit
|
||
|
// breaker is triggered and the request is rejected directly.
|
||
|
func Client(opts ...Option) middleware.Middleware {
|
||
|
opt := &options{
|
||
|
group: group.NewGroup(func() interface{} {
|
||
|
return sre.NewBreaker()
|
||
|
}),
|
||
|
}
|
||
|
for _, o := range opts {
|
||
|
o(opt)
|
||
|
}
|
||
|
return func(handler middleware.Handler) middleware.Handler {
|
||
|
return func(ctx context.Context, req interface{}) (interface{}, error) {
|
||
|
info, _ := transport.FromClientContext(ctx)
|
||
|
breaker := opt.group.Get(info.Operation()).(circuitbreaker.CircuitBreaker)
|
||
|
if err := breaker.Allow(); err != nil {
|
||
|
// rejected
|
||
|
// NOTE: when client reject requets locally,
|
||
|
// continue add counter let the drop ratio higher.
|
||
|
breaker.MarkFailed()
|
||
|
return nil, ErrNotAllowed
|
||
|
}
|
||
|
// allowed
|
||
|
reply, err := handler(ctx, req)
|
||
|
if err != nil && (errors.IsInternalServer(err) || errors.IsServiceUnavailable(err) || errors.IsGatewayTimeout(err)) {
|
||
|
breaker.MarkFailed()
|
||
|
} else {
|
||
|
breaker.MarkSuccess()
|
||
|
}
|
||
|
return reply, err
|
||
|
}
|
||
|
}
|
||
|
}
|