Merge pull request #34 from zalando-incubator/pagination

implement pagination for aws rrs and lb

pkg/aws/client.go

@@ -48,25 +48,28 @@ func New(o Options) *Client {
 
 //ListRecordSets retrieve all records existing in the specified hosted zone
 func (c *Client) ListRecordSets(zoneID string) ([]*route53.ResourceRecordSet, error) {
+	records := make([]*route53.ResourceRecordSet, 0)
+
 	client, err := c.initRoute53Client()
 	if err != nil {
 		return nil, err
 	}
 
-	// TODO: implement paging
 	params := &route53.ListResourceRecordSetsInput{
 		HostedZoneId: aws.String(zoneID),
 	}
-	rsp, err := client.ListResourceRecordSets(params)
+
+	err = client.ListResourceRecordSetsPages(params, func(resp *route53.ListResourceRecordSetsOutput, lastPage bool) bool {
+		log.Debugf("Getting a list of AWS RRS of length: %d", len(resp.ResourceRecordSets))
+		records = append(records, resp.ResourceRecordSets...)
+		return !lastPage
+	})
+
 	if err != nil {
 		return nil, err
 	}
 
-	if rsp == nil {
-		return nil, ErrInvalidAWSResponse
-	}
-
-	return rsp.ResourceRecordSets, nil
+	return records, nil
 }
 
 //ChangeRecordSets creates and submits the record set change against the AWS API

pkg/aws/elb.go

@@ -1,6 +1,9 @@
 package aws
 
 import (
+	"sync"
+
+	log "github.com/Sirupsen/logrus"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/elb"
@@ -14,6 +17,10 @@ type LoadBalancer struct {
 	CanonicalZoneID string
 }
 
+//GetLoadBalancerFunc is a func type to represent the interface of functions that retrieve the list
+//of load balancers from AWS
+type GetLoadBalancerFunc func(*session.Session) ([]*LoadBalancer, error)
+
 //getCanonicalZoneIDs returns the map of LB (ALB + ELB classic) mapped to its CanonicalHostedZoneId
 func (c *Client) getCanonicalZoneIDs(endpoints []*pkg.Endpoint) (map[string]string, error) {
 	session, err := session.NewSessionWithOptions(session.Options{
@@ -27,17 +34,28 @@ func (c *Client) getCanonicalZoneIDs(endpoints []*pkg.Endpoint) (map[string]stri
 		return nil, err
 	}
 
-	albs, err := c.getALBs(session)
-	if err != nil {
-		return nil, err
-	}
+	loadBalancers := make([]*LoadBalancer, 0)
 
-	elbs, err := c.getELBs(session)
-	if err != nil {
-		return nil, err
+	var addLBMutex sync.Mutex
+	var wg sync.WaitGroup
+
+	for _, getLBs := range []GetLoadBalancerFunc{c.getALBs, c.getELBs} {
+		wg.Add(1)
+		go func(getLBs GetLoadBalancerFunc) {
+			defer wg.Done()
+			lbs, err := getLBs(session)
+			if err != nil {
+				log.Errorf("Error getting LBs: %v. Skipping...", err)
+				return
+			}
+			addLBMutex.Lock()
+			loadBalancers = append(loadBalancers, lbs...)
+			addLBMutex.Unlock()
+		}(getLBs)
 	}
 
-	loadBalancers := append(albs, elbs...)
+	wg.Wait()
+
 	loadBalancersMap := map[string]string{} //map LB Dns to its canonical hosted zone id
 
 	for _, endpoint := range endpoints {
@@ -52,39 +70,53 @@ func (c *Client) getCanonicalZoneIDs(endpoints []*pkg.Endpoint) (map[string]stri
 }
 
 func (c *Client) getELBs(session *session.Session) ([]*LoadBalancer, error) {
+	result := make([]*LoadBalancer, 0)
+
 	client := elb.New(session)
 
-	resp, err := client.DescribeLoadBalancers(nil)
+	params := &elb.DescribeLoadBalancersInput{}
+
+	err := client.DescribeLoadBalancersPages(params, func(resp *elb.DescribeLoadBalancersOutput, lastPage bool) bool {
+		loadBalancers := resp.LoadBalancerDescriptions
+		log.Debugf("Getting a page of ELBs of length: %d", len(resp.LoadBalancerDescriptions))
+		for _, loadbalancer := range loadBalancers {
+			result = append(result, &LoadBalancer{
+				DNSName:         aws.StringValue(loadbalancer.DNSName),
+				CanonicalZoneID: aws.StringValue(loadbalancer.CanonicalHostedZoneNameID),
+			})
+		}
+		return !lastPage
+	})
+
 	if err != nil {
 		return nil, err
 	}
 
-	loadBalancers := resp.LoadBalancerDescriptions
-	result := make([]*LoadBalancer, len(loadBalancers))
-	for i, loadbalancer := range loadBalancers {
-		result[i] = &LoadBalancer{
-			DNSName:         aws.StringValue(loadbalancer.DNSName),
-			CanonicalZoneID: aws.StringValue(loadbalancer.CanonicalHostedZoneNameID),
-		}
-	}
 	return result, nil
 }
 
 func (c *Client) getALBs(session *session.Session) ([]*LoadBalancer, error) {
+	result := make([]*LoadBalancer, 0)
+
 	client := elbv2.New(session)
 
-	resp, err := client.DescribeLoadBalancers(nil)
+	params := &elbv2.DescribeLoadBalancersInput{}
+
+	err := client.DescribeLoadBalancersPages(params, func(resp *elbv2.DescribeLoadBalancersOutput, lastPage bool) bool {
+		loadBalancers := resp.LoadBalancers
+		log.Debugf("Getting a page of ALBs of length: %d", len(resp.LoadBalancers))
+		for _, loadbalancer := range loadBalancers {
+			result = append(result, &LoadBalancer{
+				DNSName:         aws.StringValue(loadbalancer.DNSName),
+				CanonicalZoneID: aws.StringValue(loadbalancer.CanonicalHostedZoneId),
+			})
+		}
+		return !lastPage
+	})
+
 	if err != nil {
 		return nil, err
 	}
 
-	loadBalancers := resp.LoadBalancers
-	result := make([]*LoadBalancer, len(loadBalancers))
-	for i, loadbalancer := range loadBalancers {
-		result[i] = &LoadBalancer{
-			DNSName:         aws.StringValue(loadbalancer.DNSName),
-			CanonicalZoneID: aws.StringValue(loadbalancer.CanonicalHostedZoneId),
-		}
-	}
 	return result, nil
 }