tasks/client/matching/partition_counts_test.go at rescue/uncommitted

mirror of https://github.com/hanzoai/tasks.git synced 2026-10-04 02:04:03 +00:00

Files

T

David ReissandGitHub dd3c0ac452 Dynamic partitioning: client side (#9732 )

## What changed?
Add client-side of dynamic partition scaling: `partitionCache`,
`PartitionCounts`, `StalePartitionCounts` error,
`invokeWithPartitionCounts`.

The intended behavior is:

For partition-aware calls (add+poll tasks), the client caches the
partition count for active task queues, and uses those partition counts
for load balancing (no change to load balancing algorithm yet). The
client sends its cached partition counts to the server in a grpc header,
and receives updated partition counts with the response in a grpc
trailer. If the updated counts are different, it updates its cache it
for subsequent requests. If the server indicates that the client's view
of partition count is stale, it returns a special `StalePartitionCounts`
error and then client makes one immediate retry with the newly received
counts.

With just this PR by itself, the server will never send counts, the
cache will always be empty, and the client will always fall back to
dynamic config for partition counts, so there's no change in behavior
yet.

## Why?
Implement half of dynamic partition scaling.

## How did you test it?
- [ ] built
- [ ] run locally and tested manually
- [ ] covered by existing tests
- [x] added new unit test(s)
- [x] added new functional test(s) – tests are in future PRs

2026-04-17 16:54:19 -07:00

266 lines

8.5 KiB

Go

Raw Permalink Blame History

 package matching
 import (
 	"context"
 	"errors"
 	"testing"
 	"github.com/stretchr/testify/require"
 	enumspb "go.temporal.io/api/enums/v1"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/metrics"
 	serviceerrors "go.temporal.io/server/common/serviceerror"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
 )
 // setTrailerInOpts finds the grpc.TrailerCallOption in opts and populates it.
 func setTrailerInOpts(opts []grpc.CallOption, pc PartitionCounts) {
 	v, _ := pc.encode()
 	md := metadata.Pairs(partitionCountsTrailerName, v)
 	for _, opt := range opts {
 		if t, ok := opt.(grpc.TrailerCallOption); ok {
 			*t.TrailerAddr = md
 			return
 		}
 	}
 }
 type hpcReq struct{}
 type hpcRes struct{ value string }
 func newTestCache(t *testing.T) *partitionCache {
 	cache := newPartitionCache(metrics.NoopMetricsHandler)
 	cache.Start()
 	t.Cleanup(cache.Stop)
 	return cache
 }
 func TestInvokeWithPartitionCounts_CacheMissSuccess(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	serverPC := PartitionCounts{Read: 4, Write: 4}
 	calls := 0
 	op := func(ctx context.Context, pc PartitionCounts, req *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		calls++
 		require.Equal(t, PartitionCounts{}, pc) // cache miss
 		setTrailerInOpts(opts, serverPC)
 		return &hpcRes{value: "ok"}, nil
 	}
 	res, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.NoError(t, err)
 	require.Equal(t, "ok", res.value)
 	require.Equal(t, 1, calls)
 	// cache should be updated
 	require.Equal(t, serverPC, cache.lookup(pkey))
 }
 func TestInvokeWithPartitionCounts_CacheHitSuccess(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	cachedPC := PartitionCounts{Read: 4, Write: 4}
 	cache.put(pkey, cachedPC)
 	op := func(ctx context.Context, pc PartitionCounts, req *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		require.Equal(t, cachedPC, pc)
 		// server confirms same counts
 		setTrailerInOpts(opts, cachedPC)
 		return &hpcRes{value: "ok"}, nil
 	}
 	res, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.NoError(t, err)
 	require.Equal(t, "ok", res.value)
 	require.Equal(t, cachedPC, cache.lookup(pkey))
 }
 func TestInvokeWithPartitionCounts_ServerUpdatesCount(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	cache.put(pkey, PartitionCounts{Read: 4, Write: 4})
 	newPC := PartitionCounts{Read: 8, Write: 8}
 	op := func(ctx context.Context, pc PartitionCounts, req *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		// server returns different counts
 		setTrailerInOpts(opts, newPC)
 		return &hpcRes{value: "ok"}, nil
 	}
 	_, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.NoError(t, err)
 	// cache should be updated
 	require.Equal(t, newPC, cache.lookup(pkey))
 }
 func TestInvokeWithPartitionCounts_StaleRetry_Succeeds(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	serverPC := PartitionCounts{Read: 8, Write: 8}
 	calls := 0
 	op := func(ctx context.Context, pc PartitionCounts, req *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		calls++
 		setTrailerInOpts(opts, serverPC)
 		if calls == 1 {
 			require.Equal(t, PartitionCounts{}, pc) // cache miss
 			return nil, serviceerrors.NewStalePartitionCounts("stale")
 		}
 		require.Equal(t, serverPC, pc) // retry with updated counts
 		return &hpcRes{value: "ok"}, nil
 	}
 	res, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.NoError(t, err)
 	require.Equal(t, "ok", res.value)
 	require.Equal(t, 2, calls)
 	require.Equal(t, serverPC, cache.lookup(pkey))
 }
 func TestInvokeWithPartitionCounts_StaleRetry_Fails(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	serverPC := PartitionCounts{Read: 8, Write: 8}
 	calls := 0
 	op := func(ctx context.Context, pc PartitionCounts, req *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		calls++
 		setTrailerInOpts(opts, serverPC)
 		if calls == 1 {
 			return nil, serviceerrors.NewStalePartitionCounts("stale first")
 		}
 		// second attempt: different non-stale error
 		return nil, errors.New("error")
 	}
 	_, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.Error(t, err)
 	require.Equal(t, 2, calls)
 }
 func TestInvokeWithPartitionCounts_OtherErrorNoRetry(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	serverPC := PartitionCounts{Read: 4, Write: 4}
 	calls := 0
 	op := func(ctx context.Context, pc PartitionCounts, req *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		calls++
 		// even on error, server sends trailer
 		setTrailerInOpts(opts, serverPC)
 		return nil, errors.New("error")
 	}
 	_, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.Error(t, err)
 	require.Equal(t, 1, calls) // no retry
 	// cache should still be updated from trailer
 	require.Equal(t, serverPC, cache.lookup(pkey))
 }
 func TestInvokeWithPartitionCounts_ZeroTrailerRemovesCache(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	cache.put(pkey, PartitionCounts{Read: 4, Write: 4})
 	op := func(ctx context.Context, pc PartitionCounts, req *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		// server signals "dynamic partitioning off"
 		setTrailerInOpts(opts, PartitionCounts{Read: 0, Write: 0})
 		return &hpcRes{value: "ok"}, nil
 	}
 	_, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.NoError(t, err)
 	// cache entry should be removed
 	require.Equal(t, PartitionCounts{}, cache.lookup(pkey))
 }
 func TestInvokeWithPartitionCounts_NoTrailerRemovesCache(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	originalPC := PartitionCounts{Read: 4, Write: 4}
 	cache.put(pkey, originalPC)
 	op := func(ctx context.Context, pc PartitionCounts, req *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		// no trailer set
 		return &hpcRes{value: "ok"}, nil
 	}
 	_, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.NoError(t, err)
 	// cache entry should be removed
 	require.Equal(t, PartitionCounts{}, cache.lookup(pkey))
 }
 func TestInvokeWithPartitionCounts_ParseErrorRemovesCache(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	originalPC := PartitionCounts{Read: 4, Write: 4}
 	cache.put(pkey, originalPC)
 	op := func(ctx context.Context, pc PartitionCounts, req *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		for _, opt := range opts {
 			if t, ok := opt.(grpc.TrailerCallOption); ok {
 				*t.TrailerAddr = metadata.Pairs(partitionCountsTrailerName, "this is an invalid proto message")
 			}
 		}
 		return &hpcRes{value: "ok"}, nil
 	}
 	_, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.NoError(t, err)
 	// cache entry should be removed
 	require.Equal(t, PartitionCounts{}, cache.lookup(pkey))
 }
 func TestInvokeWithPartitionCounts_OutgoingContextHasHeader(t *testing.T) {
 	t.Parallel()
 	cache := newTestCache(t)
 	pkey := cache.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
 	cachedPC := PartitionCounts{Read: 6, Write: 4}
 	cache.put(pkey, cachedPC)
 	op := func(ctx context.Context, _ PartitionCounts, _ *hpcReq, opts []grpc.CallOption) (*hpcRes, error) {
 		// verify the outgoing context has the partition counts header
 		md, ok := metadata.FromOutgoingContext(ctx)
 		require.True(t, ok)
 		vals := md.Get(partitionCountsHeaderName)
 		require.Len(t, vals, 1)
 		parsed, err := parsePartitionCounts(vals[0])
 		require.NoError(t, err)
 		require.Equal(t, cachedPC, parsed)
 		setTrailerInOpts(opts, cachedPC)
 		return &hpcRes{value: "ok"}, nil
 	}
 	_, err := invokeWithPartitionCounts(context.Background(), log.NewNoopLogger(), cache, pkey, &hpcReq{}, nil, op)
 	require.NoError(t, err)
 }

266 lines 8.5 KiB Go Raw Permalink Blame History

266 lines

8.5 KiB

Go

Raw Permalink Blame History