Skip to content

Commit

Permalink
Add in-country biasing for Locate V2 (#94)
Browse files Browse the repository at this point in the history 
* Add in-country biasing

* Make test parameter explicit

* Move check

* Remove unnecessary line
  • Loading branch information
@cristinaleonr
cristinaleonr authored Nov 23, 2022
1 parent 8ae6c47 commit 0687d2f
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 12 deletions.
5 changes: 3 additions & 2 deletions handler/handler.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ type Locator interface {
// LocatorV2 defines how the Nearest handler requests machines nearest to the
// client.
type LocatorV2 interface {
Nearest(service, typ string, lat, lon float64) ([]v2.Target, []url.URL, error)
Nearest(service, typ, country string, lat, lon float64) ([]v2.Target, []url.URL, error)
heartbeat.StatusTracker
}

Expand Down Expand Up @@ -188,7 +188,8 @@ func (c *Client) Nearest(rw http.ResponseWriter, req *http.Request) {

// Find the nearest targets using the client parameters.
t := req.URL.Query().Get("machine-type")
targets, urls, err := c.LocatorV2.Nearest(service, t, lat, lon)
country := req.Header.Get("X-AppEngine-Country")
targets, urls, err := c.LocatorV2.Nearest(service, t, country, lat, lon)
if err != nil {
result.Error = v2.NewError("nearest", "Failed to lookup nearest machines", http.StatusInternalServerError)
writeResult(rw, result.Error.Status, &result)
Expand Down
2 changes: 1 addition & 1 deletion handler/handler_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ type fakeLocatorV2 struct {
urls []url.URL
}

func (l *fakeLocatorV2) Nearest(service, typ string, lat, lon float64) ([]v2.Target, []url.URL, error) {
func (l *fakeLocatorV2) Nearest(service, typ, country string, lat, lon float64) ([]v2.Target, []url.URL, error) {
if l.err != nil {
return nil, nil, l.err
}
Expand Down
23 changes: 17 additions & 6 deletions heartbeat/location.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,9 @@ func NewServerLocator(tracker StatusTracker) *Locator {

// Nearest discovers the nearest machines for the target service, using
// an exponentially distributed function based on distance.
func (l *Locator) Nearest(service, typ string, lat, lon float64) ([]v2.Target, []url.URL, error) {
func (l *Locator) Nearest(service, typ, country string, lat, lon float64) ([]v2.Target, []url.URL, error) {
// Filter.
sites := filterSites(service, typ, lat, lon, l.Instances())
sites := filterSites(service, typ, country, lat, lon, l.Instances())

// Sort.
sortSites(sites)
Expand All @@ -73,7 +73,7 @@ func (l *Locator) Nearest(service, typ string, lat, lon float64) ([]v2.Target, [

// filterSites groups the v2.HeartbeatMessage instances into sites and returns
// only those that can serve the client request.
func filterSites(service, typ string, lat, lon float64, instances map[string]v2.HeartbeatMessage) []site {
func filterSites(service, typ, country string, lat, lon float64, instances map[string]v2.HeartbeatMessage) []site {
m := make(map[string]*site)

for _, v := range instances {
Expand All @@ -86,7 +86,7 @@ func filterSites(service, typ string, lat, lon float64, instances map[string]v2.
s, ok := m[r.Site]
if !ok {
s = &site{
distance: distance,
distance: biasedDistance(country, r, distance),
registration: *r,
machines: make([]machine, 0),
}
Expand Down Expand Up @@ -133,8 +133,6 @@ func isValidInstance(service, typ string, lat, lon float64, v v2.HeartbeatMessag
return false, host.Name{}, 0
}

// TODO(cristinaleon): Add in-country biasing for distance.
// It might require implementing a reverse geocoder.
distance := mathx.GetHaversineDistance(lat, lon, r.Latitude, r.Longitude)
if distance > static.EarthHalfCircumferenceKm {
return false, host.Name{}, 0
Expand Down Expand Up @@ -213,3 +211,16 @@ func getURLs(service string, registration v2.Registration) []url.URL {

return result
}

func biasedDistance(country string, r *v2.Registration, distance float64) float64 {
// The 'ZZ' country code is used for unknown or unspecified countries.
if country == "" || country == "ZZ" {
return distance
}

if country == r.CountryCode {
return distance
}

return 2 * distance
}
72 changes: 70 additions & 2 deletions heartbeat/location_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -239,6 +239,7 @@ func TestNearest(t *testing.T) {
name string
service string
typ string
country string
lat float64
lon float64
instances []v2.HeartbeatMessage
Expand All @@ -250,6 +251,7 @@ func TestNearest(t *testing.T) {
name: "NDT7-any-type",
service: "ndt/ndt7",
typ: "",
country: "US",
lat: 43.1988,
lon: -75.3242,
expectedTargets: []v2.Target{virtualTarget, physicalTarget},
Expand All @@ -260,6 +262,7 @@ func TestNearest(t *testing.T) {
name: "NDT7-physical",
service: "ndt/ndt7",
typ: "physical",
country: "US",
lat: 43.1988,
lon: -75.3242,
expectedTargets: []v2.Target{physicalTarget},
Expand All @@ -270,6 +273,7 @@ func TestNearest(t *testing.T) {
name: "NDT7-virtual",
service: "ndt/ndt7",
typ: "virtual",
country: "US",
lat: 43.1988,
lon: -75.3242,
expectedTargets: []v2.Target{virtualTarget},
Expand All @@ -280,6 +284,7 @@ func TestNearest(t *testing.T) {
name: "wehe",
service: "wehe/replay",
typ: "",
country: "US",
lat: 43.1988,
lon: -75.3242,
expectedTargets: []v2.Target{weheTarget},
Expand All @@ -305,7 +310,7 @@ func TestNearest(t *testing.T) {
locator.UpdateHealth(i.Registration.Hostname, *i.Health)
}

gotTargets, gotURLs, err := locator.Nearest(tt.service, tt.typ, tt.lat, tt.lon)
gotTargets, gotURLs, err := locator.Nearest(tt.service, tt.typ, "", tt.lat, tt.lon)

if !reflect.DeepEqual(gotTargets, tt.expectedTargets) {
t.Errorf("Nearest() targets got: %+v, want %+v", gotTargets, tt.expectedTargets)
Expand Down Expand Up @@ -335,6 +340,7 @@ func TestFilterSites(t *testing.T) {
name string
service string
typ string
country string
lat float64
lon float64
expected []site
Expand All @@ -343,6 +349,7 @@ func TestFilterSites(t *testing.T) {
name: "NDT7-any-type",
service: "ndt/ndt7",
typ: "",
country: "US",
lat: 43.1988,
lon: -75.3242,
expected: []site{virtualSite, physicalSite},
Expand All @@ -351,6 +358,7 @@ func TestFilterSites(t *testing.T) {
name: "NDT7-physical",
service: "ndt/ndt7",
typ: "physical",
country: "US",
lat: 43.1988,
lon: -75.3242,
expected: []site{physicalSite},
Expand All @@ -359,6 +367,7 @@ func TestFilterSites(t *testing.T) {
name: "NDT7-virtual",
service: "ndt/ndt7",
typ: "virtual",
country: "US",
lat: 43.1988,
lon: -75.3242,
expected: []site{virtualSite},
Expand All @@ -367,6 +376,7 @@ func TestFilterSites(t *testing.T) {
name: "wehe",
service: "wehe/replay",
typ: "",
country: "US",
lat: 43.1988,
lon: -75.3242,
expected: []site{weheSite},
Expand All @@ -375,6 +385,7 @@ func TestFilterSites(t *testing.T) {
name: "too-far",
service: "ndt-ndt7",
typ: "",
country: "",
lat: 1000,
lon: 1000,
expected: []site{},
Expand All @@ -383,7 +394,7 @@ func TestFilterSites(t *testing.T) {

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := filterSites(tt.service, tt.typ, tt.lat, tt.lon, instances)
got := filterSites(tt.service, tt.typ, tt.country, tt.lat, tt.lon, instances)

sortSites(got)
for _, v := range got {
Expand Down Expand Up @@ -766,3 +777,60 @@ func TestPickWithProbability(t *testing.T) {
})
}
}

func TestBiasedDistance(t *testing.T) {
tests := []struct {
name string
country string
r *v2.Registration
distance float64
want float64
}{
{
name: "empty-country",
country: "",
r: &v2.Registration{
CountryCode: "foo",
},
distance: 100,
want: 100,
},
{
name: "unknown-country",
country: "ZZ",
r: &v2.Registration{
CountryCode: "foo",
},
distance: 100,
want: 100,
},
{
name: "same-country",
country: "foo",
r: &v2.Registration{
CountryCode: "foo",
},
distance: 100,
want: 100,
},
{
name: "different-country",
country: "bar",
r: &v2.Registration{
CountryCode: "foo",
},
distance: 100,
want: 200,
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := biasedDistance(tt.country, tt.r, tt.distance)

if got != tt.want {
t.Errorf("biasedDistance() got: %f, want: %f", got, tt.want)
}
})
}
}
2 changes: 1 addition & 1 deletion locatetest/locatetest.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ type LocatorV2 struct {
}

// Nearest returns the pre-configured LocatorV2 Servers or Err.
func (l *LocatorV2) Nearest(service, typ string, lat, lon float64) ([]v2.Target, []url.URL, error) {
func (l *LocatorV2) Nearest(service, typ, country string, lat, lon float64) ([]v2.Target, []url.URL, error) {
if l.Err != nil {
return nil, nil, l.Err
}
Expand Down

0 comments on commit 0687d2f

Please sign in to comment.