Add in-country biasing for Locate V2 (#94)

* Add in-country biasing

* Make test parameter explicit

* Move check

* Remove unnecessary line

Author: cristinaleonr

handler/handler.go

@@ -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
 }
 
@@ -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)

handler/handler_test.go

@@ -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
 	}

heartbeat/location.go

@@ -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)
@@ -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 {
@@ -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),
 			}
@@ -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
@@ -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
+}

heartbeat/location_test.go

@@ -239,6 +239,7 @@ func TestNearest(t *testing.T) {
 		name            string
 		service         string
 		typ             string
+		country         string
 		lat             float64
 		lon             float64
 		instances       []v2.HeartbeatMessage
@@ -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},
@@ -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},
@@ -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},
@@ -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},
@@ -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)
@@ -335,6 +340,7 @@ func TestFilterSites(t *testing.T) {
 		name     string
 		service  string
 		typ      string
+		country  string
 		lat      float64
 		lon      float64
 		expected []site
@@ -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},
@@ -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},
@@ -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},
@@ -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},
@@ -375,6 +385,7 @@ func TestFilterSites(t *testing.T) {
 			name:     "too-far",
 			service:  "ndt-ndt7",
 			typ:      "",
+			country:  "",
 			lat:      1000,
 			lon:      1000,
 			expected: []site{},
@@ -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 {
@@ -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)
+			}
+		})
+	}
+}

locatetest/locatetest.go

@@ -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
 	}