perf: Reduce memory by using DescribeImages paginated call (#4546)

aws · Sep 5, 2023 · 3552184 · 3552184
1 parent 8159aea
commit 3552184
Show file tree

Hide file tree

Showing 7 changed files with 165 additions and 153 deletions.
diff --git a/pkg/cloudprovider/drift.go b/pkg/cloudprovider/drift.go
@@ -84,9 +84,9 @@ func (c *CloudProvider) isAMIDrifted(ctx context.Context, nodeClaim *corev1beta1
 	if len(amis) == 0 {
 		return "", fmt.Errorf("no amis exist given constraints")
 	}
-	mappedAMIs := amifamily.MapInstanceTypes(amis, []*cloudprovider.InstanceType{nodeInstanceType})
+	mappedAMIs := amis.MapToInstanceTypes([]*cloudprovider.InstanceType{nodeInstanceType})
 	if len(mappedAMIs) == 0 {
-		return "", fmt.Errorf("no instance types satisfy requirements of amis %v,", amis)
+		return "", fmt.Errorf("no instance types satisfy requirements of amis %v", amis)
 	}
 	if !lo.Contains(lo.Keys(mappedAMIs), instance.ImageID) {
 		return AMIDrift, nil

diff --git a/pkg/fake/ec2api.go b/pkg/fake/ec2api.go
@@ -343,6 +343,15 @@ func (e *EC2API) DescribeImagesWithContext(_ context.Context, input *ec2.Describ
 	}, nil
 }
 
+func (e *EC2API) DescribeImagesPagesWithContext(ctx context.Context, input *ec2.DescribeImagesInput, fn func(*ec2.DescribeImagesOutput, bool) bool, _ ...request.Option) error {
+	out, err := e.DescribeImagesWithContext(ctx, input)
+	if err != nil {
+		return err
+	}
+	fn(out, false)
+	return nil
+}
+
 func (e *EC2API) DescribeLaunchTemplatesWithContext(_ context.Context, input *ec2.DescribeLaunchTemplatesInput, _ ...request.Option) (*ec2.DescribeLaunchTemplatesOutput, error) {
 	if !e.NextError.IsNil() {
 		defer e.NextError.Reset()
@@ -471,29 +480,35 @@ func (e *EC2API) DescribeAvailabilityZonesWithContext(context.Context, *ec2.Desc
 	}}, nil
 }
 
-func (e *EC2API) DescribeInstanceTypesPagesWithContext(_ context.Context, _ *ec2.DescribeInstanceTypesInput, fn func(*ec2.DescribeInstanceTypesOutput, bool) bool, _ ...request.Option) error {
+func (e *EC2API) DescribeInstanceTypesWithContext(_ context.Context, _ *ec2.DescribeInstanceTypesInput, _ ...request.Option) (*ec2.DescribeInstanceTypesOutput, error) {
 	if !e.NextError.IsNil() {
 		defer e.NextError.Reset()
-		return e.NextError.Get()
+		return nil, e.NextError.Get()
 	}
 	if !e.DescribeInstanceTypesOutput.IsNil() {
-		fn(e.DescribeInstanceTypesOutput.Clone(), false)
-		return nil
+		return e.DescribeInstanceTypesOutput.Clone(), nil
 	}
-	fn(defaultDescribeInstanceTypesOutput, false)
+	return defaultDescribeInstanceTypesOutput, nil
+}
+
+func (e *EC2API) DescribeInstanceTypesPagesWithContext(ctx context.Context, input *ec2.DescribeInstanceTypesInput, fn func(*ec2.DescribeInstanceTypesOutput, bool) bool, _ ...request.Option) error {
+	out, err := e.DescribeInstanceTypesWithContext(ctx, input)
+	if err != nil {
+		return err
+	}
+	fn(out, false)
 	return nil
 }
 
-func (e *EC2API) DescribeInstanceTypeOfferingsPagesWithContext(_ context.Context, _ *ec2.DescribeInstanceTypeOfferingsInput, fn func(*ec2.DescribeInstanceTypeOfferingsOutput, bool) bool, _ ...request.Option) error {
+func (e *EC2API) DescribeInstanceTypeOfferingsWithContext(_ context.Context, _ *ec2.DescribeInstanceTypeOfferingsInput, _ ...request.Option) (*ec2.DescribeInstanceTypeOfferingsOutput, error) {
 	if !e.NextError.IsNil() {
 		defer e.NextError.Reset()
-		return e.NextError.Get()
+		return nil, e.NextError.Get()
 	}
 	if !e.DescribeInstanceTypeOfferingsOutput.IsNil() {
-		fn(e.DescribeInstanceTypeOfferingsOutput.Clone(), false)
-		return nil
+		return e.DescribeInstanceTypeOfferingsOutput.Clone(), nil
 	}
-	fn(&ec2.DescribeInstanceTypeOfferingsOutput{
+	return &ec2.DescribeInstanceTypeOfferingsOutput{
 		InstanceTypeOfferings: []*ec2.InstanceTypeOffering{
 			{
 				InstanceType: aws.String("m5.large"),
@@ -588,20 +603,36 @@ func (e *EC2API) DescribeInstanceTypeOfferingsPagesWithContext(_ context.Context
 				Location:     aws.String("test-zone-1c"),
 			},
 		},
-	}, false)
+	}, nil
+}
+
+func (e *EC2API) DescribeInstanceTypeOfferingsPagesWithContext(ctx context.Context, input *ec2.DescribeInstanceTypeOfferingsInput, fn func(*ec2.DescribeInstanceTypeOfferingsOutput, bool) bool, _ ...request.Option) error {
+	out, err := e.DescribeInstanceTypeOfferingsWithContext(ctx, input)
+	if err != nil {
+		return err
+	}
+	fn(out, false)
 	return nil
 }
 
-func (e *EC2API) DescribeSpotPriceHistoryPagesWithContext(_ aws.Context, in *ec2.DescribeSpotPriceHistoryInput, fn func(*ec2.DescribeSpotPriceHistoryOutput, bool) bool, _ ...request.Option) error {
-	e.DescribeSpotPriceHistoryInput.Set(in)
+func (e *EC2API) DescribeSpotPriceHistoryWithContext(_ aws.Context, input *ec2.DescribeSpotPriceHistoryInput, _ ...request.Option) (*ec2.DescribeSpotPriceHistoryOutput, error) {
+	e.DescribeSpotPriceHistoryInput.Set(input)
 	if !e.NextError.IsNil() {
 		defer e.NextError.Reset()
-		return e.NextError.Get()
+		return nil, e.NextError.Get()
 	}
 	if !e.DescribeSpotPriceHistoryOutput.IsNil() {
-		fn(e.DescribeSpotPriceHistoryOutput.Clone(), false)
-		return nil
+		return e.DescribeSpotPriceHistoryOutput.Clone(), nil
 	}
 	// fail if the test doesn't provide specific data which causes our pricing provider to use its static price list
-	return errors.New("no pricing data provided")
+	return nil, errors.New("no pricing data provided")
+}
+
+func (e *EC2API) DescribeSpotPriceHistoryPagesWithContext(ctx aws.Context, input *ec2.DescribeSpotPriceHistoryInput, fn func(*ec2.DescribeSpotPriceHistoryOutput, bool) bool, _ ...request.Option) error {
+	out, err := e.DescribeSpotPriceHistoryWithContext(ctx, input)
+	if err != nil {
+		return err
+	}
+	fn(out, false)
+	return nil
 }
diff --git a/pkg/operator/operator.go b/pkg/operator/operator.go
@@ -127,7 +127,7 @@ func NewOperator(ctx context.Context, operator *operator.Operator) (context.Cont
 		*sess.Config.Region,
 	)
 	amiProvider := amifamily.NewProvider(operator.GetClient(), operator.KubernetesInterface, ssm.New(sess), ec2api,
-		cache.New(awscache.DefaultTTL, awscache.DefaultCleanupInterval), cache.New(awscache.DefaultTTL, awscache.DefaultCleanupInterval), cache.New(awscache.DefaultTTL, awscache.DefaultCleanupInterval))
+		cache.New(awscache.DefaultTTL, awscache.DefaultCleanupInterval), cache.New(awscache.DefaultTTL, awscache.DefaultCleanupInterval))
 	amiResolver := amifamily.New(amiProvider)
 	launchTemplateProvider := launchtemplate.NewProvider(
 		ctx,