update the code

Author: cqi1217

pkg/networking/az_id_translator.go

@@ -36,109 +36,90 @@ func NewDefaultAZIDTranslator(ec2Client services.EC2, logger logr.Logger) *defau
 
 var _ AZIDTranslator = &defaultAZIDTranslator{}
 
-// zoneCacheKey scopes a cached zone lookup to a single account.
-// Availability zone names are only meaningful within an account, so entries for different accounts
-// must never be shared. scope is empty for the cluster's own account and the assumed role ARN
-// otherwise.
-type zoneCacheKey struct {
-	scope  string
-	lookup string
+// zoneMapping holds every availability zone of a single account, indexed by name and by ID.
+type zoneMapping struct {
+	idByName map[string]string
+	nameByID map[string]string
 }
 
 type defaultAZIDTranslator struct {
 	ec2Client services.EC2
 
+	// zoneCache maps an account scope to its zoneMapping. Availability zone names are only
+	// meaningful within an account, so entries for different accounts must never be shared. The
+	// scope is empty for the cluster's own account and the assumed role ARN otherwise.
 	zoneCache      *cache.Expiring
 	zoneCacheMutex sync.RWMutex
 
 	logger logr.Logger
 }
 
 func (t *defaultAZIDTranslator) TranslateAZName(ctx context.Context, assumeRoleArn string, externalId string, srcZoneName string) (*string, error) {
-	zoneID, err := t.resolveZoneID(ctx, srcZoneName)
+	srcZones, err := t.fetchZoneMapping(ctx, t.ec2Client, "")
 	if err != nil {
 		return nil, err
 	}
-	if zoneID == nil {
+	zoneID, exists := srcZones.idByName[srcZoneName]
+	if !exists {
+		t.logger.Info("unable to resolve availability zone ID", "zoneName", srcZoneName)
 		return nil, nil
 	}
 
 	assumedRoleEC2, err := t.ec2Client.AssumeRole(ctx, assumeRoleArn, externalId)
 	if err != nil {
 		return nil, err
 	}
-	return t.resolveZoneName(ctx, assumedRoleEC2, assumeRoleArn, *zoneID)
-}
-
-// resolveZoneID resolves a zone name to its zone ID within the cluster's own account.
-func (t *defaultAZIDTranslator) resolveZoneID(ctx context.Context, zoneName string) (*string, error) {
-	cacheKey := zoneCacheKey{lookup: zoneName}
-	if cachedZoneID, exists := t.fetchZoneFromCache(cacheKey); exists {
-		return &cachedZoneID, nil
-	}
-
-	resp, err := t.ec2Client.DescribeAvailabilityZonesWithContext(ctx, &ec2sdk.DescribeAvailabilityZonesInput{
-		ZoneNames: []string{zoneName},
-	})
+	dstZones, err := t.fetchZoneMapping(ctx, assumedRoleEC2, assumeRoleArn)
 	if err != nil {
 		return nil, err
 	}
-	for _, azInfo := range resp.AvailabilityZones {
-		if awssdk.ToString(azInfo.ZoneName) != zoneName {
-			continue
-		}
-		zoneID := awssdk.ToString(azInfo.ZoneId)
-		if zoneID == "" {
-			continue
-		}
-		t.saveZoneToCache(cacheKey, zoneID)
-		return &zoneID, nil
+	zoneName, exists := dstZones.nameByID[zoneID]
+	if !exists {
+		t.logger.Info("unable to resolve availability zone name", "zoneID", zoneID, "scope", assumeRoleArn)
+		return nil, nil
 	}
-	t.logger.Info("unable to resolve availability zone ID", "zoneName", zoneName)
-	return nil, nil
+	return &zoneName, nil
 }
 
-// resolveZoneName resolves a zone ID to the zone name used by the account reachable via ec2Client.
-func (t *defaultAZIDTranslator) resolveZoneName(ctx context.Context, ec2Client services.EC2, scope string, zoneID string) (*string, error) {
-	cacheKey := zoneCacheKey{scope: scope, lookup: zoneID}
-	if cachedZoneName, exists := t.fetchZoneFromCache(cacheKey); exists {
-		return &cachedZoneName, nil
+// fetchZoneMapping returns the availability zones visible to ec2Client, cached under scope.
+// The zones of an account are fetched in a single unfiltered call so that a zone missing from the
+// mapping is known to be unresolvable without issuing another call for it.
+func (t *defaultAZIDTranslator) fetchZoneMapping(ctx context.Context, ec2Client services.EC2, scope string) (*zoneMapping, error) {
+	if cachedMapping, exists := t.fetchZoneMappingFromCache(scope); exists {
+		return cachedMapping, nil
 	}
 
-	resp, err := ec2Client.DescribeAvailabilityZonesWithContext(ctx, &ec2sdk.DescribeAvailabilityZonesInput{
-		ZoneIds: []string{zoneID},
-	})
+	resp, err := ec2Client.DescribeAvailabilityZonesWithContext(ctx, &ec2sdk.DescribeAvailabilityZonesInput{})
 	if err != nil {
 		return nil, err
 	}
+
+	mapping := &zoneMapping{
+		idByName: make(map[string]string, len(resp.AvailabilityZones)),
+		nameByID: make(map[string]string, len(resp.AvailabilityZones)),
+	}
 	for _, azInfo := range resp.AvailabilityZones {
-		if awssdk.ToString(azInfo.ZoneId) != zoneID {
-			continue
-		}
 		zoneName := awssdk.ToString(azInfo.ZoneName)
-		if zoneName == "" {
+		zoneID := awssdk.ToString(azInfo.ZoneId)
+		if zoneName == "" || zoneID == "" {
 			continue
 		}
-		t.saveZoneToCache(cacheKey, zoneName)
-		return &zoneName, nil
+		mapping.idByName[zoneName] = zoneID
+		mapping.nameByID[zoneID] = zoneName
 	}
-	t.logger.Info("unable to resolve availability zone name", "zoneID", zoneID, "scope", scope)
-	return nil, nil
+
+	t.zoneCacheMutex.Lock()
+	defer t.zoneCacheMutex.Unlock()
+	t.zoneCache.Set(scope, mapping, defaultAZIDTranslationCacheTTL)
+	return mapping, nil
 }
 
-func (t *defaultAZIDTranslator) fetchZoneFromCache(key zoneCacheKey) (string, bool) {
+func (t *defaultAZIDTranslator) fetchZoneMappingFromCache(scope string) (*zoneMapping, bool) {
 	t.zoneCacheMutex.RLock()
 	defer t.zoneCacheMutex.RUnlock()
 
-	if rawCacheItem, exists := t.zoneCache.Get(key); exists {
-		return rawCacheItem.(string), true
+	if rawCacheItem, exists := t.zoneCache.Get(scope); exists {
+		return rawCacheItem.(*zoneMapping), true
 	}
-	return "", false
-}
-
-func (t *defaultAZIDTranslator) saveZoneToCache(key zoneCacheKey, value string) {
-	t.zoneCacheMutex.Lock()
-	defer t.zoneCacheMutex.Unlock()
-
-	t.zoneCache.Set(key, value, defaultAZIDTranslationCacheTTL)
+	return nil, false
 }