dynamic partitioning: use probability of AddTask to root partition to estimate overall add task rate

client/matching/client.go

@@ -175,10 +175,11 @@ func (c *clientImpl) addActivityTask(
 	opts []grpc.CallOption,
 ) (*matchingservice.AddActivityTaskResponse, error) {
 	request = common.CloneProto(request)
-	client, err := c.pickClientForWrite(request.GetTaskQueue(), p, loadBalance, pc)
+	client, estimatedTasksAllPartitions, err := c.pickClientForWrite(request.GetTaskQueue(), p, loadBalance, pc)
 	if err != nil {
 		return nil, err
 	}
+	ctx = appendEstimatedTasksAllPartitions(ctx, estimatedTasksAllPartitions, p.IsRoot())
 	ctx, cancel := c.createContext(ctx)
 	defer cancel()
 
@@ -207,10 +208,11 @@ func (c *clientImpl) addWorkflowTask(
 	opts []grpc.CallOption,
 ) (*matchingservice.AddWorkflowTaskResponse, error) {
 	request = common.CloneProto(request)
-	client, err := c.pickClientForWrite(request.GetTaskQueue(), p, loadBalance, pc)
+	client, estimatedTasksAllPartitions, err := c.pickClientForWrite(request.GetTaskQueue(), p, loadBalance, pc)
 	if err != nil {
 		return nil, err
 	}
+	ctx = appendEstimatedTasksAllPartitions(ctx, estimatedTasksAllPartitions, p.IsRoot())
 	ctx, cancel := c.createContext(ctx)
 	defer cancel()
 	return client.AddWorkflowTask(ctx, request, opts...)
@@ -317,10 +319,11 @@ func (c *clientImpl) queryWorkflow(
 		ForwardInfo:      request.ForwardInfo,
 		Priority:         request.Priority,
 	}
-	client, err := c.pickClientForWrite(request.GetTaskQueue(), p, loadBalance, pc)
+	client, estimatedTasksAllPartitions, err := c.pickClientForWrite(request.GetTaskQueue(), p, loadBalance, pc)
 	if err != nil {
 		return nil, err
 	}
+	ctx = appendEstimatedTasksAllPartitions(ctx, estimatedTasksAllPartitions, p.IsRoot())
 	ctx, cancel := c.createContext(ctx)
 	defer cancel()
 	return client.QueryWorkflow(ctx, request, opts...)
@@ -355,10 +358,11 @@ func (c *clientImpl) dispatchNexusTask(
 		Request:     request.Request,
 		ForwardInfo: request.ForwardInfo,
 	}
-	client, err := c.pickClientForWrite(request.GetTaskQueue(), p, loadBalance, pc)
+	client, estimatedTasksAllPartitions, err := c.pickClientForWrite(request.GetTaskQueue(), p, loadBalance, pc)
 	if err != nil {
 		return nil, err
 	}
+	ctx = appendEstimatedTasksAllPartitions(ctx, estimatedTasksAllPartitions, p.IsRoot())
 	ctx, cancel := c.createContext(ctx)
 	defer cancel()
 	return client.DispatchNexusTask(ctx, request, opts...)
@@ -422,12 +426,14 @@ func (c *clientImpl) pickClientForWrite(
 	p tqid.Partition,
 	loadBalance bool,
 	pc PartitionCounts,
-) (matchingservice.MatchingServiceClient, error) {
+) (matchingservice.MatchingServiceClient, int, error) {
+	estimatedTasksAllPartitions := 0
 	if loadBalance {
-		p = c.loadBalancer.PickWritePartition(p.TaskQueue(), pc)
+		p, estimatedTasksAllPartitions = c.loadBalancer.PickWritePartition(p.TaskQueue(), pc)
 	}
 	proto.Name = p.RpcName()
-	return c.getClientForTaskQueuePartition(p)
+	client, err := c.getClientForTaskQueuePartition(p)
+	return client, estimatedTasksAllPartitions, err
 }
 
 // pickClientForRead mutates the given proto. Callers should copy the proto before if necessary.

client/matching/loadbalancer.go

@@ -1,6 +1,7 @@
 package matching
 
 import (
+	"math"
 	"math/rand"
 	"sync"
 
@@ -18,21 +19,26 @@ import (
 // when another partition has encoded-greater-than-zero backlog. 2x feels like a good ratio.
 var readPartitionWeightFloor = number.DecodeCompact8(1)
 
+// Keep a small root sample for estimating total write rate when backlog-aware routing would
+// otherwise make the root probability too small.
+const writePartitionRootProbabilityFloor = 0.01
+
 type (
 	// LoadBalancer is the interface for implementers of
 	// component that distributes add/poll api calls across
 	// available task queue partitions when possible
 	LoadBalancer interface {
-		// PickWritePartition returns the task queue partition for adding
-		// an activity or workflow task. The input is the name of the
+		// PickWritePartition returns the task queue partition for adding an
+		// activity or workflow task and the estimated number of tasks added
+		// across all partitions per root task. The input is the name of the
 		// original task queue (with no partition info). When forwardedFrom
 		// is non-empty, this call is forwardedFrom from a child partition
 		// to a parent partition in which case, no load balancing should be
 		// performed
 		PickWritePartition(
 			taskQueue *tqid.TaskQueue,
 			pc PartitionCounts,
-		) *tqid.NormalPartition
+		) (*tqid.NormalPartition, int)
 
 		// PickReadPartition returns the task queue partition to send a poller to.
 		// Input is name of the original task queue as specified by caller. When
@@ -86,14 +92,19 @@ func NewLoadBalancer(
 func (lb *defaultLoadBalancer) PickWritePartition(
 	taskQueue *tqid.TaskQueue,
 	pc PartitionCounts,
-) *tqid.NormalPartition {
+) (*tqid.NormalPartition, int) {
 	if n, ok := testhooks.Get(lb.testHooks, testhooks.MatchingLBForceWritePartition, namespace.ID(taskQueue.NamespaceId())); ok {
-		return taskQueue.NormalPartition(n)
+		partition := taskQueue.NormalPartition(n)
+		if partition.IsRoot() {
+			return partition, 1
+		}
+		// probability of reaching root is 0, so root can't know how many tasks were added overall
+		return partition, 0
 	}
 
 	nsName, err := lb.namespaceIDToName(namespace.ID(taskQueue.NamespaceId()))
 	if err != nil {
-		return taskQueue.RootPartition()
+		return taskQueue.RootPartition(), 1
 	}
 
 	var partitionCount int
@@ -103,22 +114,29 @@ func (lb *defaultLoadBalancer) PickWritePartition(
 		partitionCount = max(1, lb.nWritePartitions(nsName.String(), taskQueue.Name(), taskQueue.TaskType()))
 	}
 
-	return taskQueue.NormalPartition(pickWritePartitionByGap(
+	partitionID, estimatedTasksAllPartitions := pickWritePartitionByGap(
 		pc.BacklogCount,
 		partitionCount,
 		number.DecodeCompact8(pc.BacklogCap),
-	))
+	)
+	partition := taskQueue.NormalPartition(partitionID)
+
+	return partition, estimatedTasksAllPartitions
 }
 
 // pickWritePartitionByGap picks a partition with probability proportional to how far its backlog
 // is below backlogCap. Falls back to uniform random if any of these are true:
 //   - every partition is at or above the backlogCap
 //   - backlogCap is 0
 //   - when backlog data is not available for all write partitions
-func pickWritePartitionByGap(counts []number.Compact8, partitionCount int, backlogCap int64) int {
+func pickWritePartitionByGap(
+	counts []number.Compact8,
+	partitionCount int,
+	backlogCap int64,
+) (partitionID int, estimatedTasksAllPartitions int) {
 	if backlogCap == 0 ||
 		len(counts) < partitionCount {
-		return rand.Intn(partitionCount)
+		return rand.Intn(partitionCount), partitionCount
 	}
 
 	var total int64
@@ -128,20 +146,43 @@ func pickWritePartitionByGap(counts []number.Compact8, partitionCount int, backl
 		}
 	}
 	if total <= 0 { // all partitions are at or above cap
-		return rand.Intn(partitionCount)
+		return rand.Intn(partitionCount), partitionCount
 	}
-	r := rand.Int63n(total)
-	for i := range partitionCount {
-		gap := backlogCap - number.DecodeCompact8(counts[i])
-		if gap <= 0 { // this partition is at or above cap
-			continue
+
+	// if rootProbability < 0.01, choose the root with p=0.01
+	rootGap := max(int64(0), backlogCap-number.DecodeCompact8(counts[0]))
+	rootProbability := float64(rootGap) / float64(total)
+	if rootProbability < writePartitionRootProbabilityFloor {
+		if rand.Float64() < writePartitionRootProbabilityFloor {
+			return 0, int(math.Round(1 / writePartitionRootProbabilityFloor))
 		}
+		// choose between non-root partitions according to gap from backlog cap
+		partitionID := pickPartitionByGap(counts[1:partitionCount], backlogCap, total-rootGap) + 1
+		return partitionID, int(math.Round(1 / writePartitionRootProbabilityFloor))
+	}
+
+	return pickPartitionByGap(counts[:partitionCount], backlogCap, total), randomRound(1 / rootProbability)
+}
+
+// randomRound rounds without biasing the expected value.
+func randomRound(x float64) int {
+	n := math.Floor(x)
+	if rand.Float64() < x-n {
+		n++
+	}
+	return int(n)
+}
+
+func pickPartitionByGap(counts []number.Compact8, backlogCap int64, total int64) int {
+	r := rand.Int63n(total)
+	for i, count := range counts {
+		gap := max(int64(0), backlogCap-number.DecodeCompact8(count))
 		if r < gap { // more likely to be true the bigger this partition's gap is
 			return i
 		}
 		r -= gap
 	}
-	return partitionCount - 1 // unreachable in practice; guard against compact8 rounding
+	return len(counts) - 1
 }
 
 // PickReadPartition picks a partition for poller to poll task from, and keeps load balanced between partitions.

client/matching/loadbalancer_test.go

@@ -1,6 +1,7 @@
 package matching
 
 import (
+	"math"
 	"math/rand"
 	"sync"
 	"testing"
@@ -365,18 +366,24 @@ func TestPickWritePartition_BacklogAware(t *testing.T) {
 		BacklogCap:   200,
 		BacklogCount: []number.Compact8{0, number.EncodeCompact8(13_000_000)},
 	}
+	gap0 := backlogCap - number.DecodeCompact8(0)
+	gap1 := backlogCap - number.DecodeCompact8(number.EncodeCompact8(13_000_000))
 	counts := make([]int, 2)
+	estimatedTasks := 0
 	const n = 3000
 	for range n {
-		p := lb.PickWritePartition(taskQueue, pc)
+		p, estimatedTasksAllPartitions := lb.PickWritePartition(taskQueue, pc)
 		counts[p.PartitionId()]++
+		if p.IsRoot() {
+			estimatedTasks += estimatedTasksAllPartitions
+		}
 	}
 	require.Greater(t, counts[0], counts[1], "emptier partition should receive more writes")
 	require.Positive(t, counts[1], "the below-cap partition should still receive some writes")
-	gap0 := backlogCap - number.DecodeCompact8(0)
-	gap1 := backlogCap - number.DecodeCompact8(number.EncodeCompact8(13_000_000))
 	require.InDelta(t, float64(n)*float64(gap0)/float64(gap0+gap1), counts[0], float64(n)*0.05,
 		"writes split in proportion to each partition's gap to cap")
+	require.InDelta(t, n, estimatedTasks, float64(n)*0.05,
+		"root samples should estimate total writes without bias")
 
 	// Now, every partition at/above cap -> no gap to weight by, so the picker declines and the caller
 	// falls back to uniform random.
@@ -388,14 +395,44 @@ func TestPickWritePartition_BacklogAware(t *testing.T) {
 	}
 	atCap := make([]int, 2)
 	for range n {
-		p := lb.PickWritePartition(taskQueue, pcAtCap)
+		p, estimatedTasksAllPartitions := lb.PickWritePartition(taskQueue, pcAtCap)
 		atCap[p.PartitionId()]++
+		require.Equal(t, 2, estimatedTasksAllPartitions)
 	}
 	for i := range atCap {
 		require.InDelta(t, n/2, atCap[i], float64(n)*0.1, "at-cap partition %d roughly uniform", i)
 	}
 }
 
+func TestPickWritePartition_RootProbabilityFloor(t *testing.T) {
+	f, err := tqid.NewTaskQueueFamily("fake-namespace-id", "fake-taskqueue")
+	require.NoError(t, err)
+	taskQueue := f.TaskQueue(enumspb.TASK_QUEUE_TYPE_ACTIVITY)
+	lb := &defaultLoadBalancer{
+		namespaceIDToName: func(namespace.ID) (namespace.Name, error) { return "fake-namespace", nil },
+		taskQueueLBs:      make(map[tqid.TaskQueue]*tqLoadBalancer),
+	}
+
+	pc := PartitionCounts{
+		Read:         2,
+		Write:        2,
+		BacklogCap:   number.EncodeCompact8(1000),
+		BacklogCount: []number.Compact8{number.EncodeCompact8(1000), 0},
+	}
+	const attempts = 100_000
+	rootPicks := 0
+	expectedProbability := writePartitionRootProbabilityFloor
+	expectedTasksAllPartitions := int(math.Round(1 / expectedProbability))
+	for range attempts {
+		partition, estimatedTasksAllPartitions := lb.PickWritePartition(taskQueue, pc)
+		if partition.IsRoot() {
+			rootPicks++
+		}
+		require.Equal(t, expectedTasksAllPartitions, estimatedTasksAllPartitions)
+	}
+	require.InDelta(t, attempts*expectedProbability, rootPicks, attempts*expectedProbability*0.2)
+}
+
 func TestPickWritePartition_NoBacklogUniform(t *testing.T) {
 	f, err := tqid.NewTaskQueueFamily("fake-namespace-id", "fake-taskqueue")
 	require.NoError(t, err)
@@ -411,8 +448,9 @@ func TestPickWritePartition_NoBacklogUniform(t *testing.T) {
 	counts := make([]int, 4)
 	const n = 4000
 	for range n {
-		p := lb.PickWritePartition(taskQueue, pc)
+		p, estimatedTasksAllPartitions := lb.PickWritePartition(taskQueue, pc)
 		counts[p.PartitionId()]++
+		require.Equal(t, 4, estimatedTasksAllPartitions)
 	}
 	for i := range counts {
 		require.InDelta(t, n/4, counts[i], float64(n)*0.1, "partition %d roughly uniform", i)

client/matching/partition_counts.go

@@ -3,6 +3,7 @@ package matching
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"errors"
 	"slices"
 
@@ -20,6 +21,7 @@ import (
 // The "-bin" suffix instructs grpc to base64-encode the value, so we can use binary.
 const partitionCountsHeaderName = "pcnt-bin"
 const partitionCountsTrailerName = "pcnt-bin"
+const estimatedTasksAllPartitionsHeaderName = "etap-bin"
 
 // PartitionCounts is a smaller version of taskqueuespb.ClientPartitionCounts that we can more
 // easily pass around and put in a map.
@@ -94,6 +96,27 @@ func ParsePartitionCountsFromIncomingContext(ctx context.Context) (PartitionCoun
 	return parsePartitionCounts(vals[0])
 }
 
+func appendEstimatedTasksAllPartitions(ctx context.Context, estimatedTasksAllPartitions int, isRoot bool) context.Context {
+	if estimatedTasksAllPartitions <= 0 || !isRoot {
+		return ctx
+	}
+	b := make([]byte, 8)
+	binary.LittleEndian.PutUint64(b, uint64(estimatedTasksAllPartitions))
+	return metadata.AppendToOutgoingContext(ctx, estimatedTasksAllPartitionsHeaderName, string(b))
+}
+
+func ParseEstimatedTasksAllPartitions(ctx context.Context) int {
+	vals := metadata.ValueFromIncomingContext(ctx, estimatedTasksAllPartitionsHeaderName)
+	if len(vals) == 0 || len(vals[0]) != 8 {
+		return 0
+	}
+	estimatedTasksAllPartitions := binary.LittleEndian.Uint64([]byte(vals[0]))
+	if estimatedTasksAllPartitions == 0 || estimatedTasksAllPartitions > uint64(^uint(0)>>1) {
+		return 0
+	}
+	return int(estimatedTasksAllPartitions)
+}
+
 func parsePartitionCountsFromTrailer(trailer metadata.MD) (PartitionCounts, error) {
 	vals := trailer.Get(partitionCountsTrailerName)
 	if len(vals) == 0 {

client/matching/partition_counts_test.go

@@ -15,6 +15,24 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
+func TestEstimatedTasksAllPartitionsMetadata(t *testing.T) {
+	for _, estimatedTasksAllPartitions := range []int{1, 2, 100, 102} {
+		outgoingCtx := appendEstimatedTasksAllPartitions(context.Background(), estimatedTasksAllPartitions, true)
+		md, ok := metadata.FromOutgoingContext(outgoingCtx)
+		require.True(t, ok)
+		incomingCtx := metadata.NewIncomingContext(context.Background(), md)
+		require.Equal(t, estimatedTasksAllPartitions, ParseEstimatedTasksAllPartitions(incomingCtx))
+	}
+
+	for _, estimatedTasksAllPartitions := range []int{0, -1} {
+		outgoingCtx := appendEstimatedTasksAllPartitions(context.Background(), estimatedTasksAllPartitions, true)
+		md, ok := metadata.FromOutgoingContext(outgoingCtx)
+		require.False(t, ok)
+		incomingCtx := metadata.NewIncomingContext(context.Background(), md)
+		require.Zero(t, ParseEstimatedTasksAllPartitions(incomingCtx))
+	}
+}
+
 // setTrailerInOpts finds the grpc.TrailerCallOption in opts and populates it.
 func setTrailerInOpts(opts []grpc.CallOption, pc PartitionCounts) {
 	v, _ := pc.encode(true) // trailer path (server -> client) includes backlog info

service/matching/partition_scaler_interface.go

@@ -38,7 +38,7 @@ type PartitionScaler interface {
 }
 
 type PartitionScalerInput struct {
-	NumTasks      int // new tasks added since last call
+	NumTasks      int // estimated tasks added since last call
 	CurrentTarget int
 	BacklogCounts []byte
 	PrivateState  *anypb.Any