Add test for global rate limiting with load balancing

functional_test.go

@@ -21,6 +21,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"math/rand"
 	"net/http"
 	"os"
 	"strings"
@@ -35,6 +36,10 @@ import (
 	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/resolver"
 	json "google.golang.org/protobuf/encoding/protojson"
 )
 
@@ -1017,6 +1022,70 @@ func TestGlobalRateLimits(t *testing.T) {
 	sendHit(peers[4].MustClient(), guber.Status_OVER_LIMIT, 1, 0)
 }
 
+// Ensure global broadcast updates all peers when GetRateLimits is called on
+// either owner or non-owner peer.
+func TestGlobalRateLimitsWithLoadBalancing(t *testing.T) {
+	ctx := context.Background()
+	const name = "test_global"
+	key := fmt.Sprintf("key:%016x", rand.Int())
+
+	// Determine owner and non-owner peers.
+	ownerPeerInfo, err := cluster.FindOwningPeer(name, key)
+	require.NoError(t, err)
+	owner := ownerPeerInfo.GRPCAddress
+	nonOwner := cluster.PeerAt(0).GRPCAddress
+	if nonOwner == owner {
+		nonOwner = cluster.PeerAt(1).GRPCAddress
+	}
+	require.NotEqual(t, owner, nonOwner)
+
+	// Connect to owner and non-owner peers in round robin.
+	dialOpts := []grpc.DialOption{
+		grpc.WithResolvers(newStaticBuilder()),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	address := fmt.Sprintf("static:///%s,%s", owner, nonOwner)
+	conn, err := grpc.DialContext(ctx, address, dialOpts...)
+	require.NoError(t, err)
+	client := guber.NewV1Client(conn)
+
+	sendHit := func(status guber.Status, i int) {
+		ctx, cancel := context.WithTimeout(ctx, 10*clock.Second)
+		defer cancel()
+		resp, err := client.GetRateLimits(ctx, &guber.GetRateLimitsReq{
+			Requests: []*guber.RateLimitReq{
+				{
+					Name:      name,
+					UniqueKey: key,
+					Algorithm: guber.Algorithm_LEAKY_BUCKET,
+					Behavior:  guber.Behavior_GLOBAL,
+					Duration:  guber.Minute * 5,
+					Hits:      1,
+					Limit:     2,
+				},
+			},
+		})
+		require.NoError(t, err, i)
+		item := resp.Responses[0]
+		assert.Equal(t, "", item.GetError(), fmt.Sprintf("mismatch error, iteration %d", i))
+		assert.Equal(t, status, item.GetStatus(), fmt.Sprintf("mismatch status, iteration %d", i))
+	}
+
+	// Send two hits that should be processed by the owner and non-owner and
+	// deplete the limit consistently.
+	sendHit(guber.Status_UNDER_LIMIT, 1)
+	sendHit(guber.Status_UNDER_LIMIT, 2)
+
+	// Sleep to ensure the global broadcast occurs (every 100ms).
+	time.Sleep(150 * time.Millisecond)
+
+	// All successive hits should return OVER_LIMIT.
+	for i := 2; i <= 10; i++ {
+		sendHit(guber.Status_OVER_LIMIT, i)
+	}
+}
+
 func TestGlobalRateLimitsPeerOverLimit(t *testing.T) {
 	const (
 		name = "test_global_token_limit"
@@ -1768,3 +1837,46 @@ func waitForBroadcast(timeout clock.Duration, d *guber.Daemon, expect int) error
 		}
 	}
 }
+
+// staticBuilder implements the `resolver.Builder` interface.
+type staticBuilder struct{}
+
+func newStaticBuilder() resolver.Builder {
+	return &staticBuilder{}
+}
+
+func (sb *staticBuilder) Build(target resolver.Target, cc resolver.ClientConn, _ resolver.BuildOptions) (resolver.Resolver, error) {
+	var resolverAddrs []resolver.Address
+	for _, address := range strings.Split(target.Endpoint(), ",") {
+		resolverAddrs = append(resolverAddrs, resolver.Address{
+			Addr:       address,
+			ServerName: address,
+		})
+
+	}
+	r, err := newStaticResolver(cc, resolverAddrs)
+	if err != nil {
+		return nil, err
+	}
+	return r, nil
+}
+
+func (sb *staticBuilder) Scheme() string {
+	return "static"
+}
+
+type staticResolver struct {
+	cc resolver.ClientConn
+}
+
+func newStaticResolver(cc resolver.ClientConn, addresses []resolver.Address) (resolver.Resolver, error) {
+	err := cc.UpdateState(resolver.State{Addresses: addresses})
+	if err != nil {
+		return nil, err
+	}
+	return &staticResolver{cc: cc}, nil
+}
+
+func (sr *staticResolver) ResolveNow(_ resolver.ResolveNowOptions) {}
+
+func (sr *staticResolver) Close() {}