mirror of
https://github.com/hanzoai/tasks.git
synced 2026-10-04 02:04:03 +00:00
David ReissandGitHub
dd3c0ac452
## 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
98 lines
2.5 KiB
Go
98 lines
2.5 KiB
Go
package matching
|
|
|
|
import (
|
|
"sync"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
enumspb "go.temporal.io/api/enums/v1"
|
|
"go.temporal.io/server/common/metrics"
|
|
)
|
|
|
|
// Using a fixed UUID so tests are deterministic. Must be valid UUID format for makeKey.
|
|
const testNsID = "f47ac10b-58cc-4372-a567-0e02b2c3d479"
|
|
|
|
func TestPartitionCache_BasicPutLookup(t *testing.T) {
|
|
t.Parallel()
|
|
c := newPartitionCache(metrics.NoopMetricsHandler)
|
|
c.Start()
|
|
defer c.Stop()
|
|
|
|
key := c.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
|
|
keyb := c.makeKey(testNsID, "my-tq-b", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
|
|
|
|
pc4 := PartitionCounts{Read: 4, Write: 4}
|
|
pc8 := PartitionCounts{Read: 8, Write: 8}
|
|
|
|
c.put(key, pc4)
|
|
c.put(keyb, pc8)
|
|
require.Equal(t, pc4, c.lookup(key))
|
|
require.Equal(t, pc8, c.lookup(keyb))
|
|
|
|
c.put(key, pc8)
|
|
require.Equal(t, pc8, c.lookup(key))
|
|
|
|
// missing key
|
|
keyc := c.makeKey(testNsID, "nonexistent", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
|
|
require.Equal(t, PartitionCounts{}, c.lookup(keyc))
|
|
|
|
// invalid PartitionCounts removes
|
|
c.put(key, PartitionCounts{Read: -3, Write: -5})
|
|
require.Equal(t, PartitionCounts{}, c.lookup(key))
|
|
}
|
|
|
|
func TestPartitionCache_Rotate(t *testing.T) {
|
|
t.Parallel()
|
|
c := newPartitionCache(metrics.NoopMetricsHandler)
|
|
c.Start()
|
|
defer c.Stop()
|
|
|
|
key := c.makeKey(testNsID, "my-tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
|
|
pc := PartitionCounts{Read: 8, Write: 8}
|
|
c.put(key, pc)
|
|
|
|
// Rotate the shard — entry moves to prev
|
|
c.shards[c.shardFromKey(key)].rotate()
|
|
|
|
// Lookup should still find it (promotes from prev to active)
|
|
require.Equal(t, pc, c.lookup(key))
|
|
|
|
// After promotion, rotate again — it should still be in active
|
|
c.shards[c.shardFromKey(key)].rotate()
|
|
require.Equal(t, pc, c.lookup(key))
|
|
|
|
// Rotate twice without any lookup — entry should be gone
|
|
c.shards[c.shardFromKey(key)].rotate()
|
|
c.shards[c.shardFromKey(key)].rotate()
|
|
|
|
require.Equal(t, PartitionCounts{}, c.lookup(key))
|
|
|
|
// Put another value after rotate
|
|
pc4 := PartitionCounts{Read: 4, Write: 4}
|
|
c.put(key, pc4)
|
|
|
|
// Lookup should find it
|
|
require.Equal(t, pc4, c.lookup(key))
|
|
}
|
|
|
|
func TestPartitionCache_ConcurrentAccess(t *testing.T) {
|
|
t.Parallel()
|
|
c := newPartitionCache(metrics.NoopMetricsHandler)
|
|
c.Start()
|
|
defer c.Stop()
|
|
|
|
var wg sync.WaitGroup
|
|
for g := range 20 {
|
|
wg.Go(func() {
|
|
for range 1000 {
|
|
key := c.makeKey(testNsID, "tq", enumspb.TASK_QUEUE_TYPE_WORKFLOW)
|
|
pc := PartitionCounts{Read: int32(g + 1), Write: int32(g + 1)}
|
|
c.put(key, pc)
|
|
c.lookup(key)
|
|
}
|
|
})
|
|
}
|
|
wg.Wait()
|
|
// no panic or data races detected
|
|
}
|