Merge pull request #24 from tetratelabs/add-ctx

Check for context cancellation in loop
tetratelabs · May 30, 2024 · c0284d3 · c0284d3
2 parents 773c935 + 0253ab7
commit c0284d3
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 12 deletions.
diff --git a/pkg/healthcheck/http.go b/pkg/healthcheck/http.go
@@ -88,6 +88,7 @@ func HealthCheck(ctx context.Context, data *HttpHealthArgs, diag *diag.Diagnosti
 	}
 
 	window := helpers.RetryWindow{
+		Context:              ctx,
 		Timeout:              time.Duration(data.Timeout) * time.Millisecond,
 		Interval:             time.Duration(data.Interval) * time.Millisecond,
 		ConsecutiveSuccesses: int(data.ConsecutiveSuccesses),

diff --git a/pkg/helpers/retry.go b/pkg/helpers/retry.go
@@ -14,9 +14,13 @@
 
 package helpers
 
-import "time"
+import (
+	"context"
+	"time"
+)
 
 type RetryWindow struct {
+	Context              context.Context
 	Timeout              time.Duration
 	Interval             time.Duration
 	ConsecutiveSuccesses int
@@ -27,33 +31,42 @@ type RetryResult int
 const (
 	Success RetryResult = iota
 	TimeoutExceeded
+	Failure
 )
 
 func (r *RetryWindow) Do(action func(attempt int, successes int) bool) RetryResult {
-	success := make(chan bool)
+	success := make(chan struct{})
+	failure := make(chan struct{})
 	go func() {
 		attempt := 0
 		successCount := 0
 		// run a while true loop, exiting when the timeout expires
 		for {
-			attempt++
-			if action(attempt, successCount) {
-				successCount++
-				if successCount >= r.ConsecutiveSuccesses {
-					success <- true
-					return
+			select {
+			case <-r.Context.Done():
+				failure <- struct{}{}
+				return
+			default:
+				attempt++
+				if action(attempt, successCount) {
+					successCount++
+					if successCount >= r.ConsecutiveSuccesses {
+						success <- struct{}{}
+						return
+					}
+				} else {
+					successCount = 0
 				}
-			} else {
-				successCount = 0
+				time.Sleep(r.Interval)
 			}
-			time.Sleep(r.Interval)
 		}
-
 	}()
 
 	select {
 	case <-success:
 		return Success
+	case <-failure:
+		return Failure
 	case <-time.After(r.Timeout):
 		return TimeoutExceeded
 	}

diff --git a/pkg/provider/resource_local_command.go b/pkg/provider/resource_local_command.go
@@ -202,6 +202,7 @@ func (r *LocalCommandResource) RunCommand(ctx context.Context, data *LocalComman
 	data.Stderr = types.StringNull()
 
 	window := helpers.RetryWindow{
+		Context:              ctx,
 		Timeout:              time.Duration(data.Timeout.ValueInt64()) * time.Millisecond,
 		Interval:             time.Duration(data.Interval.ValueInt64()) * time.Millisecond,
 		ConsecutiveSuccesses: int(data.ConsecutiveSuccesses.ValueInt64()),

diff --git a/pkg/provider/resource_tcp_echo.go b/pkg/provider/resource_tcp_echo.go
@@ -184,6 +184,7 @@ func (r *TCPEchoResource) TCPEcho(ctx context.Context, data *TCPEchoResourceMode
 	data.Passed = types.BoolValue(false)
 
 	window := helpers.RetryWindow{
+		Context:              ctx,
 		Timeout:              time.Duration(data.Timeout.ValueInt64()) * time.Millisecond,
 		Interval:             time.Duration(data.Interval.ValueInt64()) * time.Millisecond,
 		ConsecutiveSuccesses: int(data.ConsecutiveSuccesses.ValueInt64()),