Skip to content

Commit b5097b5

Browse files
shraddhabangshraddhabang
committed
[feat aga] Implement endpoint loader with DNS resolution
1 parent 9f9852e commit b5097b5

File tree

9 files changed

+1961
-0
lines changed

9 files changed

+1961
-0
lines changed

‎go.mod‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ require (
104104
github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79 // indirect
105105
github.com/hashicorp/errwrap v1.1.0 // indirect
106106
github.com/hashicorp/go-multierror v1.1.1 // indirect
107+
github.com/hashicorp/golang-lru v1.0.2 // indirect
107108
github.com/huandu/xstrings v1.5.0 // indirect
108109
github.com/imkira/go-interpol v1.1.0 // indirect
109110
github.com/inconshreveable/mousetrap v1.1.0 // indirect

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,8 @@ github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY
228228
github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
229229
github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo=
230230
github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM=
231+
github.com/hashicorp/golang-lru v1.0.2 h1:dV3g9Z/unq5DpblPpw+Oqcv4dU/1omnb4Ok8iPY6p1c=
232+
github.com/hashicorp/golang-lru v1.0.2/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
231233
github.com/hashicorp/golang-lru/arc/v2 v2.0.5 h1:l2zaLDubNhW4XO3LnliVj0GXO3+/CGNJAg1dcN2Fpfw=
232234
github.com/hashicorp/golang-lru/arc/v2 v2.0.5/go.mod h1:ny6zBSQZi2JxIeYcv7kt2sH2PXJtirBN7RDhRpxPkxU=
233235
github.com/hashicorp/golang-lru/v2 v2.0.5 h1:wW7h1TG88eUIJ2i69gaE3uNVtEPIagzhGvHgwfx2Vm4=

‎pkg/aga/dns_resolver.go‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package aga
2+
3+
import (
4+
"context"
5+
"fmt"
6+
awssdk "github.com/aws/aws-sdk-go-v2/aws"
7+
"sync"
8+
"time"
9+
10+
elbv2sdk "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
11+
"github.com/hashicorp/golang-lru"
12+
"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/services"
13+
)
14+
15+
// DNSResolver resolves load balancer DNS names to ARNs
16+
type DNSResolver struct {
17+
elbv2Client services.ELBV2
18+
cache *lru.Cache
19+
cacheMutex sync.RWMutex
20+
ttl time.Duration
21+
}
22+
23+
type cacheEntry struct {
24+
arn string
25+
expireAt time.Time
26+
}
27+
28+
// NewDNSResolver creates a new DNSResolver
29+
func NewDNSResolver(elbv2Client services.ELBV2) (*DNSResolver, error) {
30+
// AWS Global Accelerator has a quota of 420 endpoints per AWS account (can be increased)
31+
// Using 420 provides headroom while efficiently caching DNS-to-ARN resolutions
32+
cache, err := lru.New(420)
33+
if err != nil {
34+
return nil, err
35+
}
36+
37+
return &DNSResolver{
38+
elbv2Client: elbv2Client,
39+
cache: cache,
40+
ttl: 5 * time.Minute, // Default TTL of 5 minutes
41+
}, nil
42+
}
43+
44+
// ResolveDNSToARN resolves a load balancer DNS name to an ARN
45+
func (r *DNSResolver) ResolveDNSToARN(ctx context.Context, dnsName string) (string, error) {
46+
if dnsName == "" {
47+
return "", fmt.Errorf("empty DNS name")
48+
}
49+
50+
// Check cache first
51+
r.cacheMutex.RLock()
52+
if value, found := r.cache.Get(dnsName); found {
53+
entry := value.(cacheEntry)
54+
// Check if the cache entry is still valid
55+
if time.Now().Before(entry.expireAt) {
56+
r.cacheMutex.RUnlock()
57+
return entry.arn, nil
58+
}
59+
// Entry has expired, remove from cache
60+
r.cache.Remove(dnsName)
61+
}
62+
r.cacheMutex.RUnlock()
63+
64+
req := &elbv2sdk.DescribeLoadBalancersInput{}
65+
lbs, err := r.elbv2Client.DescribeLoadBalancersAsList(ctx, req)
66+
if err != nil {
67+
return "", fmt.Errorf("failed to describe load balancers: %w", err)
68+
}
69+
if len(lbs) == 0 {
70+
return "", fmt.Errorf("no load balancers found")
71+
}
72+
arn := ""
73+
for _, lb := range lbs {
74+
if awssdk.ToString(lb.DNSName) == dnsName {
75+
arn = awssdk.ToString(lb.LoadBalancerArn)
76+
break
77+
}
78+
}
79+
if arn == "" {
80+
return "", fmt.Errorf("no load balancer found for dns %s", dnsName)
81+
}
82+
83+
// Cache the result
84+
r.cacheMutex.Lock()
85+
r.cache.Add(dnsName, cacheEntry{
86+
arn: arn,
87+
expireAt: time.Now().Add(r.ttl),
88+
})
89+
r.cacheMutex.Unlock()
90+
91+
return arn, nil
92+
}
93+
94+
// Ensure DNSResolver implements DNSResolverInterface
95+
var _ DNSResolverInterface = (*DNSResolver)(nil)

‎pkg/aga/dns_resolver_test.go‎

Lines changed: 273 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,273 @@
1+
package aga
2+
3+
import (
4+
"context"
5+
awssdk "github.com/aws/aws-sdk-go-v2/aws"
6+
elbv2sdk "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
7+
"github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
8+
"github.com/golang/mock/gomock"
9+
"github.com/pkg/errors"
10+
"github.com/stretchr/testify/assert"
11+
"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/services"
12+
"testing"
13+
"time"
14+
)
15+
16+
func TestDNSResolver_ResolveDNSToARN(t *testing.T) {
17+
type describeLoadBalancersAsListCall struct {
18+
req *elbv2sdk.DescribeLoadBalancersInput
19+
resp []types.LoadBalancer
20+
err error
21+
}
22+
23+
type fields struct {
24+
elbv2Client *services.MockELBV2
25+
describeLoadBalancersCalls []describeLoadBalancersAsListCall
26+
}
27+
28+
tests := []struct {
29+
name string
30+
fields fields
31+
dnsName string
32+
wantARN string
33+
wantErr bool
34+
setupFields func(fields fields)
35+
}{
36+
{
37+
name: "successfully resolves DNS to ARN",
38+
fields: fields{
39+
elbv2Client: services.NewMockELBV2(gomock.NewController(t)),
40+
describeLoadBalancersCalls: []describeLoadBalancersAsListCall{
41+
{
42+
req: &elbv2sdk.DescribeLoadBalancersInput{},
43+
resp: []types.LoadBalancer{
44+
{
45+
DNSName: awssdk.String("test-lb.us-west-2.elb.amazonaws.com"),
46+
LoadBalancerArn: awssdk.String("arn:aws:elasticloadbalancing:us-west-2:123456789012:loadbalancer/app/test-lb/1234567890abcdef"),
47+
},
48+
{
49+
DNSName: awssdk.String("another-lb.us-west-2.elb.amazonaws.com"),
50+
LoadBalancerArn: awssdk.String("arn:aws:elasticloadbalancing:us-west-2:123456789012:loadbalancer/app/another-lb/0987654321fedcba"),
51+
},
52+
},
53+
err: nil,
54+
},
55+
},
56+
},
57+
dnsName: "test-lb.us-west-2.elb.amazonaws.com",
58+
wantARN: "arn:aws:elasticloadbalancing:us-west-2:123456789012:loadbalancer/app/test-lb/1234567890abcdef",
59+
wantErr: false,
60+
setupFields: func(fields fields) {
61+
gomock.InOrder(
62+
fields.elbv2Client.EXPECT().
63+
DescribeLoadBalancersAsList(gomock.Any(), fields.describeLoadBalancersCalls[0].req).
64+
Return(fields.describeLoadBalancersCalls[0].resp, fields.describeLoadBalancersCalls[0].err),
65+
)
66+
},
67+
},
68+
{
69+
name: "uses cached ARN on second call",
70+
fields: fields{
71+
elbv2Client: services.NewMockELBV2(gomock.NewController(t)),
72+
describeLoadBalancersCalls: []describeLoadBalancersAsListCall{
73+
{
74+
req: &elbv2sdk.DescribeLoadBalancersInput{},
75+
resp: []types.LoadBalancer{
76+
{
77+
DNSName: awssdk.String("test-lb.us-west-2.elb.amazonaws.com"),
78+
LoadBalancerArn: awssdk.String("arn:aws:elasticloadbalancing:us-west-2:123456789012:loadbalancer/app/test-lb/1234567890abcdef"),
79+
},
80+
},
81+
err: nil,
82+
},
83+
},
84+
},
85+
dnsName: "test-lb.us-west-2.elb.amazonaws.com",
86+
wantARN: "arn:aws:elasticloadbalancing:us-west-2:123456789012:loadbalancer/app/test-lb/1234567890abcdef",
87+
wantErr: false,
88+
setupFields: func(fields fields) {
89+
gomock.InOrder(
90+
fields.elbv2Client.EXPECT().
91+
DescribeLoadBalancersAsList(gomock.Any(), fields.describeLoadBalancersCalls[0].req).
92+
Return(fields.describeLoadBalancersCalls[0].resp, fields.describeLoadBalancersCalls[0].err).
93+
Times(1),
94+
)
95+
},
96+
},
97+
{
98+
name: "returns error for empty DNS name",
99+
fields: fields{
100+
elbv2Client: services.NewMockELBV2(gomock.NewController(t)),
101+
describeLoadBalancersCalls: []describeLoadBalancersAsListCall{},
102+
},
103+
dnsName: "",
104+
wantARN: "",
105+
wantErr: true,
106+
setupFields: func(fields fields) {
107+
// No calls expected for empty DNS name
108+
},
109+
},
110+
{
111+
name: "returns error when no load balancers found",
112+
fields: fields{
113+
elbv2Client: services.NewMockELBV2(gomock.NewController(t)),
114+
describeLoadBalancersCalls: []describeLoadBalancersAsListCall{
115+
{
116+
req: &elbv2sdk.DescribeLoadBalancersInput{},
117+
resp: []types.LoadBalancer{},
118+
err: nil,
119+
},
120+
},
121+
},
122+
dnsName: "test-lb.us-west-2.elb.amazonaws.com",
123+
wantARN: "",
124+
wantErr: true,
125+
setupFields: func(fields fields) {
126+
gomock.InOrder(
127+
fields.elbv2Client.EXPECT().
128+
DescribeLoadBalancersAsList(gomock.Any(), fields.describeLoadBalancersCalls[0].req).
129+
Return(fields.describeLoadBalancersCalls[0].resp, fields.describeLoadBalancersCalls[0].err),
130+
)
131+
},
132+
},
133+
{
134+
name: "returns error when no matching load balancer found",
135+
fields: fields{
136+
elbv2Client: services.NewMockELBV2(gomock.NewController(t)),
137+
describeLoadBalancersCalls: []describeLoadBalancersAsListCall{
138+
{
139+
req: &elbv2sdk.DescribeLoadBalancersInput{},
140+
resp: []types.LoadBalancer{
141+
{
142+
DNSName: awssdk.String("another-lb.us-west-2.elb.amazonaws.com"),
143+
LoadBalancerArn: awssdk.String("arn:aws:elasticloadbalancing:us-west-2:123456789012:loadbalancer/app/another-lb/0987654321fedcba"),
144+
},
145+
},
146+
err: nil,
147+
},
148+
},
149+
},
150+
dnsName: "test-lb.us-west-2.elb.amazonaws.com",
151+
wantARN: "",
152+
wantErr: true,
153+
setupFields: func(fields fields) {
154+
gomock.InOrder(
155+
fields.elbv2Client.EXPECT().
156+
DescribeLoadBalancersAsList(gomock.Any(), fields.describeLoadBalancersCalls[0].req).
157+
Return(fields.describeLoadBalancersCalls[0].resp, fields.describeLoadBalancersCalls[0].err),
158+
)
159+
},
160+
},
161+
{
162+
name: "returns error when API call fails",
163+
fields: fields{
164+
elbv2Client: services.NewMockELBV2(gomock.NewController(t)),
165+
describeLoadBalancersCalls: []describeLoadBalancersAsListCall{
166+
{
167+
req: &elbv2sdk.DescribeLoadBalancersInput{},
168+
resp: nil,
169+
err: errors.New("API error"),
170+
},
171+
},
172+
},
173+
dnsName: "test-lb.us-west-2.elb.amazonaws.com",
174+
wantARN: "",
175+
wantErr: true,
176+
setupFields: func(fields fields) {
177+
gomock.InOrder(
178+
fields.elbv2Client.EXPECT().
179+
DescribeLoadBalancersAsList(gomock.Any(), fields.describeLoadBalancersCalls[0].req).
180+
Return(fields.describeLoadBalancersCalls[0].resp, fields.describeLoadBalancersCalls[0].err),
181+
)
182+
},
183+
},
184+
}
185+
186+
// Add a test case for cache expiration
187+
t.Run("cache expiration", func(t *testing.T) {
188+
ctrl := gomock.NewController(t)
189+
elbv2Client := services.NewMockELBV2(ctrl)
190+
dnsName := "expired-lb.us-west-2.elb.amazonaws.com"
191+
originalARN := "arn:aws:elasticloadbalancing:us-west-2:123456789012:loadbalancer/app/expired-lb/original"
192+
updatedARN := "arn:aws:elasticloadbalancing:us-west-2:123456789012:loadbalancer/app/expired-lb/updated"
193+
194+
// Create resolver with a small TTL for testing
195+
resolver, err := NewDNSResolver(elbv2Client)
196+
assert.NoError(t, err)
197+
198+
// Override the TTL for testing
199+
resolver.ttl = 10 * time.Millisecond
200+
201+
// First call, should resolve through API
202+
elbv2Client.EXPECT().
203+
DescribeLoadBalancersAsList(gomock.Any(), &elbv2sdk.DescribeLoadBalancersInput{}).
204+
Return([]types.LoadBalancer{
205+
{
206+
DNSName: awssdk.String(dnsName),
207+
LoadBalancerArn: awssdk.String(originalARN),
208+
},
209+
}, nil).
210+
Times(1)
211+
212+
gotARN1, err := resolver.ResolveDNSToARN(context.Background(), dnsName)
213+
assert.NoError(t, err)
214+
assert.Equal(t, originalARN, gotARN1)
215+
216+
// Wait for cache to expire
217+
time.Sleep(15 * time.Millisecond)
218+
219+
// Second call after cache expiry, should resolve through API again
220+
elbv2Client.EXPECT().
221+
DescribeLoadBalancersAsList(gomock.Any(), &elbv2sdk.DescribeLoadBalancersInput{}).
222+
Return([]types.LoadBalancer{
223+
{
224+
DNSName: awssdk.String(dnsName),
225+
LoadBalancerArn: awssdk.String(updatedARN), // Different ARN to verify re-resolution
226+
},
227+
}, nil).
228+
Times(1)
229+
230+
gotARN2, err := resolver.ResolveDNSToARN(context.Background(), dnsName)
231+
assert.NoError(t, err)
232+
assert.Equal(t, updatedARN, gotARN2, "ARN should be updated after cache expiry")
233+
})
234+
235+
for _, tt := range tests {
236+
t.Run(tt.name, func(t *testing.T) {
237+
tt.setupFields(tt.fields)
238+
239+
resolver, err := NewDNSResolver(tt.fields.elbv2Client)
240+
assert.NoError(t, err)
241+
242+
// For cache test, we need to call it twice
243+
if tt.name == "uses cached ARN on second call" {
244+
// First call
245+
gotARN, err := resolver.ResolveDNSToARN(context.Background(), tt.dnsName)
246+
if tt.wantErr {
247+
assert.Error(t, err)
248+
} else {
249+
assert.NoError(t, err)
250+
assert.Equal(t, tt.wantARN, gotARN)
251+
}
252+
253+
// Second call - should use cache
254+
gotARN, err = resolver.ResolveDNSToARN(context.Background(), tt.dnsName)
255+
if tt.wantErr {
256+
assert.Error(t, err)
257+
} else {
258+
assert.NoError(t, err)
259+
assert.Equal(t, tt.wantARN, gotARN)
260+
}
261+
} else {
262+
// Regular test
263+
gotARN, err := resolver.ResolveDNSToARN(context.Background(), tt.dnsName)
264+
if tt.wantErr {
265+
assert.Error(t, err)
266+
} else {
267+
assert.NoError(t, err)
268+
assert.Equal(t, tt.wantARN, gotARN)
269+
}
270+
}
271+
})
272+
}
273+
}

0 commit comments

Comments
 (0)