Merge pull request #42 from matope/add-callcontext-on-breaker

Add Breaker.CallContext to handle context.Canceled

Author: rubyist

circuitbreaker.go

@@ -31,6 +31,7 @@
 package circuit
 
 import (
+	"context"
 	"errors"
 	"sync"
 	"sync/atomic"
@@ -329,6 +330,12 @@ func (cb *Breaker) Ready() bool {
 // whenever the function returns an error. If the called function takes longer
 // than timeout to run, a failure will be recorded.
 func (cb *Breaker) Call(circuit func() error, timeout time.Duration) error {
+	return cb.CallContext(context.Background(), circuit, timeout)
+}
+
+// CallContext is same as Call but if the ctx is canceled after the circuit returned an error,
+// the error will not be marked as a failure because the call was canceled intentionally.
+func (cb *Breaker) CallContext(ctx context.Context, circuit func() error, timeout time.Duration) error {
 	var err error
 
 	if !cb.Ready() {
@@ -353,7 +360,9 @@ func (cb *Breaker) Call(circuit func() error, timeout time.Duration) error {
 	}
 
 	if err != nil {
-		cb.Fail()
+		if ctx.Err() != context.Canceled {
+			cb.Fail()
+		}
 		return err
 	}
 

circuitbreaker_test.go

@@ -1,6 +1,7 @@
 package circuit
 
 import (
+	"context"
 	"fmt"
 	"sync/atomic"
 	"testing"
@@ -252,6 +253,42 @@ func TestThresholdBreakerCalling(t *testing.T) {
 	}
 }
 
+func TestThresholdBreakerCallingContext(t *testing.T) {
+	circuit := func() error {
+		return fmt.Errorf("error")
+	}
+
+	cb := NewThresholdBreaker(2)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	err := cb.CallContext(ctx, circuit, 0) // First failure
+	if err == nil {
+		t.Fatal("expected threshold breaker to error")
+	}
+	if cb.Tripped() {
+		t.Fatal("expected threshold breaker to be open")
+	}
+
+	// Cancel the next Call.
+	cancel()
+
+	err = cb.CallContext(ctx, circuit, 0) // Second failure but it's canceled
+	if err == nil {
+		t.Fatal("expected threshold breaker to error")
+	}
+	if cb.Tripped() {
+		t.Fatal("expected threshold breaker to be open")
+	}
+
+	err = cb.CallContext(context.Background(), circuit, 0) // Thirt failure trips
+	if err == nil {
+		t.Fatal("expected threshold breaker to error")
+	}
+	if !cb.Tripped() {
+		t.Fatal("expected threshold breaker to be tripped")
+	}
+}
+
 func TestThresholdBreakerResets(t *testing.T) {
 	called := 0
 	success := false