From 4bbc6f59d95734cf06e181348e21bc080073002c Mon Sep 17 00:00:00 2001
From: Marc Cirauqui <marc@tetrate.io>
Date: Thu, 30 May 2024 11:41:53 +0200
Subject: [PATCH 1/2] Check for context cancellation in loop

---
 pkg/healthcheck/http.go                |  1 +
 pkg/helpers/retry.go                   | 22 +++++++++++++++++++---
 pkg/provider/resource_local_command.go |  1 +
 pkg/provider/resource_tcp_echo.go      |  1 +
 4 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/pkg/healthcheck/http.go b/pkg/healthcheck/http.go
index 392cf89..de2a107 100644
--- 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
index 75c01b5..375eb4d 100644
--- 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
index 5c67006..a09955f 100644
--- 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
index fb333f8..76dc6c6 100644
--- 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()),

From 0253ab78b454d90fe47e76d6c8ea652ec33e9cf8 Mon Sep 17 00:00:00 2001
From: Marc Cirauqui <marc@tetrate.io>
Date: Thu, 30 May 2024 11:52:07 +0200
Subject: [PATCH 2/2] use channel

---
 pkg/helpers/retry.go | 29 +++++++++++++----------------
 1 file changed, 13 insertions(+), 16 deletions(-)

diff --git a/pkg/helpers/retry.go b/pkg/helpers/retry.go
index 375eb4d..0f12347 100644
--- a/pkg/helpers/retry.go
+++ b/pkg/helpers/retry.go
@@ -42,27 +42,24 @@ func (r *RetryWindow) Do(action func(attempt int, successes int) bool) RetryResu
 		successCount := 0
 		// run a while true loop, exiting when the timeout expires
 		for {
-			attempt++
-			if action(attempt, successCount) {
-				successCount++
-				if successCount >= r.ConsecutiveSuccesses {
-					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{}{}
+			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
 				}
+				time.Sleep(r.Interval)
 			}
-			time.Sleep(r.Interval)
 		}
-
 	}()
 
 	select {