From 27909592b88ba7059d2645aa6e88018e472f652f Mon Sep 17 00:00:00 2001
From: ashmrtn <ashmrtnz@alcion.ai>
Date: Mon, 10 Apr 2023 11:38:14 -0700
Subject: [PATCH] Exit early if context is marked as done (#2883)

Keeps the code from blocking and then failing immediately without sending a retry.

---

#### Does this PR need a docs update or release note?

- [ ] :white_check_mark: Yes, it's included
- [ ] :clock1: Yes, but in a later PR
- [x] :no_entry: No

#### Type of change

- [ ] :sunflower: Feature
- [ ] :bug: Bugfix
- [ ] :world_map: Documentation
- [ ] :robot: Supportability/Tests
- [ ] :computer: CI/Deployment
- [x] :broom: Tech Debt/Cleanup

#### Issue(s)

* closes #2882

#### Test Plan

- [ ] :muscle: Manual
- [ ] :zap: Unit test
- [ ] :green_heart: E2E
---
 src/internal/connector/graph/retry_middleware.go | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/internal/connector/graph/retry_middleware.go b/src/internal/connector/graph/retry_middleware.go
index 2a4ebb73e..33c228523 100644
--- a/src/internal/connector/graph/retry_middleware.go
+++ b/src/internal/connector/graph/retry_middleware.go
@@ -42,9 +42,20 @@ func (middleware RetryHandler) retryRequest(
 
 		cumulativeDelay += delay
 
-		req.Header.Set(retryAttemptHeader, strconv.Itoa(executionCount))
+		timer := time.NewTimer(delay)
 
-		time.Sleep(delay)
+		select {
+		case <-ctx.Done():
+			// Don't retry if the context is marked as done, it will just error out
+			// when we attempt to send the retry anyway.
+			return resp, ctx.Err()
+
+		// Will exit switch-block so the remainder of the code doesn't need to be
+		// indented.
+		case <-timer.C:
+		}
+
+		req.Header.Set(retryAttemptHeader, strconv.Itoa(executionCount))
 
 		response, err := pipeline.Next(req, middlewareIndex)
 		if err != nil && !IsErrTimeout(err) {