From fe831e4f82748280199d8061f0a1f40a33f6bd2d Mon Sep 17 00:00:00 2001
From: Vasil Averyanau <vasil.averyanau@scylladb.com>
Date: Wed, 18 Dec 2024 13:45:15 +0100
Subject: [PATCH] fix(cloudmeta): fixes handling of context cancellation.

This fixes the issue when context that was passed to GetInstanceMetadata is
canceled before any of provider's functions returned.
---
 pkg/cloudmeta/metadata.go      | 14 +++++++++-----
 pkg/cloudmeta/metadata_test.go | 30 ++++++++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/pkg/cloudmeta/metadata.go b/pkg/cloudmeta/metadata.go
index 15aee9a17..785d8ecad 100644
--- a/pkg/cloudmeta/metadata.go
+++ b/pkg/cloudmeta/metadata.go
@@ -98,12 +98,16 @@ func (cloud *CloudMeta) GetInstanceMetadata(ctx context.Context) (InstanceMetada
 	// Return the first non error result or wait until all providers return err.
 	var mErr error
 	for range len(cloud.providers) {
-		res := <-results
-		if res.err != nil {
-			mErr = multierr.Append(mErr, res.err)
-			continue
+		select {
+		case <-ctx.Done():
+			return InstanceMetadata{}, ctx.Err()
+		case res := <-results:
+			if res.err != nil {
+				mErr = multierr.Append(mErr, res.err)
+				continue
+			}
+			return res.meta, nil
 		}
-		return res.meta, nil
 	}
 	return InstanceMetadata{}, mErr
 }
diff --git a/pkg/cloudmeta/metadata_test.go b/pkg/cloudmeta/metadata_test.go
index a2d5b782d..21a85adcd 100644
--- a/pkg/cloudmeta/metadata_test.go
+++ b/pkg/cloudmeta/metadata_test.go
@@ -77,7 +77,8 @@ func TestGetInstanceMetadata(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			cloudmeta := &CloudMeta{
-				providers: tc.providers,
+				providers:       tc.providers,
+				providerTimeout: 1 * time.Second,
 			}
 
 			meta, err := cloudmeta.GetInstanceMetadata(context.Background())
@@ -101,6 +102,31 @@ func TestGetInstanceMetadata(t *testing.T) {
 	}
 }
 
+func TestGetInstanceMetadataWithCancelledContext(t *testing.T) {
+	cloudmeta := &CloudMeta{
+		providers: []CloudMetadataProvider{
+			newTestProvider(t, "test_provider_1", "x-test-1", 1*time.Second, nil),
+		},
+		providerTimeout: 100 * time.Millisecond,
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	meta, err := cloudmeta.GetInstanceMetadata(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context.Canceled, got %v", err)
+	}
+
+	if meta.CloudProvider != "" {
+		t.Fatalf("meta.CloudProvider should be empty, got %s", meta.CloudProvider)
+	}
+
+	if meta.InstanceType != "" {
+		t.Fatalf("meta.InstanceType should be empty, got %s", meta.InstanceType)
+	}
+}
+
 func newTestProvider(t *testing.T, providerName, instanceType string, latency time.Duration, err error) *testProvider {
 	t.Helper()
 
@@ -128,5 +154,5 @@ func (tp testProvider) Metadata(ctx context.Context) (InstanceMetadata, error) {
 	return InstanceMetadata{
 		CloudProvider: tp.name,
 		InstanceType:  tp.instanceType,
-	}, nil
+	}, ctx.Err()
 }