diff --git a/pkg/providers/instancetype/instancetype.go b/pkg/providers/instancetype/instancetype.go index 3607782b4482..3c0f67d081d7 100644 --- 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