Check for context cancellation in loop

tetratelabs · May 30, 2024 · 4bbc6f5 · 4bbc6f5
1 parent 773c935
commit 4bbc6f5
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 3 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,10 +31,12 @@ 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
@@ -40,12 +46,20 @@ func (r *RetryWindow) Do(action func(attempt int, successes int) bool) RetryResu
 			if action(attempt, successCount) {
 				successCount++
 				if successCount >= r.ConsecutiveSuccesses {
-					success <- true
+					success <- struct{}{}
 					return
 				}
 			} else {
 				successCount = 0
 			}
+			if r.Context != nil {
+				if err := r.Context.Err(); err != nil {
+					if err == context.Canceled || err == context.DeadlineExceeded {
+						failure <- struct{}{}
+						return
+					}
+				}
+			}
 			time.Sleep(r.Interval)
 		}
 
@@ -54,6 +68,8 @@ func (r *RetryWindow) Do(action func(attempt int, successes int) bool) RetryResu
 	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()),