fix: msk collector fanout by region (#900)

Author: stephan-rayner

pkg/aws/msk/msk.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log/slog"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -147,13 +148,8 @@ func New(ctx context.Context, config *Config) (*Collector, error) {
 func (c *Collector) Collect(ctx context.Context, ch chan<- prometheus.Metric) error {
 	snapshot := c.pricingStore.Snapshot()
 
+	var wg sync.WaitGroup
 	for _, region := range c.regions {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		default:
-		}
-
 		if region.RegionName == nil || *region.RegionName == "" {
 			c.logger.Warn("skipping region with empty name")
 			continue
@@ -166,18 +162,21 @@ func (c *Collector) Collect(ctx context.Context, ch chan<- prometheus.Metric) er
 			continue
 		}
 
-		clusters, err := regionClient.ListMSKClusters(ctx)
-		if err != nil {
-			c.logger.Error("error listing MSK clusters", "region", regionName, "error", err)
-			continue
-		}
-
-		for _, cluster := range clusters {
-			c.collectCluster(ch, snapshot, regionName, cluster)
-		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			clusters, err := regionClient.ListMSKClusters(ctx)
+			if err != nil {
+				c.logger.Error("error listing MSK clusters", "region", regionName, "error", err)
+				return
+			}
+			for _, cluster := range clusters {
+				c.collectCluster(ch, snapshot, regionName, cluster)
+			}
+		}()
 	}
-
-	return nil
+	wg.Wait()
+	return ctx.Err()
 }
 
 func (c *Collector) collectCluster(ch chan<- prometheus.Metric, snapshot pricingstore.Snapshot, region string, cluster msktypes.Cluster) {

pkg/aws/msk/msk_test.go

@@ -1,6 +1,7 @@
 package msk
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"log/slog"
@@ -285,6 +286,82 @@ func TestCollectorCollectContinuesWhenRegionListingFails(t *testing.T) {
 	assert.Len(t, results, 2)
 }
 
+func TestCollectorCollectReturnsContextErrWhenContextCancelled(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	regionClient := mockclient.NewMockClient(ctrl)
+	pricingClient := mockclient.NewMockClient(ctrl)
+
+	regionClient.EXPECT().
+		ListMSKClusters(gomock.Any()).
+		Return(nil, context.Canceled).
+		AnyTimes()
+	expectPricingLoad(pricingClient, "us-east-1", "USE1", "0.2100000000", "0.1000000000")
+
+	collector, err := New(t.Context(), &Config{
+		Regions:   []ec2types.Region{{RegionName: aws.String("us-east-1")}},
+		RegionMap: map[string]client.Client{"us-east-1": regionClient},
+		Client:    pricingClient,
+		Logger:    testLogger(),
+		AccountID: "123456789012",
+	})
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	ch := make(chan prometheus.Metric, 10)
+	err = collector.Collect(ctx, ch)
+	close(ch)
+
+	assert.ErrorIs(t, err, context.Canceled)
+}
+
+func TestCollectorCollectContinuesOnContextDeadlineExceeded(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	failingClient := mockclient.NewMockClient(ctrl)
+	healthyClient := mockclient.NewMockClient(ctrl)
+	pricingClient := mockclient.NewMockClient(ctrl)
+	cluster := newProvisionedCluster(
+		"test-cluster",
+		"arn:aws:kafka:us-west-2:123456789012:cluster/test-cluster",
+		"kafka.m5.large", 3, 100,
+	)
+
+	failingClient.EXPECT().
+		ListMSKClusters(gomock.Any()).
+		Return(nil, context.DeadlineExceeded).
+		Times(1)
+	healthyClient.EXPECT().
+		ListMSKClusters(gomock.Any()).
+		Return([]msktypes.Cluster{cluster}, nil).
+		Times(1)
+	expectPricingLoad(pricingClient, "us-east-1", "USE1", "0.2100000000", "0.1000000000")
+	expectPricingLoad(pricingClient, "us-west-2", "USW2", "0.2100000000", "0.1000000000")
+
+	collector, err := New(t.Context(), &Config{
+		Regions: []ec2types.Region{
+			{RegionName: aws.String("us-east-1")},
+			{RegionName: aws.String("us-west-2")},
+		},
+		RegionMap: map[string]client.Client{
+			"us-east-1": failingClient,
+			"us-west-2": healthyClient,
+		},
+		Client:    pricingClient,
+		Logger:    testLogger(),
+		AccountID: "123456789012",
+	})
+	require.NoError(t, err)
+
+	results, err := collectMetricResults(t, collector)
+	require.NoError(t, err)
+	assert.Len(t, results, 2) // metrics from us-west-2 still collected
+}
+
 func newProvisionedCluster(name, arn, instanceType string, brokerCount, volumeSizeGiB int32) msktypes.Cluster {
 	return msktypes.Cluster{
 		ClusterArn:  aws.String(arn),