Skip to content

Commit b84aaa2

Browse files
gyuhogyuho
authored
fix(login): prefer provider region for login requests (#1246)
1 parent 51edfea commit b84aaa2

12 files changed

Lines changed: 569 additions & 16 deletions

File tree

‎pkg/machine-info/login_request.go‎

Lines changed: 35 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package machineinfo
22

33
import (
44
"fmt"
5+
"strings"
56

67
corev1 "k8s.io/api/core/v1"
78
"k8s.io/apimachinery/pkg/api/resource"
@@ -45,17 +46,20 @@ func createLoginRequest(
4546
donec := make(chan struct{})
4647
defer close(donec)
4748

48-
// deciding machine location can take awhile
49-
// depending on the network latency
50-
// run async
51-
machineLocationCh := make(chan *apiv1.MachineLocation, 1)
52-
go func() {
53-
select {
54-
case <-donec:
55-
return
56-
case machineLocationCh <- getMachineLocationFunc():
57-
}
58-
}()
49+
var machineLocationCh <-chan *apiv1.MachineLocation
50+
startMachineLocationLookup := func() {
51+
// Deciding machine location can take awhile depending on network latency,
52+
// so keep the DERP fallback asynchronous while the rest of login is built.
53+
ch := make(chan *apiv1.MachineLocation, 1)
54+
machineLocationCh = ch
55+
go func() {
56+
select {
57+
case <-donec:
58+
return
59+
case ch <- getMachineLocationFunc():
60+
}
61+
}()
62+
}
5963

6064
req := &apiv1.LoginRequest{
6165
Token: token,
@@ -80,6 +84,10 @@ func createLoginRequest(
8084

8185
req.Provider = detectedProvider.Provider
8286
req.ProviderInstanceID = detectedProvider.InstanceID
87+
providerLocation := getProviderLocation(detectedProvider)
88+
if providerLocation == nil {
89+
startMachineLocationLookup()
90+
}
8391

8492
// we always prioritize the provider's public IP and private IP
8593
// even if the local network interface has a private IP
@@ -138,7 +146,22 @@ func createLoginRequest(
138146
req.Resources["nvidia.com/gpu"] = gpuCnt
139147
}
140148

141-
req.Location = <-machineLocationCh
149+
if providerLocation != nil {
150+
req.Location = providerLocation
151+
} else if machineLocationCh != nil {
152+
req.Location = <-machineLocationCh
153+
}
142154

143155
return req, nil
144156
}
157+
158+
func getProviderLocation(providerInfo *providers.Info) *apiv1.MachineLocation {
159+
if providerInfo == nil {
160+
return nil
161+
}
162+
region := strings.TrimSpace(providerInfo.Region)
163+
if region == "" {
164+
return nil
165+
}
166+
return &apiv1.MachineLocation{Region: region}
167+
}

‎pkg/machine-info/login_request_test.go‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1141,3 +1141,78 @@ func TestCreateLoginRequest_NscaleUsesLatencyLocationWithoutMetadataOverride(t *
11411141
assert.Equal(t, "eu-west-1", req.Location.Region)
11421142
assert.Equal(t, "eu-west-1a", req.Location.Zone)
11431143
}
1144+
1145+
func TestCreateLoginRequest_ProviderRegionOverridesLatencyLocation(t *testing.T) {
1146+
getMachineInfoFunc := func(nvidianvml.Instance) (*apiv1.MachineInfo, error) {
1147+
return &apiv1.MachineInfo{
1148+
CPUInfo: &apiv1.MachineCPUInfo{LogicalCores: 4},
1149+
MemoryInfo: &apiv1.MachineMemoryInfo{TotalBytes: 16 * 1024 * 1024 * 1024},
1150+
NICInfo: &apiv1.MachineNICInfo{},
1151+
}, nil
1152+
}
1153+
getProviderFunc := func(ip string) *providers.Info {
1154+
return &providers.Info{
1155+
Provider: "aws",
1156+
PublicIP: ip,
1157+
PrivateIP: "172.31.0.10",
1158+
Region: "us-east-1",
1159+
InstanceID: "i-00001923",
1160+
}
1161+
}
1162+
1163+
machineLocationCalled := false
1164+
req, err := createLoginRequest(
1165+
"token",
1166+
"machine-id",
1167+
"",
1168+
"1",
1169+
&mockNvmlInstance{},
1170+
func() (string, error) { return "54.123.45.67", nil },
1171+
func() *apiv1.MachineLocation {
1172+
machineLocationCalled = true
1173+
return &apiv1.MachineLocation{Region: "us-west-2", Zone: "us-west-2a"}
1174+
},
1175+
getMachineInfoFunc,
1176+
getProviderFunc,
1177+
func() (string, error) { return "100Gi", nil },
1178+
func(nvidianvml.Instance) (string, error) { return "1", nil },
1179+
)
1180+
1181+
assert.NoError(t, err)
1182+
assert.NotNil(t, req)
1183+
assert.NotNil(t, req.Location)
1184+
assert.Equal(t, "us-east-1", req.Location.Region)
1185+
assert.Empty(t, req.Location.Zone)
1186+
assert.False(t, machineLocationCalled)
1187+
}
1188+
1189+
func TestGetProviderLocation(t *testing.T) {
1190+
tests := []struct {
1191+
name string
1192+
provider *providers.Info
1193+
want *apiv1.MachineLocation
1194+
}{
1195+
{
1196+
name: "nil provider",
1197+
},
1198+
{
1199+
name: "empty region",
1200+
provider: &providers.Info{Provider: "aws"},
1201+
},
1202+
{
1203+
name: "whitespace region",
1204+
provider: &providers.Info{Provider: "aws", Region: " "},
1205+
},
1206+
{
1207+
name: "trims region",
1208+
provider: &providers.Info{Provider: "aws", Region: " us-east-1 "},
1209+
want: &apiv1.MachineLocation{Region: "us-east-1"},
1210+
},
1211+
}
1212+
1213+
for _, tc := range tests {
1214+
t.Run(tc.name, func(t *testing.T) {
1215+
assert.Equal(t, tc.want, getProviderLocation(tc.provider))
1216+
})
1217+
}
1218+
}

‎pkg/providers/all/all.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,15 @@ func Detect(ctx context.Context) (*pkgproviders.Info, error) {
7070
log.Logger.Infow("successfully detected private IP", "provider", detector.Name(), "privateIP", privateIP)
7171
}
7272

73+
if regionDetector, ok := detector.(pkgproviders.RegionDetector); ok {
74+
region, err := regionDetector.Region(ctx)
75+
if err != nil {
76+
log.Logger.Warnw("failed to get region", "provider", detector.Name(), "error", err)
77+
} else {
78+
info.Region = region
79+
}
80+
}
81+
7382
vmEnvironment, err := detector.VMEnvironment(ctx)
7483
if err != nil {
7584
log.Logger.Warnw("failed to get VM environment", "provider", detector.Name(), "error", err)

‎pkg/providers/all/all_test.go‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,20 +17,35 @@ type mockDetector struct {
1717
provider string
1818
publicIP string
1919
privateIP string
20+
region string
2021
provErr error
2122
publicErr error
2223
privateErr error
24+
regionErr error
2325
vmEnv string
2426
vmEnvErr error
2527
instanceID string
2628
instanceIDErr error
2729
delay time.Duration
2830
}
2931

32+
type legacyDetector struct {
33+
name string
34+
provider string
35+
publicIP string
36+
privateIP string
37+
vmEnv string
38+
instanceID string
39+
}
40+
3041
func (m *mockDetector) Name() string {
3142
return m.name
3243
}
3344

45+
func (m *legacyDetector) Name() string {
46+
return m.name
47+
}
48+
3449
func (m *mockDetector) Provider(ctx context.Context) (string, error) {
3550
if m.delay > 0 {
3651
select {
@@ -42,22 +57,46 @@ func (m *mockDetector) Provider(ctx context.Context) (string, error) {
4257
return m.provider, m.provErr
4358
}
4459

60+
func (m *legacyDetector) Provider(ctx context.Context) (string, error) {
61+
return m.provider, nil
62+
}
63+
4564
func (m *mockDetector) PublicIPv4(ctx context.Context) (string, error) {
4665
return m.publicIP, m.publicErr
4766
}
4867

68+
func (m *legacyDetector) PublicIPv4(ctx context.Context) (string, error) {
69+
return m.publicIP, nil
70+
}
71+
4972
func (m *mockDetector) PrivateIPv4(ctx context.Context) (string, error) {
5073
return m.privateIP, m.privateErr
5174
}
5275

76+
func (m *legacyDetector) PrivateIPv4(ctx context.Context) (string, error) {
77+
return m.privateIP, nil
78+
}
79+
80+
func (m *mockDetector) Region(ctx context.Context) (string, error) {
81+
return m.region, m.regionErr
82+
}
83+
5384
func (m *mockDetector) VMEnvironment(ctx context.Context) (string, error) {
5485
return m.vmEnv, m.vmEnvErr
5586
}
5687

88+
func (m *legacyDetector) VMEnvironment(ctx context.Context) (string, error) {
89+
return m.vmEnv, nil
90+
}
91+
5792
func (m *mockDetector) InstanceID(ctx context.Context) (string, error) {
5893
return m.instanceID, m.instanceIDErr
5994
}
6095

96+
func (m *legacyDetector) InstanceID(ctx context.Context) (string, error) {
97+
return m.instanceID, nil
98+
}
99+
61100
// withTemporaryDetectors runs the provided function with a temporary replacement for All
62101
// and restores the original value when done
63102
func withTemporaryDetectors(tempDetectors []providers.Detector, fn func()) {
@@ -76,6 +115,56 @@ func TestDetect_Success(t *testing.T) {
76115
provider: "aws",
77116
publicIP: "1.2.3.4",
78117
privateIP: "10.0.1.100",
118+
region: "us-east-1",
119+
vmEnv: "AWS",
120+
instanceID: "i-abc",
121+
},
122+
}
123+
124+
withTemporaryDetectors(testDetectors, func() {
125+
info, err := Detect(context.Background())
126+
assert.NoError(t, err)
127+
assert.Equal(t, "aws", info.Provider)
128+
assert.Equal(t, "1.2.3.4", info.PublicIP)
129+
assert.Equal(t, "10.0.1.100", info.PrivateIP)
130+
assert.Equal(t, "us-east-1", info.Region)
131+
assert.Equal(t, "AWS", info.VMEnvironment)
132+
assert.Equal(t, "i-abc", info.InstanceID)
133+
})
134+
}
135+
136+
func TestDetect_SuccessWithoutRegionDetector(t *testing.T) {
137+
testDetectors := []providers.Detector{
138+
&legacyDetector{
139+
name: "legacy",
140+
provider: "legacy",
141+
publicIP: "1.2.3.4",
142+
privateIP: "10.0.1.100",
143+
vmEnv: "LEGACY",
144+
instanceID: "i-legacy",
145+
},
146+
}
147+
148+
withTemporaryDetectors(testDetectors, func() {
149+
info, err := Detect(context.Background())
150+
assert.NoError(t, err)
151+
assert.Equal(t, "legacy", info.Provider)
152+
assert.Equal(t, "1.2.3.4", info.PublicIP)
153+
assert.Equal(t, "10.0.1.100", info.PrivateIP)
154+
assert.Empty(t, info.Region)
155+
assert.Equal(t, "LEGACY", info.VMEnvironment)
156+
assert.Equal(t, "i-legacy", info.InstanceID)
157+
})
158+
}
159+
160+
func TestDetect_RegionError(t *testing.T) {
161+
testDetectors := []providers.Detector{
162+
&mockDetector{
163+
name: "aws",
164+
provider: "aws",
165+
publicIP: "1.2.3.4",
166+
privateIP: "10.0.1.100",
167+
regionErr: errors.New("region error"),
79168
vmEnv: "AWS",
80169
instanceID: "i-abc",
81170
},
@@ -87,6 +176,7 @@ func TestDetect_Success(t *testing.T) {
87176
assert.Equal(t, "aws", info.Provider)
88177
assert.Equal(t, "1.2.3.4", info.PublicIP)
89178
assert.Equal(t, "10.0.1.100", info.PrivateIP)
179+
assert.Empty(t, info.Region)
90180
assert.Equal(t, "AWS", info.VMEnvironment)
91181
assert.Equal(t, "i-abc", info.InstanceID)
92182
})

‎pkg/providers/aws/aws.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import (
1010
const Name = "aws"
1111

1212
func New() providers.Detector {
13-
return providers.New(Name, detectProvider, imds.FetchPublicIPv4, imds.FetchLocalIPv4, nil, imds.FetchInstanceID)
13+
return providers.NewWithRegion(Name, detectProvider, imds.FetchPublicIPv4, imds.FetchLocalIPv4, imds.FetchRegion, nil, imds.FetchInstanceID)
1414
}
1515

1616
func detectProvider(ctx context.Context) (string, error) {

0 commit comments

Comments
 (0)