From d0a9ace27c354b7e936ad9c1da44a26b0610f07e Mon Sep 17 00:00:00 2001 From: Hao Zhou Date: Mon, 14 Dec 2020 14:11:17 -0800 Subject: [PATCH] Replace DescribeNetworkInterfaces with paginated version --- pkg/awsutils/awsutils.go | 30 +++++++++++++-- pkg/awsutils/awsutils_test.go | 48 +++++++++++++++--------- pkg/ec2wrapper/client.go | 1 + pkg/ec2wrapper/mocks/ec2wrapper_mocks.go | 19 ++++++++++ 4 files changed, 77 insertions(+), 21 deletions(-) diff --git a/pkg/awsutils/awsutils.go b/pkg/awsutils/awsutils.go index 65826ab9ea2..243541231e3 100644 --- a/pkg/awsutils/awsutils.go +++ b/pkg/awsutils/awsutils.go @@ -66,6 +66,9 @@ const ( // Http client timeout env for sessions httpTimeoutEnv = "HTTP_TIMEOUT" + + // the default page size when paginating the DescribeNetworkInterfaces call + describeENIPageSize = 1000 ) var ( @@ -355,7 +358,7 @@ func New(useCustomNetworking bool) (*EC2InstanceMetadataCache, error) { sess, err := session.NewSession( &aws.Config{ - Region: aws.String(cache.region), + Region: aws.String(cache.region), MaxRetries: aws.Int(15), HTTPClient: &http.Client{ Timeout: httpTimeoutValue, @@ -1425,14 +1428,16 @@ func (cache *EC2InstanceMetadataCache) getFilteredListOfNetworkInterfaces() ([]* input := &ec2.DescribeNetworkInterfacesInput{ Filters: []*ec2.Filter{tagFilter, statusFilter}, + MaxResults: aws.Int64(describeENIPageSize), } - result, err := cache.ec2SVC.DescribeNetworkInterfacesWithContext(context.Background(), input, userAgent) + + outputENIs, err := cache.getENIsFromPaginatedDescribeNetworkInterfaces(input) if err != nil { return nil, errors.Wrap(err, "awsutils: unable to obtain filtered list of network interfaces") } networkInterfaces := make([]*ec2.NetworkInterface, 0) - for _, networkInterface := range result.NetworkInterfaces { + for _, networkInterface := range outputENIs { // Verify the description starts with "aws-K8S-" if !strings.HasPrefix(aws.StringValue(networkInterface.Description), eniDescriptionPrefix) { continue @@ -1515,3 +1520,22 @@ func (cache *EC2InstanceMetadataCache) IsUnmanagedENI(eniID string) bool { } return false } + +func (cache *EC2InstanceMetadataCache) getENIsFromPaginatedDescribeNetworkInterfaces( + input *ec2.DescribeNetworkInterfacesInput) ([]*ec2.NetworkInterface, error) { + outputENIs := make([]*ec2.NetworkInterface, 0) + pageNum := 0 + log.Debugf("Paginating describe ENI has page size: %d", *input.MaxResults) + pageFn := func(output *ec2.DescribeNetworkInterfacesOutput, lastPage bool) (nextPage bool) { + pageNum++ + log.Debugf("EC2 DescribeNetworkInterfaces succeeded with %d results on page %d", + len(output.NetworkInterfaces), pageNum) + + outputENIs = append(outputENIs, output.NetworkInterfaces...) + // Loop is guided by nextToken, the func expect a false to exit. + return output.NextToken != nil + } + + err := cache.ec2SVC.DescribeNetworkInterfacesPagesWithContext(context.Background(), input, pageFn, userAgent) + return outputENIs, err +} diff --git a/pkg/awsutils/awsutils_test.go b/pkg/awsutils/awsutils_test.go index 5385ab82271..daa2725086f 100644 --- a/pkg/awsutils/awsutils_test.go +++ b/pkg/awsutils/awsutils_test.go @@ -24,6 +24,7 @@ import ( "testing" "time" + "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/service/ec2" "github.com/golang/mock/gomock" "github.com/stretchr/testify/assert" @@ -711,10 +712,8 @@ func TestEC2InstanceMetadataCache_getFilteredListOfNetworkInterfaces_OneResult(t attachment := &ec2.NetworkInterfaceAttachment{AttachmentId: &attachmentID} cureniID := eniID - result := &ec2.DescribeNetworkInterfacesOutput{ - NetworkInterfaces: []*ec2.NetworkInterface{{Attachment: attachment, Status: &status, TagSet: tag, Description: &description, NetworkInterfaceId: &cureniID}}} - mockEC2.EXPECT().DescribeNetworkInterfacesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil) - + interfaces := []*ec2.NetworkInterface{{Attachment: attachment, Status: &status, TagSet: tag, Description: &description, NetworkInterfaceId: &cureniID}} + helpCallMockedPagedDescribeNetworkInterface(t, mockEC2, interfaces, nil, 1) ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2} got, err := ins.getFilteredListOfNetworkInterfaces() assert.NotNil(t, got) @@ -725,10 +724,7 @@ func TestEC2InstanceMetadataCache_getFilteredListOfNetworkInterfaces_NoResult(t ctrl, mockEC2 := setup(t) defer ctrl.Finish() - result := &ec2.DescribeNetworkInterfacesOutput{ - NetworkInterfaces: []*ec2.NetworkInterface{}} - mockEC2.EXPECT().DescribeNetworkInterfacesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil) - + helpCallMockedPagedDescribeNetworkInterface(t, mockEC2, []*ec2.NetworkInterface{}, nil, 1) ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2} got, err := ins.getFilteredListOfNetworkInterfaces() assert.Nil(t, got) @@ -739,7 +735,12 @@ func TestEC2InstanceMetadataCache_getFilteredListOfNetworkInterfaces_Error(t *te ctrl, mockEC2 := setup(t) defer ctrl.Finish() - mockEC2.EXPECT().DescribeNetworkInterfacesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("dummy error")) + interfaces := []*ec2.NetworkInterface{{ + TagSet: []*ec2.Tag{ + {Key: aws.String("foo"), Value: aws.String("foo-value")}, + }, + }} + helpCallMockedPagedDescribeNetworkInterface(t, mockEC2, interfaces, errors.New("dummy error"), 1) ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2} got, err := ins.getFilteredListOfNetworkInterfaces() @@ -857,19 +858,30 @@ func TestEC2InstanceMetadataCache_cleanUpLeakedENIsInternal(t *testing.T) { defer ctrl.Finish() description := eniDescriptionPrefix + "test" - result := &ec2.DescribeNetworkInterfacesOutput{ - NetworkInterfaces: []*ec2.NetworkInterface{{ - Description: &description, - TagSet: []*ec2.Tag{ - {Key: aws.String(eniNodeTagKey), Value: aws.String("test-value")}, - }, - }}, - } + interfaces := []*ec2.NetworkInterface{{ + Description: &description, + TagSet: []*ec2.Tag{ + {Key: aws.String(eniNodeTagKey), Value: aws.String("test-value")}, + }, + }} - mockEC2.EXPECT().DescribeNetworkInterfacesWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil) + helpCallMockedPagedDescribeNetworkInterface(t, mockEC2, interfaces, nil, 1) mockEC2.EXPECT().CreateTagsWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil) ins := &EC2InstanceMetadataCache{ec2SVC: mockEC2} // Test checks that both mocks gets called. ins.cleanUpLeakedENIsInternal(time.Millisecond) } + +func helpCallMockedPagedDescribeNetworkInterface( + t *testing.T, mockEC2 *mock_ec2wrapper.MockEC2, interfaces []*ec2.NetworkInterface, err error, times int) { + mockEC2.EXPECT(). + DescribeNetworkInterfacesPagesWithContext(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(times). + DoAndReturn(func(_ context.Context, _ *ec2.DescribeNetworkInterfacesInput, + fn func(*ec2.DescribeNetworkInterfacesOutput, bool) bool, userAgent request.Option) error { + assert.Equal(t, false, fn(&ec2.DescribeNetworkInterfacesOutput{ + NetworkInterfaces: interfaces, + }, true)) + return err + }) +} diff --git a/pkg/ec2wrapper/client.go b/pkg/ec2wrapper/client.go index be1dad1e568..77dd743e741 100644 --- a/pkg/ec2wrapper/client.go +++ b/pkg/ec2wrapper/client.go @@ -33,6 +33,7 @@ type EC2 interface { DescribeNetworkInterfacesWithContext(ctx aws.Context, input *ec2svc.DescribeNetworkInterfacesInput, opts ...request.Option) (*ec2svc.DescribeNetworkInterfacesOutput, error) ModifyNetworkInterfaceAttributeWithContext(ctx aws.Context, input *ec2svc.ModifyNetworkInterfaceAttributeInput, opts ...request.Option) (*ec2svc.ModifyNetworkInterfaceAttributeOutput, error) CreateTagsWithContext(ctx aws.Context, input *ec2svc.CreateTagsInput, opts ...request.Option) (*ec2svc.CreateTagsOutput, error) + DescribeNetworkInterfacesPagesWithContext(ctx aws.Context, input *ec2svc.DescribeNetworkInterfacesInput, fn func(*ec2svc.DescribeNetworkInterfacesOutput, bool) bool, opts ...request.Option) error } // New creates a new EC2 wrapper diff --git a/pkg/ec2wrapper/mocks/ec2wrapper_mocks.go b/pkg/ec2wrapper/mocks/ec2wrapper_mocks.go index ecc2ebfdca7..f4791e0abd1 100644 --- a/pkg/ec2wrapper/mocks/ec2wrapper_mocks.go +++ b/pkg/ec2wrapper/mocks/ec2wrapper_mocks.go @@ -190,6 +190,25 @@ func (mr *MockEC2MockRecorder) DescribeInstancesWithContext(arg0, arg1 interface return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DescribeInstancesWithContext", reflect.TypeOf((*MockEC2)(nil).DescribeInstancesWithContext), varargs...) } +// DescribeNetworkInterfacesPagesWithContext mocks base method +func (m *MockEC2) DescribeNetworkInterfacesPagesWithContext(arg0 context.Context, arg1 *ec2.DescribeNetworkInterfacesInput, arg2 func(*ec2.DescribeNetworkInterfacesOutput, bool) bool, arg3 ...request.Option) error { + m.ctrl.T.Helper() + varargs := []interface{}{arg0, arg1, arg2} + for _, a := range arg3 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "DescribeNetworkInterfacesPagesWithContext", varargs...) + ret0, _ := ret[0].(error) + return ret0 +} + +// DescribeNetworkInterfacesPagesWithContext indicates an expected call of DescribeNetworkInterfacesPagesWithContext +func (mr *MockEC2MockRecorder) DescribeNetworkInterfacesPagesWithContext(arg0, arg1, arg2 interface{}, arg3 ...interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]interface{}{arg0, arg1, arg2}, arg3...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DescribeNetworkInterfacesPagesWithContext", reflect.TypeOf((*MockEC2)(nil).DescribeNetworkInterfacesPagesWithContext), varargs...) +} + // DescribeNetworkInterfacesWithContext mocks base method func (m *MockEC2) DescribeNetworkInterfacesWithContext(arg0 context.Context, arg1 *ec2.DescribeNetworkInterfacesInput, arg2 ...request.Option) (*ec2.DescribeNetworkInterfacesOutput, error) { m.ctrl.T.Helper()