Merge pull request #15 from m-lab/sandbox-soltesz

Enable rate-limit using a threshold of 40 requests per day.

appengine/rate_table/rate_table.go

@@ -24,6 +24,11 @@ import (
 	"google.golang.org/appengine/memcache"
 )
 
+// threshold defines the maximum requests per day for a client before being rate
+// limited. Start with a conservative threshold. Ideally, we want a lower limit
+// for batch jobs than interactive users.
+const threshold = 40
+
 // 1.  Datastore stuff
 // 2.  Bigquery stuff
 // 3.  Bloom filter stuff
@@ -88,7 +93,6 @@ func NewDataset(ctx context.Context, project, dataset string, clientOpts ...opti
 func Update(w http.ResponseWriter, r *http.Request) {
 	// TODO - load threshold from flags or env-vars (see Peter's code?)
 	// TODO - move env var loading to init() ?
-	threshold := 12                             // requests per day
 	projectID, ok := os.LookupEnv("PROJECT_ID") // Datastore output project
 	if ok != true {
 		// metrics.FailCount.WithLabelValues("environ").Inc()
@@ -130,7 +134,7 @@ func Update(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	keys, endpoints, err := endpoint.MakeKeysAndStats(rows)
+	keys, endpoints, err := endpoint.MakeKeysAndStats(rows, threshold)
 	if err != nil {
 		logWarning(ctx, "MakeKeysAndStats: %v", err)
 		// metrics.FailCount.WithLabelValues("make").Inc()

endpoint/endpoint.go

@@ -36,17 +36,15 @@ type Stats struct {
 }
 
 // StatsFromMap creates a Key and Stats object from a bigquery result map.
-func StatsFromMap(row map[string]bigquery.Value) (string, Stats) {
+func StatsFromMap(row map[string]bigquery.Value, threshold int) (string, Stats) {
 	var stats Stats
 	rpd, ok := row["RequestsPerDay"]
 
 	if ok && rpd != nil {
 		stats.RequestsPerDay = rpd.(int64)
 		if stats.RequestsPerDay > 0 {
-			// TODO: make probability proportional to usage,
-			// e.g. 6.0 / float32(stats.RequestsPerDay)
-			// Probability of zero guarantees that all requests are offloaded or blocked.
-			stats.Probability = 0.0
+			// TODO: allow probability of zero to guarantee that all requests are blocked.
+			stats.Probability = float32(threshold) / float32(stats.RequestsPerDay)
 		}
 	}
 
@@ -90,12 +88,12 @@ func (ep *Stats) Save(ctx context.Context, client datastore.Client, key string)
 }
 
 // MakeKeysAndStats converts slice of bigquery rows into DSKeys and Stats objects.
-func MakeKeysAndStats(rows []map[string]bigquery.Value) ([]*datastore.Key, []Stats, error) {
+func MakeKeysAndStats(rows []map[string]bigquery.Value, threshold int) ([]*datastore.Key, []Stats, error) {
 	// preallocate to match number of rows, to avoid reallocation.
 	keys := make([]*datastore.Key, 0, len(rows))
 	endpoints := make([]Stats, 0, len(rows))
 	for i := range rows {
-		key, ep := StatsFromMap(rows[i])
+		key, ep := StatsFromMap(rows[i], threshold)
 		endpoints = append(endpoints, ep)
 		keys = append(keys, DSKey(key))
 	}
@@ -107,7 +105,7 @@ func MakeKeysAndStats(rows []map[string]bigquery.Value) ([]*datastore.Key, []Sta
 // endpoint signatures and request counts.
 // TODO - move the body (excluding simpleQuery) into go/bqext
 func FetchEndpointStats(ctx context.Context, dsExt *bqext.Dataset, threshold int) ([]map[string]bigquery.Value, error) {
-	qString := strings.Replace(sixHourQuery, "${THRESHOLD}", fmt.Sprint(threshold), 1)
+	qString := strings.Replace(simpleQuery, "${THRESHOLD}", fmt.Sprint(threshold), 1)
 	qString = strings.Replace(qString, "${DATE}", fmt.Sprint(threshold), 1)
 
 	query := dsExt.ResultQuery(qString, false)
@@ -261,7 +259,8 @@ FROM (
   WHERE
     (_table_suffix = FORMAT_DATE("%Y%m%d", CURRENT_DATE())
     OR _table_suffix = FORMAT_DATE("%Y%m%d", DATE_SUB(CURRENT_DATE(), INTERVAL 1 DAY)))
-    AND protoPayload.starttime > TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 1 DAY)
+		AND protoPayload.starttime > TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 1 DAY)
+		AND (REGEXP_CONTAINS(protoPayload.resource, '/neubot') OR REGEXP_CONTAINS(protoPayload.resource, '/ndt'))
   GROUP BY
     RequesterIP, userAgent, resource )
 WHERE
@@ -271,27 +270,3 @@ GROUP BY
 ORDER BY
   RequestsPerDay DESC
 LIMIT 20000`
-
-// sixHourQuery looks for clients that run every six hours and issue requests
-// to both /ndt and /neubot.
-var sixHourQuery = `
-SELECT
-    protoPayload.ip AS RequesterIP,
-    protoPayload.resource as resource,
-    protoPayload.userAgent as userAgent,
-    COUNT(*) as RequestsPerDay
-FROM
-  ` + "`mlab-ns.exports.appengine_googleapis_com_request_log_*`" + `
-WHERE
-     (_table_suffix = FORMAT_DATE("%Y%m%d", CURRENT_DATE())
-  OR  _table_suffix = FORMAT_DATE("%Y%m%d", DATE_SUB(CURRENT_DATE(), INTERVAL 1 DAY))
-  OR  _table_suffix = FORMAT_DATE("%Y%m%d", DATE_SUB(CURRENT_DATE(), INTERVAL 2 DAY))
-  AND protoPayload.starttime > TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 2 DAY))
-  AND (protoPayload.resource = '/neubot' OR protoPayload.resource = '/ndt')
-  AND protoPayload.userAgent IS NULL
-  AND protoPayload.ip IN ( SELECT ip FROM ` + "`mlab-ns.library.unique_ips_in_six_hour_periods`" + ` )
-GROUP BY
-  RequesterIP, protoPayload.resource, protoPayload.userAgent
-ORDER BY
-  RequesterIP, RequestsPerDay DESC
-`

endpoint/endpoint_test.go

@@ -82,7 +82,7 @@ func TestLiveBQQuery(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	keys, _, err := endpoint.MakeKeysAndStats(rows)
+	keys, _, err := endpoint.MakeKeysAndStats(rows, 20)
 	if err != nil {
 		t.Fatalf("Failed: %v", err)
 	}
@@ -102,7 +102,7 @@ func TestCreateTestEntries(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	keys, endpoints, err := endpoint.MakeKeysAndStats(rows)
+	keys, endpoints, err := endpoint.MakeKeysAndStats(rows, 20)
 	if err != nil {
 		t.Fatalf("Failed: %v", err)
 	}