Add -early-exit-clients command-line flag (#223)

* Add -early-exit-client command-line argument.

* Make flag configurable via cloudbuild vars

* Use consts for early exit parameter and value

Author: robertodauria

cloudbuild/app.yaml.mlab-ns.template

@@ -44,3 +44,4 @@ env_variables:
   RATE_LIMIT_IP_MAX: {{RATE_LIMIT_IP_MAX}}
   PROMETHEUSX_LISTEN_ADDRESS: ':9090' # Must match one of the forwarded_ports above.
   PROMETHEUS_URL: 'https://prometheus-basicauth.{{PLATFORM_PROJECT}}.measurementlab.net/'
+  EARLY_EXIT_CLIENTS: {{EARLY_EXIT_CLIENTS}}

cloudbuild/app.yaml.template

@@ -43,3 +43,4 @@ env_variables:
   RATE_LIMIT_IP_MAX: {{RATE_LIMIT_IP_MAX}}
   PROMETHEUSX_LISTEN_ADDRESS: ':9090' # Must match one of the forwarded_ports above.
   PROMETHEUS_URL: 'https://prometheus-basicauth.{{PLATFORM_PROJECT}}.measurementlab.net/'
+  EARLY_EXIT_CLIENTS: {{EARLY_EXIT_CLIENTS}}

cloudbuild/cloudbuild.yaml

@@ -34,6 +34,7 @@ steps:
         -e 's/{{RATE_LIMIT_MAX}}/$_RATE_LIMIT_MAX/'
         -e 's/{{RATE_LIMIT_IP_INTERVAL}}/$_RATE_LIMIT_IP_INTERVAL/'
         -e 's/{{RATE_LIMIT_IP_MAX}}/$_RATE_LIMIT_IP_MAX/'
+        -e 's/{{EARLY_EXIT_CLIENTS}}/$_EARLY_EXIT_CLIENTS/'
         app.yaml
       - gcloud --project $PROJECT_ID app deploy --promote app.yaml
       # After deploying the new service, deploy the openapi spec.
@@ -57,6 +58,7 @@ steps:
         -e 's/{{RATE_LIMIT_MAX}}/$_RATE_LIMIT_MAX/'
         -e 's/{{RATE_LIMIT_IP_INTERVAL}}/$_RATE_LIMIT_IP_INTERVAL/'
         -e 's/{{RATE_LIMIT_IP_MAX}}/$_RATE_LIMIT_IP_MAX/'
+        -e 's/{{EARLY_EXIT_CLIENTS}}/$_EARLY_EXIT_CLIENTS/'
         app.yaml
       - gcloud --project $PROJECT_ID app deploy --no-promote app.yaml
       # After deploying the new service, deploy the openapi spec.

handler/handler.go

@@ -54,9 +54,10 @@ type Client struct {
 	LocatorV2
 	ClientLocator
 	PrometheusClient
-	targetTmpl  *template.Template
-	agentLimits limits.Agents
-	ipLimiter   Limiter
+	targetTmpl       *template.Template
+	agentLimits      limits.Agents
+	ipLimiter        Limiter
+	earlyExitClients map[string]bool
 }
 
 // LocatorV2 defines how the Nearest handler requests machines nearest to the
@@ -90,7 +91,12 @@ func init() {
 
 // NewClient creates a new client.
 func NewClient(project string, private Signer, locatorV2 LocatorV2, client ClientLocator,
-	prom PrometheusClient, lmts limits.Agents, limiter Limiter) *Client {
+	prom PrometheusClient, lmts limits.Agents, limiter Limiter, earlyExitClients []string) *Client {
+	// Convert slice to map for O(1) lookups
+	earlyExitMap := make(map[string]bool)
+	for _, client := range earlyExitClients {
+		earlyExitMap[client] = true
+	}
 	return &Client{
 		Signer:           private,
 		project:          project,
@@ -100,6 +106,7 @@ func NewClient(project string, private Signer, locatorV2 LocatorV2, client Clien
 		targetTmpl:       template.Must(template.New("name").Parse("{{.Hostname}}{{.Ports}}")),
 		agentLimits:      lmts,
 		ipLimiter:        limiter,
+		earlyExitClients: earlyExitMap,
 	}
 }
 
@@ -116,8 +123,9 @@ func NewClientDirect(project string, private Signer, locatorV2 LocatorV2, client
 	}
 }
 
-func extraParams(hostname string, index int, p paramOpts) url.Values {
+func (c *Client) extraParams(hostname string, index int, p paramOpts) url.Values {
 	v := url.Values{}
+
 	// Add client parameters.
 	for key := range p.raw {
 		if strings.HasPrefix(key, "client_") {
@@ -131,6 +139,12 @@ func extraParams(hostname string, index int, p paramOpts) url.Values {
 		}
 	}
 
+	// Add early_exit parameter for specified clients
+	clientName := p.raw.Get("client_name")
+	if clientName != "" && c.earlyExitClients[clientName] {
+		v.Set(static.EarlyExitParameter, static.EarlyExitDefaultValue)
+	}
+
 	// Add Locate Service version.
 	v.Set("locate_version", p.version)
 
@@ -323,7 +337,7 @@ func (c *Client) checkClientLocation(rw http.ResponseWriter, req *http.Request)
 func (c *Client) populateURLs(targets []v2.Target, ports static.Ports, exp string, pOpts paramOpts) {
 	for i, target := range targets {
 		token := c.getAccessToken(target.Machine, exp)
-		params := extraParams(target.Machine, i, pOpts)
+		params := c.extraParams(target.Machine, i, pOpts)
 		targets[i].URLs = c.getURLs(ports, target.Hostname, token, params)
 	}
 }

handler/handler_test.go

@@ -338,7 +338,7 @@ func TestClient_Nearest(t *testing.T) {
 			if tt.cl == nil {
 				tt.cl = clientgeo.NewAppEngineLocator()
 			}
-			c := NewClient(tt.project, tt.signer, tt.locator, tt.cl, prom.NewAPI(nil), tt.limits, tt.ipLimiter)
+			c := NewClient(tt.project, tt.signer, tt.locator, tt.cl, prom.NewAPI(nil), tt.limits, tt.ipLimiter, nil)
 
 			mux := http.NewServeMux()
 			mux.HandleFunc("/v2/nearest/", c.Nearest)
@@ -417,7 +417,7 @@ func TestClient_Ready(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			c := NewClient("foo", &fakeSigner{}, &fakeLocatorV2{StatusTracker: &heartbeattest.FakeStatusTracker{Err: tt.fakeErr}}, nil, nil, nil, nil)
+			c := NewClient("foo", &fakeSigner{}, &fakeLocatorV2{StatusTracker: &heartbeattest.FakeStatusTracker{Err: tt.fakeErr}}, nil, nil, nil, nil, nil)
 
 			mux := http.NewServeMux()
 			mux.HandleFunc("/ready/", c.Ready)
@@ -475,7 +475,7 @@ func TestClient_Registrations(t *testing.T) {
 		}
 
 		t.Run(tt.name, func(t *testing.T) {
-			c := NewClient("foo", &fakeSigner{}, &fakeLocatorV2{StatusTracker: fakeStatusTracker}, nil, nil, nil, nil)
+			c := NewClient("foo", &fakeSigner{}, &fakeLocatorV2{StatusTracker: fakeStatusTracker}, nil, nil, nil, nil, nil)
 
 			mux := http.NewServeMux()
 			mux.HandleFunc("/v2/siteinfo/registrations/", c.Registrations)
@@ -499,6 +499,7 @@ func TestExtraParams(t *testing.T) {
 		hostname             string
 		index                int
 		p                    paramOpts
+		client               *Client
 		earlyExitProbability float64
 		want                 url.Values
 	}{
@@ -512,13 +513,55 @@ func TestExtraParams(t *testing.T) {
 				ranks:     map[string]int{"host": 0},
 				svcParams: map[string]float64{},
 			},
+			client: &Client{},
 			want: url.Values{
 				"client_name":    []string{"client"},
 				"locate_version": []string{"v2"},
 				"metro_rank":     []string{"0"},
 				"index":          []string{"0"},
 			},
 		},
+		{
+			name:     "early-exit-client-match",
+			hostname: "host",
+			index:    0,
+			p: paramOpts{
+				raw:       map[string][]string{"client_name": {"foo"}},
+				version:   "v2",
+				ranks:     map[string]int{"host": 0},
+				svcParams: map[string]float64{},
+			},
+			client: &Client{
+				earlyExitClients: map[string]bool{"foo": true},
+			},
+			want: url.Values{
+				"client_name":    []string{"foo"},
+				"locate_version": []string{"v2"},
+				"metro_rank":     []string{"0"},
+				"index":          []string{"0"},
+				"early_exit":     []string{"250"},
+			},
+		},
+		{
+			name:     "early-exit-client-no-match",
+			hostname: "host",
+			index:    0,
+			p: paramOpts{
+				raw:       map[string][]string{"client_name": {"bar"}},
+				version:   "v2",
+				ranks:     map[string]int{"host": 0},
+				svcParams: map[string]float64{},
+			},
+			client: &Client{
+				earlyExitClients: map[string]bool{"foo": true},
+			},
+			want: url.Values{
+				"client_name":    []string{"bar"},
+				"locate_version": []string{"v2"},
+				"metro_rank":     []string{"0"},
+				"index":          []string{"0"},
+			},
+		},
 		{
 			name:     "no-client",
 			hostname: "host",
@@ -626,7 +669,7 @@ func TestExtraParams(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := extraParams(tt.hostname, tt.index, tt.p)
+			got := tt.client.extraParams(tt.hostname, tt.index, tt.p)
 			if !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("extraParams() = %v, want %v", got, tt.want)
 			}

handler/heartbeat_test.go

@@ -70,7 +70,7 @@ func TestClient_handleHeartbeats(t *testing.T) {
 func fakeClient(t heartbeat.StatusTracker) *Client {
 	locatorv2 := fakeLocatorV2{StatusTracker: t}
 	return NewClient("mlab-sandbox", &fakeSigner{}, &locatorv2,
-		clientgeo.NewAppEngineLocator(), prom.NewAPI(nil), nil, nil)
+		clientgeo.NewAppEngineLocator(), prom.NewAPI(nil), nil, nil, nil)
 }
 
 type fakeConn struct {

handler/monitoring_test.go

@@ -92,7 +92,7 @@ func TestClient_Monitoring(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			cl := clientgeo.NewAppEngineLocator()
-			c := NewClient("mlab-sandbox", tt.signer, tt.locator, cl, prom.NewAPI(nil), nil, nil)
+			c := NewClient("mlab-sandbox", tt.signer, tt.locator, cl, prom.NewAPI(nil), nil, nil, nil)
 			rw := httptest.NewRecorder()
 			req := httptest.NewRequest(http.MethodGet, "/v2/platform/monitoring/"+tt.path, nil)
 			req = req.Clone(controller.SetClaim(req.Context(), tt.claim))

locate.go

@@ -60,6 +60,8 @@ var (
 	rateLimitIPUAMax      int
 	rateLimitIPInterval   time.Duration
 	rateLimitIPMax        int
+
+	earlyExitClients flagx.StringArray
 )
 
 func init() {
@@ -88,6 +90,7 @@ func init() {
 		"Max number of events in the time window for IP-only rate limiting")
 	flag.StringVar(&rateLimitPrefix, "rate-limit-prefix", "locate:ratelimit", "Prefix for Redis keys for IP+UA rate limiting")
 	flag.StringVar(&rateLimitRedisAddr, "rate-limit-redis-address", "", "Primary endpoint for Redis instance for rate limiting")
+	flag.Var(&earlyExitClients, "early-exit-clients", "Client names that should receive early_exit parameter (can be specified multiple times)")
 	// Enable logging with line numbers to trace error locations.
 	log.SetFlags(log.LUTC | log.Llongfile)
 }
@@ -173,7 +176,8 @@ func main() {
 
 	lmts, err := limits.ParseConfig(limitsPath)
 	rtx.Must(err, "failed to parse limits config")
-	c := handler.NewClient(project, signer, srvLocatorV2, locators, promClient, lmts, ipLimiter)
+	c := handler.NewClient(project, signer, srvLocatorV2, locators, promClient,
+		lmts, ipLimiter, earlyExitClients)
 
 	go func() {
 		// Check and reload db at least once a day.

static/configs.go

@@ -30,6 +30,7 @@ const (
 	RegistrationLoadMax        = 24 * time.Hour
 	EarthHalfCircumferenceKm   = 20038
 	EarlyExitParameter         = "early_exit"
+	EarlyExitDefaultValue      = "250"
 	MaxCwndGainParameter       = "max_cwnd_gain"
 	MaxElapsedTimeParameter    = "max_elapsed_time"
 	CountryParameter           = "country"