Add breaker.GetState() method to inspect state

This is useful in a few scenarios where the user might want to check the
breaker state directly (perhaps for monitoring purposes), but also makes
the tests much clearer because we can just assert on the expected
breaker state directly.

Author: eapache

breaker/breaker.go

@@ -12,10 +12,13 @@ import (
 // because the breaker is currently open.
 var ErrBreakerOpen = errors.New("circuit breaker is open")
 
+// State is a type representing the possible states of a circuit breaker.
+type State uint32
+
 const (
-	closed uint32 = iota
-	open
-	halfOpen
+	Closed State = iota
+	Open
+	HalfOpen
 )
 
 // Breaker implements the circuit-breaker resiliency pattern
@@ -24,7 +27,7 @@ type Breaker struct {
 	timeout                          time.Duration
 
 	lock              sync.Mutex
-	state             uint32
+	state             State
 	errors, successes int
 	lastError         time.Time
 }
@@ -46,9 +49,9 @@ func New(errorThreshold, successThreshold int, timeout time.Duration) *Breaker {
 // already open, or it will run the given function and pass along its return
 // value. It is safe to call Run concurrently on the same Breaker.
 func (b *Breaker) Run(work func() error) error {
-	state := atomic.LoadUint32(&b.state)
+	state := b.GetState()
 
-	if state == open {
+	if state == Open {
 		return ErrBreakerOpen
 	}
 
@@ -61,9 +64,9 @@ func (b *Breaker) Run(work func() error) error {
 // the return value of the function. It is safe to call Go concurrently on the
 // same Breaker.
 func (b *Breaker) Go(work func() error) error {
-	state := atomic.LoadUint32(&b.state)
+	state := b.GetState()
 
-	if state == open {
+	if state == Open {
 		return ErrBreakerOpen
 	}
 
@@ -75,7 +78,13 @@ func (b *Breaker) Go(work func() error) error {
 	return nil
 }
 
-func (b *Breaker) doWork(state uint32, work func() error) error {
+// GetState returns the current State of the circuit-breaker at the moment
+// that it is called.
+func (b *Breaker) GetState() State {
+	return (State)(atomic.LoadUint32((*uint32)(&b.state)))
+}
+
+func (b *Breaker) doWork(state State, work func() error) error {
 	var panicValue interface{}
 
 	result := func() error {
@@ -85,7 +94,7 @@ func (b *Breaker) doWork(state uint32, work func() error) error {
 		return work()
 	}()
 
-	if result == nil && panicValue == nil && state == closed {
+	if result == nil && panicValue == nil && state == Closed {
 		// short-circuit the normal, success path without contending
 		// on the lock
 		return nil
@@ -108,7 +117,7 @@ func (b *Breaker) processResult(result error, panicValue interface{}) {
 	defer b.lock.Unlock()
 
 	if result == nil && panicValue == nil {
-		if b.state == halfOpen {
+		if b.state == HalfOpen {
 			b.successes++
 			if b.successes == b.successThreshold {
 				b.closeBreaker()
@@ -123,26 +132,26 @@ func (b *Breaker) processResult(result error, panicValue interface{}) {
 		}
 
 		switch b.state {
-		case closed:
+		case Closed:
 			b.errors++
 			if b.errors == b.errorThreshold {
 				b.openBreaker()
 			} else {
 				b.lastError = time.Now()
 			}
-		case halfOpen:
+		case HalfOpen:
 			b.openBreaker()
 		}
 	}
 }
 
 func (b *Breaker) openBreaker() {
-	b.changeState(open)
+	b.changeState(Open)
 	go b.timer()
 }
 
 func (b *Breaker) closeBreaker() {
-	b.changeState(closed)
+	b.changeState(Closed)
 }
 
 func (b *Breaker) timer() {
@@ -151,11 +160,11 @@ func (b *Breaker) timer() {
 	b.lock.Lock()
 	defer b.lock.Unlock()
 
-	b.changeState(halfOpen)
+	b.changeState(HalfOpen)
 }
 
-func (b *Breaker) changeState(newState uint32) {
+func (b *Breaker) changeState(newState State) {
 	b.errors = 0
 	b.successes = 0
-	atomic.StoreUint32(&b.state, newState)
+	atomic.StoreUint32((*uint32)(&b.state), (uint32)(newState))
 }

breaker/breaker_test.go

@@ -22,24 +22,36 @@ func returnsSuccess() error {
 
 func TestBreakerErrorExpiry(t *testing.T) {
 	breaker := New(2, 1, 10*time.Millisecond)
+	if breaker.GetState() != Closed {
+		t.Error("incorrect state")
+	}
 
 	for i := 0; i < 3; i++ {
 		if err := breaker.Run(returnsError); err != errSomeError {
 			t.Error(err)
 		}
 		time.Sleep(10 * time.Millisecond)
 	}
+	if breaker.GetState() != Closed {
+		t.Error("incorrect state")
+	}
 
 	for i := 0; i < 3; i++ {
 		if err := breaker.Go(returnsError); err != nil {
 			t.Error(err)
 		}
 		time.Sleep(10 * time.Millisecond)
 	}
+	if breaker.GetState() != Closed {
+		t.Error("incorrect state")
+	}
 }
 
 func TestBreakerPanicsCountAsErrors(t *testing.T) {
 	breaker := New(3, 2, 1*time.Second)
+	if breaker.GetState() != Closed {
+		t.Error("incorrect state")
+	}
 
 	// three errors opens the breaker
 	for i := 0; i < 3; i++ {
@@ -58,6 +70,9 @@ func TestBreakerPanicsCountAsErrors(t *testing.T) {
 	}
 
 	// breaker is open
+	if breaker.GetState() != Open {
+		t.Error("incorrect state")
+	}
 	for i := 0; i < 5; i++ {
 		if err := breaker.Run(returnsError); err != ErrBreakerOpen {
 			t.Error(err)
@@ -67,6 +82,9 @@ func TestBreakerPanicsCountAsErrors(t *testing.T) {
 
 func TestBreakerStateTransitions(t *testing.T) {
 	breaker := New(3, 2, 10*time.Millisecond)
+	if breaker.GetState() != Closed {
+		t.Error("incorrect state")
+	}
 
 	// three errors opens the breaker
 	for i := 0; i < 3; i++ {
@@ -76,6 +94,9 @@ func TestBreakerStateTransitions(t *testing.T) {
 	}
 
 	// breaker is open
+	if breaker.GetState() != Open {
+		t.Error("incorrect state")
+	}
 	for i := 0; i < 5; i++ {
 		if err := breaker.Run(returnsError); err != ErrBreakerOpen {
 			t.Error(err)
@@ -84,6 +105,9 @@ func TestBreakerStateTransitions(t *testing.T) {
 
 	// wait for it to half-close
 	time.Sleep(20 * time.Millisecond)
+	if breaker.GetState() != HalfOpen {
+		t.Error("incorrect state")
+	}
 	// one success works, but is not enough to fully close
 	if err := breaker.Run(returnsSuccess); err != nil {
 		t.Error(err)
@@ -93,23 +117,35 @@ func TestBreakerStateTransitions(t *testing.T) {
 		t.Error(err)
 	}
 	// breaker is open
+	if breaker.GetState() != Open {
+		t.Error("incorrect state")
+	}
 	if err := breaker.Run(returnsError); err != ErrBreakerOpen {
 		t.Error(err)
 	}
 
 	// wait for it to half-close
 	time.Sleep(20 * time.Millisecond)
+	if breaker.GetState() != HalfOpen {
+		t.Error("incorrect state")
+	}
 	// two successes is enough to close it for good
 	for i := 0; i < 2; i++ {
 		if err := breaker.Run(returnsSuccess); err != nil {
 			t.Error(err)
 		}
 	}
+	if breaker.GetState() != Closed {
+		t.Error("incorrect state")
+	}
 	// error works
 	if err := breaker.Run(returnsError); err != errSomeError {
 		t.Error(err)
 	}
 	// breaker is still closed
+	if breaker.GetState() != Closed {
+		t.Error("incorrect state")
+	}
 	if err := breaker.Run(returnsSuccess); err != nil {
 		t.Error(err)
 	}