From 091b7c6c02b83cb51ea5102e3ca838fbac4661b0 Mon Sep 17 00:00:00 2001
From: lhchavez <lhchavez@lhchavez.com>
Date: Tue, 21 Dec 2021 03:39:21 +0000
Subject: [PATCH] fix: Make SendRequestWithRetry check for canceled contexts
 twice

This change makes it deterministic to call SendRequestWithRetry with a
canceled context. This is going to be useful because some Cloud APIs
rely on the context being canceled to prevent some operations from
proceeding, like
[`storage.(*ObjectHandle).NewWriter`](https://pkg.go.dev/cloud.google.com/go/storage#ObjectHandle.NewWriter).

Fixes: #1358
---
 internal/gensupport/send.go      | 13 +++++++++++++
 internal/gensupport/send_test.go | 22 ++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/internal/gensupport/send.go b/internal/gensupport/send.go
index dab64aef367..2eff614f5d7 100644
--- a/internal/gensupport/send.go
+++ b/internal/gensupport/send.go
@@ -99,6 +99,19 @@ func sendAndRetry(ctx context.Context, client *http.Client, req *http.Request, r
 		case <-time.After(pause):
 		}
 
+		select {
+		case <-ctx.Done():
+			// Check for context cancellation once more. If more than one case in a
+			// select is satisfied at the same time, Go will choose one arbitrarily.
+			// That can cause an operation to go through even if the context was
+			// canceled before.
+			if err == nil {
+				err = ctx.Err()
+			}
+			return resp, err
+		default:
+		}
+
 		resp, err = client.Do(req.WithContext(ctx))
 
 		var status int
diff --git a/internal/gensupport/send_test.go b/internal/gensupport/send_test.go
index d6af483e66e..52c8ad09c36 100644
--- a/internal/gensupport/send_test.go
+++ b/internal/gensupport/send_test.go
@@ -6,6 +6,7 @@ package gensupport
 
 import (
 	"context"
+	"errors"
 	"net/http"
 	"testing"
 )
@@ -29,3 +30,24 @@ func TestSendRequestWithRetry(t *testing.T) {
 		t.Error("got nil, want error")
 	}
 }
+
+type brokenRoundTripper struct{}
+
+func (t *brokenRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	return nil, errors.New("this should not happen")
+}
+
+func TestCanceledContextDoesNotPerformRequest(t *testing.T) {
+	client := http.Client{
+		Transport: &brokenRoundTripper{},
+	}
+	for i := 0; i < 1000; i++ {
+		req, _ := http.NewRequest("GET", "url", nil)
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := SendRequestWithRetry(ctx, &client, req)
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("got %v, want %v", err, context.Canceled)
+		}
+	}
+}