feat: add optional mutator to instancetype provider

aws · Sep 27, 2024 · 1554502 · 1554502
1 parent 7706b2d
commit 1554502
Showing 1 changed file with 30 additions and 12 deletions.
diff --git a/pkg/providers/instancetype/instancetype.go b/pkg/providers/instancetype/instancetype.go
@@ -21,6 +21,7 @@ import (
 	"sync"
 	"sync/atomic"
 
+	"github.com/awslabs/operatorpkg/option"
 	"github.com/mitchellh/hashstructure/v2"
 	"github.com/patrickmn/go-cache"
 	"github.com/prometheus/client_golang/prometheus"
@@ -54,7 +55,19 @@ type Provider interface {
 	UpdateInstanceTypeOfferings(ctx context.Context) error
 }
 
+type defaultProviderOptions struct {
+	mutator func(*cloudprovider.InstanceType, *v1.EC2NodeClass)
+}
+
+var WithInstanceTypeMutator = func(mutator func(*cloudprovider.InstanceType, *v1.EC2NodeClass)) func(*defaultProviderOptions) {
+	return func(opts *defaultProviderOptions) {
+		opts.mutator = mutator
+	}
+}
+
 type DefaultProvider struct {
+	defaultProviderOptions
+
 	region          string
 	ec2api          ec2iface.EC2API
 	subnetProvider  subnet.Provider
@@ -82,18 +95,19 @@ type DefaultProvider struct {
 }
 
 func NewDefaultProvider(region string, instanceTypesCache *cache.Cache, ec2api ec2iface.EC2API, subnetProvider subnet.Provider,
-	unavailableOfferingsCache *awscache.UnavailableOfferings, pricingProvider pricing.Provider) *DefaultProvider {
+	unavailableOfferingsCache *awscache.UnavailableOfferings, pricingProvider pricing.Provider, opts ...option.Function[defaultProviderOptions]) *DefaultProvider {
 	return &DefaultProvider{
-		ec2api:                ec2api,
-		region:                region,
-		subnetProvider:        subnetProvider,
-		pricingProvider:       pricingProvider,
-		instanceTypesInfo:     []*ec2.InstanceTypeInfo{},
-		instanceTypeOfferings: map[string]sets.Set[string]{},
-		instanceTypesCache:    instanceTypesCache,
-		unavailableOfferings:  unavailableOfferingsCache,
-		cm:                    pretty.NewChangeMonitor(),
-		instanceTypesSeqNum:   0,
+		defaultProviderOptions: *option.Resolve(opts...),
+		ec2api:                 ec2api,
+		region:                 region,
+		subnetProvider:         subnetProvider,
+		pricingProvider:        pricingProvider,
+		instanceTypesInfo:      []*ec2.InstanceTypeInfo{},
+		instanceTypeOfferings:  map[string]sets.Set[string]{},
+		instanceTypesCache:     instanceTypesCache,
+		unavailableOfferings:   unavailableOfferingsCache,
+		cm:                     pretty.NewChangeMonitor(),
+		instanceTypesSeqNum:    0,
 	}
 }
 
@@ -164,11 +178,15 @@ func (p *DefaultProvider) List(ctx context.Context, kc *v1.KubeletConfiguration,
 		// Any changes to the values passed into the NewInstanceType method will require making updates to the cache key
 		// so that Karpenter is able to cache the set of InstanceTypes based on values that alter the set of instance types
 		// !!! Important !!!
-		return NewInstanceType(ctx, i, p.region,
+		it := NewInstanceType(ctx, i, p.region,
 			nodeClass.Spec.BlockDeviceMappings, nodeClass.Spec.InstanceStorePolicy,
 			kc.MaxPods, kc.PodsPerCore, kc.KubeReserved, kc.SystemReserved, kc.EvictionHard, kc.EvictionSoft,
 			amiFamily, p.createOfferings(ctx, i, allZones, p.instanceTypeOfferings[aws.StringValue(i.InstanceType)], nodeClass.Status.Subnets),
 		)
+		if p.mutator != nil {
+			p.mutator(it, nodeClass)
+		}
+		return it
 	})
 	p.instanceTypesCache.SetDefault(key, result)
 	return result, nil