Global ratelimiter: everything else

client/clientBean.go

@@ -41,6 +41,7 @@ type (
 	// Bean in an collection of clients
 	Bean interface {
 		GetHistoryClient() history.Client
+		GetHistoryPeers() history.PeerResolver
 		GetMatchingClient(domainIDToName DomainIDToNameFunc) (matching.Client, error)
 		GetFrontendClient() frontend.Client
 		GetRemoteAdminClient(cluster string) admin.Client
@@ -51,6 +52,7 @@ type (
 	clientBeanImpl struct {
 		sync.Mutex
 		historyClient         history.Client
+		historyPeers          history.PeerResolver
 		matchingClient        atomic.Value
 		frontendClient        frontend.Client
 		remoteAdminClients    map[string]admin.Client
@@ -62,7 +64,7 @@ type (
 // NewClientBean provides a collection of clients
 func NewClientBean(factory Factory, dispatcher *yarpc.Dispatcher, clusterMetadata cluster.Metadata) (Bean, error) {
 
-	historyClient, err := factory.NewHistoryClient()
+	historyClient, historyPeers, err := factory.NewHistoryClient()
 	if err != nil {
 		return nil, err
 	}
@@ -97,6 +99,7 @@ func NewClientBean(factory Factory, dispatcher *yarpc.Dispatcher, clusterMetadat
 	return &clientBeanImpl{
 		factory:               factory,
 		historyClient:         historyClient,
+		historyPeers:          historyPeers,
 		frontendClient:        remoteFrontendClients[clusterMetadata.GetCurrentClusterName()],
 		remoteAdminClients:    remoteAdminClients,
 		remoteFrontendClients: remoteFrontendClients,
@@ -107,6 +110,10 @@ func (h *clientBeanImpl) GetHistoryClient() history.Client {
 	return h.historyClient
 }
 
+func (h *clientBeanImpl) GetHistoryPeers() history.PeerResolver {
+	return h.historyPeers
+}
+
 func (h *clientBeanImpl) GetMatchingClient(domainIDToName DomainIDToNameFunc) (matching.Client, error) {
 	if client := h.matchingClient.Load(); client != nil {
 		return client.(matching.Client), nil

client/clientfactory.go

@@ -54,10 +54,10 @@ import (
 type (
 	// Factory can be used to create RPC clients for cadence services
 	Factory interface {
-		NewHistoryClient() (history.Client, error)
+		NewHistoryClient() (history.Client, history.PeerResolver, error)
 		NewMatchingClient(domainIDToName DomainIDToNameFunc) (matching.Client, error)
 
-		NewHistoryClientWithTimeout(timeout time.Duration) (history.Client, error)
+		NewHistoryClientWithTimeout(timeout time.Duration) (history.Client, history.PeerResolver, error)
 		NewMatchingClientWithTimeout(domainIDToName DomainIDToNameFunc, timeout time.Duration, longPollTimeout time.Duration) (matching.Client, error)
 
 		NewAdminClientWithTimeoutAndConfig(config transport.ClientConfig, timeout time.Duration, largeTimeout time.Duration) (admin.Client, error)
@@ -96,15 +96,15 @@ func NewRPCClientFactory(
 	}
 }
 
-func (cf *rpcClientFactory) NewHistoryClient() (history.Client, error) {
+func (cf *rpcClientFactory) NewHistoryClient() (history.Client, history.PeerResolver, error) {
 	return cf.NewHistoryClientWithTimeout(timeoutwrapper.HistoryDefaultTimeout)
 }
 
 func (cf *rpcClientFactory) NewMatchingClient(domainIDToName DomainIDToNameFunc) (matching.Client, error) {
 	return cf.NewMatchingClientWithTimeout(domainIDToName, timeoutwrapper.MatchingDefaultTimeout, timeoutwrapper.MatchingDefaultLongPollTimeout)
 }
 
-func (cf *rpcClientFactory) NewHistoryClientWithTimeout(timeout time.Duration) (history.Client, error) {
+func (cf *rpcClientFactory) NewHistoryClientWithTimeout(timeout time.Duration) (history.Client, history.PeerResolver, error) {
 	var rawClient history.Client
 	var namedPort = membership.PortTchannel
 
@@ -132,7 +132,7 @@ func (cf *rpcClientFactory) NewHistoryClientWithTimeout(timeout time.Duration) (
 		client = metered.NewHistoryClient(client, cf.metricsClient)
 	}
 	client = timeoutwrapper.NewHistoryClient(client, timeout)
-	return client, nil
+	return client, peerResolver, nil
 }
 
 func (cf *rpcClientFactory) NewMatchingClientWithTimeout(

client/history/client.go

@@ -22,6 +22,7 @@ package history
 
 import (
 	"context"
+	"fmt"
 	"math/rand"
 	"sync"
 	"time"
@@ -1079,6 +1080,22 @@ func (c *clientImpl) GetFailoverInfo(
 	return c.client.GetFailoverInfo(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
+func (c *clientImpl) RatelimitUpdate(ctx context.Context, request *types.RatelimitUpdateRequest, opts ...yarpc.CallOption) (*types.RatelimitUpdateResponse, error) {
+	if len(opts) == 0 {
+		// unfortunately there is not really any way to ensure "must have a shard key option"
+		// due to the closed nature of yarpc.CallOption's implementation, outside private-field-reading reflection.
+		//
+		// there are a few options to work around this, but they are currently rather high effort or
+		// run into import cycles or similar.  risk is low for now as there is only one caller, and likely
+		// never will be others.
+		return nil, fmt.Errorf("invalid arguments, missing yarpc.WithShardKey(peer) at a minimum")
+	}
+
+	// intentionally does not use peer-redirecting retries, as keys in this request
+	// could end up on multiple different hosts after a peer change.
+	return c.client.RatelimitUpdate(ctx, request, opts...)
+}
+
 func (c *clientImpl) executeWithRedirect(
 	ctx context.Context,
 	peer string,

client/history/client_test.go

@@ -454,6 +454,22 @@ func TestClient_withResponse(t *testing.T) {
 			},
 			want: &types.StartWorkflowExecutionResponse{},
 		},
+		{
+			name: "RatelimitUpdate",
+			op: func(c Client) (any, error) {
+				return c.RatelimitUpdate(context.Background(), &types.RatelimitUpdateRequest{
+					Any: &types.Any{
+						ValueType: "something",
+						Value:     []byte("data"),
+					},
+				}, yarpc.WithShardKey("test-peer"))
+			},
+			mock: func(p *MockPeerResolver, c *MockClient) {
+				c.EXPECT().RatelimitUpdate(gomock.Any(), gomock.Any(), []yarpc.CallOption{yarpc.WithShardKey("test-peer")}).
+					Return(&types.RatelimitUpdateResponse{}, nil).Times(1)
+			},
+			want: &types.RatelimitUpdateResponse{},
+		},
 		{
 			name: "StartWorkflowExecution peer resolve failure",
 			op: func(c Client) (any, error) {
@@ -906,6 +922,23 @@ func TestClient_withResponse(t *testing.T) {
 			},
 			wantError: true,
 		},
+		{
+			name: "RatelimitUpdate requires explicit shard key arg",
+			op: func(c Client) (any, error) {
+				// same as successful call...
+				return c.RatelimitUpdate(context.Background(), &types.RatelimitUpdateRequest{
+					// Peer: "", // intentionally the zero value
+					Any: &types.Any{
+						ValueType: "something",
+						Value:     []byte("data"),
+					},
+				})
+			},
+			mock: func(p *MockPeerResolver, c *MockClient) {
+				// no calls expected
+			},
+			wantError: true,
+		},
 	}
 
 	for _, tt := range tests {

client/history/interface.go

@@ -81,4 +81,12 @@ type Client interface {
 	SyncShardStatus(context.Context, *types.SyncShardStatusRequest, ...yarpc.CallOption) error
 	TerminateWorkflowExecution(context.Context, *types.HistoryTerminateWorkflowExecutionRequest, ...yarpc.CallOption) error
 	GetFailoverInfo(context.Context, *types.GetFailoverInfoRequest, ...yarpc.CallOption) (*types.GetFailoverInfoResponse, error)
+
+	// RatelimitUpdate pushes usage info for the passed ratelimit keys, and requests updated weight info from aggregating hosts.
+	// Exact semantics beyond this depend on the load-balanced ratelimit implementation.
+	//
+	// A peer (via yarpc.WithShardkey) MUST be determined before calling and passed in yarpc opts,
+	// and unlike most endpoints this will NOT be forwarded to a new peer if the ring membership changes.
+	// To correctly forward keys to the new hosts, they must be re-sharded to find their new hosts.
+	RatelimitUpdate(ctx context.Context, request *types.RatelimitUpdateRequest, opts ...yarpc.CallOption) (*types.RatelimitUpdateResponse, error)
 }

client/history/peer_resolver.go

@@ -21,6 +21,10 @@
 package history
 
 import (
+	"fmt"
+
+	"go.uber.org/yarpc"
+
 	"github.com/uber/cadence/common"
 	"github.com/uber/cadence/common/membership"
 	"github.com/uber/cadence/common/service"
@@ -38,6 +42,20 @@ type PeerResolver interface {
 	FromShardID(shardID int) (string, error)
 	FromHostAddress(hostAddress string) (string, error)
 	GetAllPeers() ([]string, error)
+
+	// GlobalRatelimitPeers partitions the ratelimit keys into map[yarpc peer][]limits_for_peer
+	GlobalRatelimitPeers(ratelimits []string) (ratelimitsByPeer map[Peer][]string, err error)
+}
+
+// Peer is used to mark a string as the routing information to a peer process.
+//
+// This is essentially the host:port address of the peer to be contacted,
+// but it is meant to be treated as an opaque blob until given to yarpc
+// via ToYarpcShardKey.
+type Peer string
+
+func (s Peer) ToYarpcShardKey() yarpc.CallOption {
+	return yarpc.WithShardKey(string(s))
 }
 
 type peerResolver struct {
@@ -110,3 +128,54 @@ func (pr peerResolver) GetAllPeers() ([]string, error) {
 	}
 	return peers, nil
 }
+
+func (pr peerResolver) GlobalRatelimitPeers(ratelimits []string) (map[Peer][]string, error) {
+	// History was selected simply because it already has a ring and an internal-only API.
+	// Any service should be fine, it just needs to be shared by both ends of the system.
+	hosts, err := pr.resolver.Members(service.History)
+	if err != nil {
+		return nil, fmt.Errorf("unable to get history peers: %w", err)
+	}
+	if len(hosts) == 0 {
+		// can occur when shutting down the only instance because it calls EvictSelf ASAP.
+		// this is common in dev, but otherwise *probably* should not be possible in a healthy system.
+		return nil, fmt.Errorf("unable to get history peers: no peers available")
+	}
+
+	results := make(map[Peer][]string, len(hosts))
+	initialCapacity := len(ratelimits) / len(hosts)
+	// add a small buffer to reduce copies, as this is only an estimate
+	initialCapacity += initialCapacity / 10
+	// but don't use zero, that'd be pointless
+	if initialCapacity < 1 {
+		initialCapacity = 1
+	}
+
+	for _, r := range ratelimits {
+		// figure out the destination for this ratelimit
+		host, err := pr.resolver.Lookup(service.History, r)
+		if err != nil {
+			return nil, fmt.Errorf(
+				"unable to find host for ratelimit key %q: %w", r, err)
+		}
+		peer, err := host.GetNamedAddress(pr.namedPort)
+		if err != nil {
+			// AFAICT should not happen unless the peer data is incomplete / the ring has no port info.
+			// retry may work?  eventually?
+			return nil, fmt.Errorf(
+				"unable to get address from host: %s, %w",
+				host.String(), // HostInfo has a readable String(), it's good enough
+				err,
+			)
+		}
+
+		// add to the peer's partition
+		current := results[Peer(peer)]
+		if len(current) == 0 {
+			current = make([]string, 0, initialCapacity)
+		}
+		current = append(current, r)
+		results[Peer(peer)] = current
+	}
+	return results, nil
+}