[common-go] Rate limiting on booleans and composite keys

Author: WVerlaek

components/common-go/grpc/ratelimit.go

@@ -33,8 +33,12 @@ type RateLimit struct {
 	// so effectively this is the rate at which requests can be made.
 	RefillInterval util.Duration `json:"refillInterval"`
 
-	KeyCacheSize uint   `json:"keyCacheSize,omitempty"`
-	Key          string `json:"key,omitempty"`
+	// Key is the proto field name to rate limit on. Each unique value of this
+	// field gets its own rate limit bucket. Must be a String, Enum, or Boolean field.
+	// Can be a composite key by separating fields by comma, e.g. `foo.bar,foo.baz`
+	Key string `json:"key,omitempty"`
+	// KeyCacheSize is the max number of buckets kept in a LRU cache.
+	KeyCacheSize uint `json:"keyCacheSize,omitempty"`
 }
 
 func (r RateLimit) Limiter() *rate.Limiter {
@@ -88,11 +92,20 @@ func fieldAccessKey(key string) keyFunc {
 			return "", status.Errorf(codes.Internal, "request was not a protobuf message")
 		}
 
-		val, ok := getFieldValue(msg.ProtoReflect(), strings.Split(key, "."))
-		if !ok {
-			return "", status.Errorf(codes.Internal, "Field %s does not exist in message. This is a rate limiting configuration error.", key)
+		fields := strings.Split(key, ",")
+		var composite string
+		for _, field := range fields {
+			val, ok := getFieldValue(msg.ProtoReflect(), strings.Split(field, "."))
+			if !ok {
+				return "", status.Errorf(codes.Internal, "Field %s does not exist in message. This is a rate limiting configuration error.", field)
+			}
+			// It's technically possible that `|` is part of one of the field values, and therefore could cause collisions
+			// in composite keys, e.g. values (`a|`, `b`), and (`a`, `|b`) would result in the same composite key `a||b`
+			// and share the rate limit. This is currently not possible though given the fields we rate limit on.
+			composite += "|" + val
 		}
-		return val, nil
+
+		return composite, nil
 	}
 }
 
@@ -120,6 +133,13 @@ func getFieldValue(msg protoreflect.Message, path []string) (val string, ok bool
 	case protoreflect.EnumKind:
 		enumNum := msg.Get(field).Enum()
 		return strconv.Itoa(int(enumNum)), true
+	case protoreflect.BoolKind:
+		if msg.Get(field).Bool() {
+			return "t", true
+		} else {
+			return "f", true
+		}
+
 	default:
 		// we only support string and enum fields
 		return "", false

components/common-go/grpc/ratelimit_test.go

@@ -51,6 +51,18 @@ func TestGetFieldValue(t *testing.T) {
 			Path:        "syntax",
 			Expectation: Expectation{Found: true, Val: "0"},
 		},
+		{
+			Name:        "bool field",
+			Message:     &apipb.Method{RequestStreaming: true},
+			Path:        "request_streaming",
+			Expectation: Expectation{Found: true, Val: "t"},
+		},
+		{
+			Name:        "empty bool field",
+			Message:     &apipb.Method{},
+			Path:        "request_streaming",
+			Expectation: Expectation{Found: true, Val: "f"},
+		},
 		{
 			Name:        "non-existent field",
 			Message:     &apipb.Api{},
@@ -79,6 +91,41 @@ func TestGetFieldValue(t *testing.T) {
 	}
 }
 
+func TestFieldAccessKey(t *testing.T) {
+	type Expectation struct {
+		Val string
+		Err error
+	}
+	tests := []struct {
+		Name        string
+		Message     proto.Message
+		Key         string
+		Expectation Expectation
+	}{
+		{
+			Name: "composite key",
+			Message: &apipb.Api{
+				SourceContext: &sourcecontextpb.SourceContext{
+					FileName: "bar",
+				},
+				Syntax: typepb.Syntax_SYNTAX_PROTO3,
+			},
+			Key:         "source_context.file_name,syntax",
+			Expectation: Expectation{Val: "|bar|1", Err: nil},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.Name, func(t *testing.T) {
+			var act Expectation
+			keyFn := fieldAccessKey(test.Key)
+			act.Val, act.Err = keyFn(test.Message)
+			if diff := cmp.Diff(test.Expectation, act); diff != "" {
+				t.Errorf("unexpected fieldAccessKey (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
 func BenchmarkGetFieldValue(b *testing.B) {
 	msg := apipb.Api{
 		SourceContext: &sourcecontextpb.SourceContext{
@@ -92,3 +139,56 @@ func BenchmarkGetFieldValue(b *testing.B) {
 		getFieldValue(msgr, path)
 	}
 }
+
+func BenchmarkFieldAccessKey_String(b *testing.B) {
+	msg := &apipb.Api{
+		Name: "bar",
+	}
+	keyFn := fieldAccessKey("name")
+	for n := 0; n < b.N; n++ {
+		if _, err := keyFn(msg); err != nil {
+			b.Logf("failed to access key: %v", err)
+			b.Fail()
+		}
+	}
+}
+
+func BenchmarkFieldAccessKey_Enum(b *testing.B) {
+	msg := &apipb.Api{
+		Syntax: typepb.Syntax_SYNTAX_PROTO3,
+	}
+	keyFn := fieldAccessKey("syntax")
+	for n := 0; n < b.N; n++ {
+		if _, err := keyFn(msg); err != nil {
+			b.Logf("failed to access key: %v", err)
+			b.Fail()
+		}
+	}
+}
+
+func BenchmarkFieldAccessKey_Bool(b *testing.B) {
+	msg := &apipb.Method{
+		RequestStreaming: true,
+	}
+	keyFn := fieldAccessKey("request_streaming")
+	for n := 0; n < b.N; n++ {
+		if _, err := keyFn(msg); err != nil {
+			b.Logf("failed to access key: %v", err)
+			b.Fail()
+		}
+	}
+}
+
+func BenchmarkFieldAccessKey_Composite(b *testing.B) {
+	msg := &apipb.Method{
+		Name:             "bar",
+		RequestStreaming: true,
+	}
+	keyFn := fieldAccessKey("name,request_streaming")
+	for n := 0; n < b.N; n++ {
+		if _, err := keyFn(msg); err != nil {
+			b.Logf("failed to access key: %v", err)
+			b.Fail()
+		}
+	}
+}