github.com/andy2046/gopie@v0.7.0/pkg/breaker/breaker.go (about)

     1  // Package breaker implements Circuit Breaker pattern.
     2  package breaker
     3  
     4  import (
     5  	"errors"
     6  	"fmt"
     7  	"log"
     8  	"sync"
     9  	"time"
    10  )
    11  
    12  // State is the type representing a state of CircuitBreaker.
    13  type State int
    14  
    15  const (
    16  	// StateClosed represents Closed State.
    17  	StateClosed State = iota
    18  	// StateHalfOpen represents HalfOpen State.
    19  	StateHalfOpen
    20  	// StateOpen represents Open State.
    21  	StateOpen
    22  )
    23  
    24  var (
    25  	// ErrTooManyRequests is returned when the state is half open
    26  	// and the requests count is more the maxRequests.
    27  	ErrTooManyRequests = errors.New("too many requests, requests count is more the maxRequests in half open state")
    28  	// ErrOpenState is returned when the state is open.
    29  	ErrOpenState = errors.New("circuit breaker is open")
    30  )
    31  
    32  // String implements stringer interface.
    33  func (s State) String() string {
    34  	switch s {
    35  	case StateClosed:
    36  		return "closed"
    37  	case StateHalfOpen:
    38  		return "half-open"
    39  	case StateOpen:
    40  		return "open"
    41  	default:
    42  		return fmt.Sprintf("unknown state: %d", s)
    43  	}
    44  }
    45  
    46  // Counts holds the numbers of requests and their successes/failures.
    47  // CircuitBreaker clears the internal Counts either
    48  // on the change of the state or at the closed-state intervals.
    49  // Counts ignores the results of the requests sent before clearing.
    50  type Counts struct {
    51  	Requests             uint64
    52  	TotalSuccesses       uint64
    53  	TotalFailures        uint64
    54  	ConsecutiveSuccesses uint64
    55  	ConsecutiveFailures  uint64
    56  }
    57  
    58  func (c *Counts) onRequest() {
    59  	c.Requests++
    60  }
    61  
    62  func (c *Counts) onSuccess() {
    63  	c.TotalSuccesses++
    64  	c.ConsecutiveSuccesses++
    65  	c.ConsecutiveFailures = 0
    66  }
    67  
    68  func (c *Counts) onFailure() {
    69  	c.TotalFailures++
    70  	c.ConsecutiveFailures++
    71  	c.ConsecutiveSuccesses = 0
    72  }
    73  
    74  func (c *Counts) clear() {
    75  	c.Requests = 0
    76  	c.TotalSuccesses = 0
    77  	c.TotalFailures = 0
    78  	c.ConsecutiveSuccesses = 0
    79  	c.ConsecutiveFailures = 0
    80  }
    81  
    82  type (
    83  	// Settings represents settings for CircuitBreaker.
    84  	Settings struct {
    85  		// Name is the name of the CircuitBreaker.
    86  		Name string
    87  		// MaxRequests is the maximum number of requests allowed to pass through
    88  		// when the CircuitBreaker is half-open.
    89  		// If MaxRequests is 0, the CircuitBreaker allows only 1 request.
    90  		MaxRequests uint64
    91  		// Interval is the cyclic period of the closed state
    92  		// for the CircuitBreaker to clear the internal Counts.
    93  		// If Interval is 0, the CircuitBreaker doesn't clear internal Counts during the closed state.
    94  		Interval time.Duration
    95  		// Timeout is the period of the open state,
    96  		// after which the state of the CircuitBreaker becomes half-open.
    97  		// If Timeout is 0, the timeout for the CircuitBreaker is 60 seconds.
    98  		Timeout time.Duration
    99  		// ShouldTrip is called with a copy of Counts whenever a request fails in the closed state.
   100  		// If ShouldTrip returns true, the CircuitBreaker will be placed into the open state.
   101  		// If ShouldTrip is nil, default ShouldTrip is used.
   102  		// Default ShouldTrip returns true when the number of consecutive failures is more than 5.
   103  		ShouldTrip func(counts Counts) bool
   104  		// OnStateChange is called whenever the state of the CircuitBreaker changes.
   105  		OnStateChange func(name string, from, to State)
   106  	}
   107  
   108  	// CircuitBreaker prevent an application repeatedly trying to execute an operation that is likely to fail.
   109  	CircuitBreaker struct {
   110  		name          string
   111  		maxRequests   uint64
   112  		interval      time.Duration
   113  		timeout       time.Duration
   114  		shouldTrip    func(counts Counts) bool
   115  		onStateChange func(name string, from, to State)
   116  
   117  		mutex      sync.Mutex
   118  		state      State
   119  		generation uint64
   120  		counts     Counts
   121  		expiry     time.Time
   122  	}
   123  
   124  	// Option applies settings to CircuitBreaker Settings.
   125  	Option = func(*Settings) error
   126  )
   127  
   128  // DefaultSettings is the default CircuitBreaker Settings.
   129  var DefaultSettings = Settings{
   130  	Name:        "CircuitBreaker",
   131  	MaxRequests: 1,
   132  	Interval:    0,
   133  	Timeout:     60 * time.Second,
   134  	ShouldTrip: func(counts Counts) bool {
   135  		return counts.ConsecutiveFailures > 5
   136  	},
   137  	OnStateChange: nil,
   138  }
   139  
   140  func setOption(s *Settings, options ...func(*Settings) error) error {
   141  	for _, opt := range options {
   142  		if err := opt(s); err != nil {
   143  			return err
   144  		}
   145  	}
   146  	return nil
   147  }
   148  
   149  // New returns a new CircuitBreaker with options applied.
   150  func New(options ...Option) *CircuitBreaker {
   151  	st := DefaultSettings
   152  	err := setOption(&st, options...)
   153  	if err != nil {
   154  		log.Panicf("fail to apply Settings -> %v\n", err)
   155  	}
   156  
   157  	cb := &CircuitBreaker{
   158  		name:          st.Name,
   159  		maxRequests:   st.MaxRequests,
   160  		interval:      st.Interval,
   161  		timeout:       st.Timeout,
   162  		shouldTrip:    st.ShouldTrip,
   163  		onStateChange: st.OnStateChange,
   164  	}
   165  
   166  	cb.toNewGeneration(time.Now())
   167  	return cb
   168  }
   169  
   170  // Name returns the name of the CircuitBreaker.
   171  func (cb *CircuitBreaker) Name() string {
   172  	return cb.name
   173  }
   174  
   175  // State returns the current state of the CircuitBreaker.
   176  func (cb *CircuitBreaker) State() State {
   177  	cb.mutex.Lock()
   178  	defer cb.mutex.Unlock()
   179  
   180  	now := time.Now()
   181  	state, _ := cb.currentState(now)
   182  	return state
   183  }
   184  
   185  // Execute runs the given request if the CircuitBreaker accepts it.
   186  // Execute returns an error instantly if the CircuitBreaker rejects the request.
   187  // Otherwise, Execute returns the result of the request.
   188  // If a panic occurs in the request, the CircuitBreaker handles it as an error
   189  // and causes the same panic again.
   190  func (cb *CircuitBreaker) Execute(request func() (interface{}, error)) (interface{}, error) {
   191  	generation, err := cb.beforeRequest()
   192  	if err != nil {
   193  		return nil, err
   194  	}
   195  
   196  	defer func() {
   197  		e := recover()
   198  		if e != nil {
   199  			cb.afterRequest(generation, false)
   200  			panic(e)
   201  		}
   202  	}()
   203  
   204  	result, err := request()
   205  	cb.afterRequest(generation, err == nil)
   206  	return result, err
   207  }
   208  
   209  func (cb *CircuitBreaker) beforeRequest() (uint64, error) {
   210  	cb.mutex.Lock()
   211  	defer cb.mutex.Unlock()
   212  
   213  	now := time.Now()
   214  	state, generation := cb.currentState(now)
   215  
   216  	if state == StateOpen {
   217  		return generation, ErrOpenState
   218  	} else if state == StateHalfOpen && cb.counts.Requests >= cb.maxRequests {
   219  		return generation, ErrTooManyRequests
   220  	}
   221  
   222  	cb.counts.onRequest()
   223  	return generation, nil
   224  }
   225  
   226  func (cb *CircuitBreaker) afterRequest(before uint64, success bool) {
   227  	cb.mutex.Lock()
   228  	defer cb.mutex.Unlock()
   229  
   230  	now := time.Now()
   231  	state, generation := cb.currentState(now)
   232  	if generation != before {
   233  		return
   234  	}
   235  
   236  	if success {
   237  		cb.onSuccess(state, now)
   238  	} else {
   239  		cb.onFailure(state, now)
   240  	}
   241  }
   242  
   243  func (cb *CircuitBreaker) onSuccess(state State, now time.Time) {
   244  	switch state {
   245  	case StateClosed:
   246  		cb.counts.onSuccess()
   247  	case StateHalfOpen:
   248  		cb.counts.onSuccess()
   249  		if cb.counts.ConsecutiveSuccesses >= cb.maxRequests {
   250  			cb.setState(StateClosed, now)
   251  		}
   252  	}
   253  }
   254  
   255  func (cb *CircuitBreaker) onFailure(state State, now time.Time) {
   256  	switch state {
   257  	case StateClosed:
   258  		cb.counts.onFailure()
   259  		if cb.shouldTrip(cb.counts) {
   260  			cb.setState(StateOpen, now)
   261  		}
   262  	case StateHalfOpen:
   263  		cb.setState(StateOpen, now)
   264  	}
   265  }
   266  
   267  func (cb *CircuitBreaker) currentState(now time.Time) (State, uint64) {
   268  	switch cb.state {
   269  	case StateClosed:
   270  		if !cb.expiry.IsZero() && cb.expiry.Before(now) {
   271  			cb.toNewGeneration(now)
   272  		}
   273  	case StateOpen:
   274  		if cb.expiry.Before(now) {
   275  			cb.setState(StateHalfOpen, now)
   276  		}
   277  	}
   278  	return cb.state, cb.generation
   279  }
   280  
   281  func (cb *CircuitBreaker) setState(state State, now time.Time) {
   282  	if cb.state == state {
   283  		return
   284  	}
   285  
   286  	prev := cb.state
   287  	cb.state = state
   288  
   289  	cb.toNewGeneration(now)
   290  
   291  	if cb.onStateChange != nil {
   292  		cb.onStateChange(cb.name, prev, state)
   293  	}
   294  }
   295  
   296  func (cb *CircuitBreaker) toNewGeneration(now time.Time) {
   297  	cb.generation++
   298  	cb.counts.clear()
   299  
   300  	var zero time.Time
   301  	switch cb.state {
   302  	case StateClosed:
   303  		if cb.interval == 0 {
   304  			cb.expiry = zero
   305  		} else {
   306  			cb.expiry = now.Add(cb.interval)
   307  		}
   308  	case StateOpen:
   309  		cb.expiry = now.Add(cb.timeout)
   310  	default: // for StateHalfOpen
   311  		cb.expiry = zero
   312  	}
   313  }