feat(compute/metadata): allow canceling GCE detection (#11786)

* Add a new `OnGCEWithContext` function that allows passing in
a cancellable context when performing runtime GCE platform detection.

* Enable the caller to pass in a context. Previously there was a possibility that
during GCE platform detection, the metadata address and the DNS resolution
steps could simultaneously block, which prevents the function from returning,
or can delay its return significantly.  Context mitigates this as the caller can
specify a deadline on the context or otherwise manage its cancellation.

Author: diurnalist

compute/metadata/metadata.go

@@ -117,80 +117,18 @@ var (
 // NOTE: True returned from `OnGCE` does not guarantee that the metadata server
 // is accessible from this process and have all the metadata defined.
 func OnGCE() bool {
-	onGCEOnce.Do(initOnGCE)
-	return onGCE
-}
-
-func initOnGCE() {
-	onGCE = testOnGCE()
+	return OnGCEWithContext(context.Background())
 }
 
-func testOnGCE() bool {
-	// The user explicitly said they're on GCE, so trust them.
-	if os.Getenv(metadataHostEnv) != "" {
-		return true
-	}
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	resc := make(chan bool, 2)
-
-	// Try two strategies in parallel.
-	// See https://github.com/googleapis/google-cloud-go/issues/194
-	go func() {
-		req, _ := http.NewRequest("GET", "http://"+metadataIP, nil)
-		req.Header.Set("User-Agent", userAgent)
-		res, err := newDefaultHTTPClient().Do(req.WithContext(ctx))
-		if err != nil {
-			resc <- false
-			return
-		}
-		defer res.Body.Close()
-		resc <- res.Header.Get("Metadata-Flavor") == "Google"
-	}()
-
-	go func() {
-		resolver := &net.Resolver{}
-		addrs, err := resolver.LookupHost(ctx, "metadata.google.internal.")
-		if err != nil || len(addrs) == 0 {
-			resc <- false
-			return
-		}
-		resc <- strsContains(addrs, metadataIP)
-	}()
-
-	tryHarder := systemInfoSuggestsGCE()
-	if tryHarder {
-		res := <-resc
-		if res {
-			// The first strategy succeeded, so let's use it.
-			return true
-		}
-		// Wait for either the DNS or metadata server probe to
-		// contradict the other one and say we are running on
-		// GCE. Give it a lot of time to do so, since the system
-		// info already suggests we're running on a GCE BIOS.
-		timer := time.NewTimer(5 * time.Second)
-		defer timer.Stop()
-		select {
-		case res = <-resc:
-			return res
-		case <-timer.C:
-			// Too slow. Who knows what this system is.
-			return false
-		}
-	}
-
-	// There's no hint from the system info that we're running on
-	// GCE, so use the first probe's result as truth, whether it's
-	// true or false. The goal here is to optimize for speed for
-	// users who are NOT running on GCE. We can't assume that
-	// either a DNS lookup or an HTTP request to a blackholed IP
-	// address is fast. Worst case this should return when the
-	// metaClient's Transport.ResponseHeaderTimeout or
-	// Transport.Dial.Timeout fires (in two seconds).
-	return <-resc
+// OnGCEWithContext reports whether this process is running on Google Compute Platforms.
+// This function's return value is memoized for better performance.
+// NOTE: True returned from `OnGCEWithContext` does not guarantee that the metadata server
+// is accessible from this process and have all the metadata defined.
+func OnGCEWithContext(ctx context.Context) bool {
+	onGCEOnce.Do(func() {
+		onGCE = defaultClient.OnGCEWithContext(ctx)
+	})
+	return onGCE
 }
 
 // Subscribe calls Client.SubscribeWithContext on the default client.
@@ -450,6 +388,84 @@ func NewWithOptions(opts *Options) *Client {
 	return &Client{hc: client, logger: logger}
 }
 
+// NOTE: metadataRequestStrategy is assigned to a variable for test stubbing purposes.
+var metadataRequestStrategy = func(ctx context.Context, httpClient *http.Client, resc chan bool) {
+	req, _ := http.NewRequest("GET", "http://"+metadataIP, nil)
+	req.Header.Set("User-Agent", userAgent)
+	res, err := httpClient.Do(req.WithContext(ctx))
+	if err != nil {
+		resc <- false
+		return
+	}
+	defer res.Body.Close()
+	resc <- res.Header.Get("Metadata-Flavor") == "Google"
+}
+
+// NOTE: dnsRequestStrategy is assigned to a variable for test stubbing purposes.
+var dnsRequestStrategy = func(ctx context.Context, resc chan bool) {
+	resolver := &net.Resolver{}
+	addrs, err := resolver.LookupHost(ctx, "metadata.google.internal.")
+	if err != nil || len(addrs) == 0 {
+		resc <- false
+		return
+	}
+	resc <- strsContains(addrs, metadataIP)
+}
+
+// OnGCEWithContext reports whether this process is running on Google Compute Platforms.
+// NOTE: True returned from `OnGCEWithContext` does not guarantee that the metadata server
+// is accessible from this process and have all the metadata defined.
+func (c *Client) OnGCEWithContext(ctx context.Context) bool {
+	// The user explicitly said they're on GCE, so trust them.
+	if os.Getenv(metadataHostEnv) != "" {
+		return true
+	}
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	resc := make(chan bool, 2)
+
+	// Try two strategies in parallel.
+	// See https://github.com/googleapis/google-cloud-go/issues/194
+	go metadataRequestStrategy(ctx, c.hc, resc)
+	go dnsRequestStrategy(ctx, resc)
+
+	tryHarder := systemInfoSuggestsGCE()
+	if tryHarder {
+		res := <-resc
+		if res {
+			// The first strategy succeeded, so let's use it.
+			return true
+		}
+
+		// Wait for either the DNS or metadata server probe to
+		// contradict the other one and say we are running on
+		// GCE. Give it a lot of time to do so, since the system
+		// info already suggests we're running on a GCE BIOS.
+		// Ensure cancellations from the calling context are respected.
+		waitContext, cancelWait := context.WithTimeout(ctx, 5*time.Second)
+		defer cancelWait()
+		select {
+		case res = <-resc:
+			return res
+		case <-waitContext.Done():
+			// Too slow. Who knows what this system is.
+			return false
+		}
+	}
+
+	// There's no hint from the system info that we're running on
+	// GCE, so use the first probe's result as truth, whether it's
+	// true or false. The goal here is to optimize for speed for
+	// users who are NOT running on GCE. We can't assume that
+	// either a DNS lookup or an HTTP request to a blackholed IP
+	// address is fast. Worst case this should return when the
+	// metaClient's Transport.ResponseHeaderTimeout or
+	// Transport.Dial.Timeout fires (in two seconds).
+	return <-resc
+}
+
 // getETag returns a value from the metadata service as well as the associated ETag.
 // This func is otherwise equivalent to Get.
 func (c *Client) getETag(ctx context.Context, suffix string) (value, etag string, err error) {

compute/metadata/metadata_test.go

@@ -27,14 +27,15 @@ import (
 )
 
 func TestOnGCE_Stress(t *testing.T) {
+	ctx := context.Background()
 	if testing.Short() {
 		t.Skip("skipping in -short mode")
 	}
 	var last bool
 	for i := 0; i < 100; i++ {
 		onGCEOnce = sync.Once{}
 
-		now := OnGCE()
+		now := OnGCEWithContext(ctx)
 		if i > 0 && now != last {
 			t.Errorf("%d. changed from %v to %v", i, last, now)
 		}
@@ -44,15 +45,70 @@ func TestOnGCE_Stress(t *testing.T) {
 }
 
 func TestOnGCE_Force(t *testing.T) {
+	ctx := context.Background()
 	onGCEOnce = sync.Once{}
 	old := os.Getenv(metadataHostEnv)
 	defer os.Setenv(metadataHostEnv, old)
 	os.Setenv(metadataHostEnv, "127.0.0.1")
-	if !OnGCE() {
+	if !OnGCEWithContext(ctx) {
 		t.Error("OnGCE() = false; want true")
 	}
 }
 
+func TestOnGCE_Cancel(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	onGCEOnce = sync.Once{}
+	if OnGCEWithContext(ctx) {
+		t.Error("OnGCE() = true; want false")
+	}
+}
+
+func TestOnGCE_CancelTryHarder(t *testing.T) {
+	// If system info suggests GCE, we allow extra time for the
+	// probe with higher latency (HTTP or DNS) to return. In this
+	// test, the system info suggest GCE, the DNS probe fails
+	// immediately, and the HTTP probe would succeed after 750ms.
+	// However, the user-provided context deadline is 500ms. GCE
+	// detection should fail, respecting the provided context.
+	//
+	// NOTE: This code could create a data race if tests are run
+	// in parallel.
+	origSystemInfoSuggestsGCE := systemInfoSuggestsGCE
+	origMetadataRequestStrategy := metadataRequestStrategy
+	origDNSRequestStrategy := dnsRequestStrategy
+	systemInfoSuggestsGCE = func() bool { return true }
+	metadataRequestStrategy = func(_ context.Context, _ *http.Client, resc chan bool) {
+		time.Sleep(750 * time.Millisecond)
+		resc <- true
+	}
+	dnsRequestStrategy = func(_ context.Context, resc chan bool) {
+		resc <- false
+	}
+	defer func() {
+		systemInfoSuggestsGCE = origSystemInfoSuggestsGCE
+		metadataRequestStrategy = origMetadataRequestStrategy
+		dnsRequestStrategy = origDNSRequestStrategy
+	}()
+
+	// Set deadline upper-limit to 500ms
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+
+	// Set HTTP deadline to 1s
+	c := NewClient(&http.Client{Transport: sleepyTransport{1 * time.Second}})
+
+	start := time.Now()
+	if c.OnGCEWithContext(ctx) {
+		t.Error("OnGCE() = true; want false")
+	}
+
+	// Should have returned around 500ms, but account for some scheduling budget
+	if time.Now().Sub(start) > 510*time.Millisecond {
+		t.Error("OnGCE() did not return within deadline")
+	}
+}
+
 func TestOverrideUserAgent(t *testing.T) {
 	ctx := context.Background()
 	const userAgent = "my-user-agent"
@@ -214,7 +270,7 @@ func TestClientGetWithContext(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			ctx, cancel := context.WithTimeout(context.Background(), tc.ctxTimeout)
 			defer cancel()
-			c := NewClient(&http.Client{Transport: sleepyTransport{}})
+			c := NewClient(&http.Client{Transport: sleepyTransport{500 * time.Millisecond}})
 			_, err := c.GetWithContext(ctx, "foo")
 			if tc.wantErr && err == nil {
 				t.Fatal("c.GetWithContext() == nil, want an error")
@@ -227,14 +283,15 @@ func TestClientGetWithContext(t *testing.T) {
 }
 
 type sleepyTransport struct {
+	delay time.Duration
 }
 
 func (s sleepyTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	req.Context().Done()
 	select {
 	case <-req.Context().Done():
 		return nil, req.Context().Err()
-	case <-time.After(500 * time.Millisecond):
+	case <-time.After(s.delay):
 	}
 	return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("I woke up"))}, nil
 }

compute/metadata/syscheck.go

@@ -20,7 +20,9 @@ package metadata
 // doing network requests) suggests that we're running on GCE. If this
 // returns true, testOnGCE tries a bit harder to reach its metadata
 // server.
-func systemInfoSuggestsGCE() bool {
+//
+// NOTE: systemInfoSuggestsGCE is assigned to a varible for test stubbing purposes.
+var systemInfoSuggestsGCE = func() bool {
 	// We don't currently have checks for other GOOS
 	return false
 }

compute/metadata/syscheck_linux.go

@@ -21,8 +21,10 @@ import (
 	"strings"
 )
 
-func systemInfoSuggestsGCE() bool {
+// NOTE: systemInfoSuggestsGCE is assigned to a varible for test stubbing purposes.
+var systemInfoSuggestsGCE = func() bool {
 	b, _ := os.ReadFile("/sys/class/dmi/id/product_name")
+
 	name := strings.TrimSpace(string(b))
 	return name == "Google" || name == "Google Compute Engine"
 }

compute/metadata/syscheck_windows.go

@@ -22,7 +22,8 @@ import (
 	"golang.org/x/sys/windows/registry"
 )
 
-func systemInfoSuggestsGCE() bool {
+// NOTE: systemInfoSuggestsGCE is assigned to a varible for test stubbing purposes.
+var systemInfoSuggestsGCE = func() bool {
 	k, err := registry.OpenKey(registry.LOCAL_MACHINE, `SYSTEM\HardwareConfig\Current`, registry.QUERY_VALUE)
 	if err != nil {
 		return false