fix: broke select in stream retrier

TBD54566975 · Oct 26, 2023 · 41e4cd8 · 41e4cd8
1 parent f60aba9
commit 41e4cd8
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/backend/common/rpc/rpc.go b/backend/common/rpc/rpc.go
@@ -156,8 +156,10 @@ func RetryStreamingClientStream[Req, Resp any](
 				logger.Infof("Stream recovered")
 				errored = false
 			}
-			if _, ok := <-ctx.Done(); ok {
+			select {
+			case <-ctx.Done():
 				return
+			default:
 			}
 			retry.Reset()
 			logLevel = log.Warn
@@ -167,8 +169,12 @@ func RetryStreamingClientStream[Req, Resp any](
 		errored = true
 		delay := retry.Duration()
 		logger.Logf(logLevel, "Stream handler failed, retrying in %s: %s", delay, err)
-		if _, ok := <-ctx.Done(); ok {
+		select {
+		case <-ctx.Done():
 			return
+
+		case <-time.After(delay):
 		}
+
 	}
 }