Github/temporal
1
0
Fork 0
mirror of https://github.com/temporalio/temporal.git synced 2026-08-30 18:41:49 -07:00
Code Issues Packages Projects Releases Wiki Activity

Dynamic partitioning: server validation (#9997)

Browse Source 
## What changed?
- Matching server propagates partition counts in ephemeral data.
- Matching server validates partition counts sent by client, and rejects
if they are too far.

## Why?
Next part of dynamic partitioning.

## 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) - covered by future integration tests

## Potential risks
No change in behavior at this point, partition scale info is always nil,
so nothing will be rejected. The server will start sending grpc trailers
though.
This commit is contained in:
David Reiss
2026-05-22 11:05:13 -07:00
committed by GitHub
parent fb31b65b9f
commit 503c6fb3c3
10 changed files with 580 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
api/taskqueue/v1/message.go-helpers.pb.go
View File

@@ -412,6 +412,43 @@ func (this *VersionedEphemeralData) Equal(that interface{}) bool {
return proto.Equal(this, that1)
}
// Marshal an object of type PartitionScaleInfo to the protobuf v3 wire format
func (val *PartitionScaleInfo) Marshal() ([]byte, error) {
return proto.Marshal(val)
}
// Unmarshal an object of type PartitionScaleInfo from the protobuf v3 wire format
func (val *PartitionScaleInfo) Unmarshal(buf []byte) error {
return proto.Unmarshal(buf, val)
}
// Size returns the size of the object, in bytes, once serialized
func (val *PartitionScaleInfo) Size() int {
return proto.Size(val)
}
// Equal returns whether two PartitionScaleInfo values are equivalent by recursively
// comparing the message's fields.
// For more information see the documentation for
// https://pkg.go.dev/google.golang.org/protobuf/proto#Equal
func (this *PartitionScaleInfo) Equal(that interface{}) bool {
if that == nil {
return this == nil
}
var that1 *PartitionScaleInfo
switch t := that.(type) {
case *PartitionScaleInfo:
that1 = t
case PartitionScaleInfo:
that1 = &t
default:
return false
}
return proto.Equal(this, that1)
}
// Marshal an object of type ClientPartitionCounts to the protobuf v3 wire format
func (val *ClientPartitionCounts) Marshal() ([]byte, error) {
return proto.Marshal(val)

182
api/taskqueue/v1/message.pb.go
View File

@@ -795,8 +795,10 @@ func (x *TaskForwardInfo) GetDispatchVersionSet() string {
// task queue family (all queues with the same name, across types), ephemeral data applies only to
// one type at a time.
type EphemeralData struct {
state protoimpl.MessageState `protogen:"open.v1"`
Partition []*EphemeralData_ByPartition `protobuf:"bytes,1,rep,name=partition,proto3" json:"partition,omitempty"`
state protoimpl.MessageState `protogen:"open.v1"`
Partition []*EphemeralData_ByPartition `protobuf:"bytes,1,rep,name=partition,proto3" json:"partition,omitempty"`
// Current state of dynamic partition scaling
Scale *PartitionScaleInfo `protobuf:"bytes,2,opt,name=scale,proto3" json:"scale,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
@@ -838,6 +840,13 @@ func (x *EphemeralData) GetPartition() []*EphemeralData_ByPartition {
return nil
}
func (x *EphemeralData) GetScale() *PartitionScaleInfo {
if x != nil {
return x.Scale
}
return nil
}
type VersionedEphemeralData struct {
state protoimpl.MessageState `protogen:"open.v1"`
Data *EphemeralData `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"`
@@ -890,7 +899,74 @@ func (x *VersionedEphemeralData) GetVersion() int64 {
return 0
}
// PartitionScaleInfo is propagated among task queue partitions in ephemeral data.
type PartitionScaleInfo struct {
state protoimpl.MessageState `protogen:"open.v1"`
Read int32 `protobuf:"varint,1,opt,name=read,proto3" json:"read,omitempty"`
Write int32 `protobuf:"varint,2,opt,name=write,proto3" json:"write,omitempty"`
// version identifies a specific version of the scale state, which changes when the target
// number of partitions (i.e. the write count) changes. It may not change for other changes
// to scale state/info. This is used by the scale manager to know that a partition is
// operating with the latest scale info, to avoid ABA problems, i.e. to differentiate from a
// previous version whose read and write counts happen to be the same numbers.
Version int64 `protobuf:"fixed64,10,opt,name=version,proto3" json:"version,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *PartitionScaleInfo) Reset() {
*x = PartitionScaleInfo{}
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[11]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *PartitionScaleInfo) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*PartitionScaleInfo) ProtoMessage() {}
func (x *PartitionScaleInfo) ProtoReflect() protoreflect.Message {
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[11]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use PartitionScaleInfo.ProtoReflect.Descriptor instead.
func (*PartitionScaleInfo) Descriptor() ([]byte, []int) {
return file_temporal_server_api_taskqueue_v1_message_proto_rawDescGZIP(), []int{11}
}
func (x *PartitionScaleInfo) GetRead() int32 {
if x != nil {
return x.Read
}
return 0
}
func (x *PartitionScaleInfo) GetWrite() int32 {
if x != nil {
return x.Write
}
return 0
}
func (x *PartitionScaleInfo) GetVersion() int64 {
if x != nil {
return x.Version
}
return 0
}
// ClientPartitionCounts is propagated from the matching service to clients in grpc headers/trailers.
// It may be a subset of PartitionScaleInfo.
type ClientPartitionCounts struct {
state protoimpl.MessageState `protogen:"open.v1"`
Read int32 `protobuf:"varint,1,opt,name=read,proto3" json:"read,omitempty"`
@@ -901,7 +977,7 @@ type ClientPartitionCounts struct {
func (x *ClientPartitionCounts) Reset() {
*x = ClientPartitionCounts{}
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[11]
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[12]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
@@ -913,7 +989,7 @@ func (x *ClientPartitionCounts) String() string {
func (*ClientPartitionCounts) ProtoMessage() {}
func (x *ClientPartitionCounts) ProtoReflect() protoreflect.Message {
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[11]
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[12]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
@@ -926,7 +1002,7 @@ func (x *ClientPartitionCounts) ProtoReflect() protoreflect.Message {
// Deprecated: Use ClientPartitionCounts.ProtoReflect.Descriptor instead.
func (*ClientPartitionCounts) Descriptor() ([]byte, []int) {
return file_temporal_server_api_taskqueue_v1_message_proto_rawDescGZIP(), []int{11}
return file_temporal_server_api_taskqueue_v1_message_proto_rawDescGZIP(), []int{12}
}
func (x *ClientPartitionCounts) GetRead() int32 {
@@ -957,7 +1033,7 @@ type EphemeralData_ByVersion struct {
func (x *EphemeralData_ByVersion) Reset() {
*x = EphemeralData_ByVersion{}
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[13]
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[14]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
@@ -969,7 +1045,7 @@ func (x *EphemeralData_ByVersion) String() string {
func (*EphemeralData_ByVersion) ProtoMessage() {}
func (x *EphemeralData_ByVersion) ProtoReflect() protoreflect.Message {
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[13]
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[14]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
@@ -1009,7 +1085,7 @@ type EphemeralData_ByPartition struct {
func (x *EphemeralData_ByPartition) Reset() {
*x = EphemeralData_ByPartition{}
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[14]
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[15]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
@@ -1021,7 +1097,7 @@ func (x *EphemeralData_ByPartition) String() string {
func (*EphemeralData_ByPartition) ProtoMessage() {}
func (x *EphemeralData_ByPartition) ProtoReflect() protoreflect.Message {
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[14]
mi := &file_temporal_server_api_taskqueue_v1_message_proto_msgTypes[15]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
@@ -1116,9 +1192,10 @@ const file_temporal_server_api_taskqueue_v1_message_proto_rawDesc = "" +
"createTime\x12Z\n" +
"\rredirect_info\x18\x03 \x01(\v25.temporal.server.api.taskqueue.v1.BuildIdRedirectInfoR\fredirectInfo\x12*\n" +
"\x11dispatch_build_id\x18\x04 \x01(\tR\x0fdispatchBuildId\x120\n" +
"\x14dispatch_version_set\x18\x05 \x01(\tR\x12dispatchVersionSet\"\x89\x03\n" +
"\x14dispatch_version_set\x18\x05 \x01(\tR\x12dispatchVersionSet\"\xd5\x03\n" +
"\rEphemeralData\x12Y\n" +
"\tpartition\x18\x01 \x03(\v2;.temporal.server.api.taskqueue.v1.EphemeralData.ByPartitionR\tpartition\x1a\x99\x01\n" +
"\tpartition\x18\x01 \x03(\v2;.temporal.server.api.taskqueue.v1.EphemeralData.ByPartitionR\tpartition\x12J\n" +
"\x05scale\x18\x02 \x01(\v24.temporal.server.api.taskqueue.v1.PartitionScaleInfoR\x05scale\x1a\x99\x01\n" +
"\tByVersion\x12T\n" +
"\aversion\x18\x01 \x01(\v2:.temporal.server.api.deployment.v1.WorkerDeploymentVersionR\aversion\x126\n" +
"\x17backlog_priority_levels\x18\x02 \x01(\x03R\x15backlogPriorityLevels\x1a\x80\x01\n" +
@@ -1127,7 +1204,12 @@ const file_temporal_server_api_taskqueue_v1_message_proto_rawDesc = "" +
"\aversion\x18\x02 \x03(\v29.temporal.server.api.taskqueue.v1.EphemeralData.ByVersionR\aversion\"w\n" +
"\x16VersionedEphemeralData\x12C\n" +
"\x04data\x18\x01 \x01(\v2/.temporal.server.api.taskqueue.v1.EphemeralDataR\x04data\x12\x18\n" +
"\aversion\x18\x02 \x01(\x03R\aversion\"A\n" +
"\aversion\x18\x02 \x01(\x03R\aversion\"X\n" +
"\x12PartitionScaleInfo\x12\x12\n" +
"\x04read\x18\x01 \x01(\x05R\x04read\x12\x14\n" +
"\x05write\x18\x02 \x01(\x05R\x05write\x12\x18\n" +
"\aversion\x18\n" +
" \x01(\x10R\aversion\"A\n" +
"\x15ClientPartitionCounts\x12\x12\n" +
"\x04read\x18\x01 \x01(\x05R\x04read\x12\x14\n" +
"\x05write\x18\x02 \x01(\x05R\x05writeB2Z0go.temporal.io/server/api/taskqueue/v1;taskqueueb\x06proto3"
@@ -1144,7 +1226,7 @@ func file_temporal_server_api_taskqueue_v1_message_proto_rawDescGZIP() []byte {
return file_temporal_server_api_taskqueue_v1_message_proto_rawDescData
}
var file_temporal_server_api_taskqueue_v1_message_proto_msgTypes = make([]protoimpl.MessageInfo, 15)
var file_temporal_server_api_taskqueue_v1_message_proto_msgTypes = make([]protoimpl.MessageInfo, 16)
var file_temporal_server_api_taskqueue_v1_message_proto_goTypes = []any{
(*TaskVersionDirective)(nil), // 0: temporal.server.api.taskqueue.v1.TaskVersionDirective
(*FairLevel)(nil), // 1: temporal.server.api.taskqueue.v1.FairLevel
@@ -1157,50 +1239,52 @@ var file_temporal_server_api_taskqueue_v1_message_proto_goTypes = []any{
(*TaskForwardInfo)(nil), // 8: temporal.server.api.taskqueue.v1.TaskForwardInfo
(*EphemeralData)(nil), // 9: temporal.server.api.taskqueue.v1.EphemeralData
(*VersionedEphemeralData)(nil), // 10: temporal.server.api.taskqueue.v1.VersionedEphemeralData
(*ClientPartitionCounts)(nil), // 11: temporal.server.api.taskqueue.v1.ClientPartitionCounts
nil, // 12: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.TaskQueueStatsByPriorityKeyEntry
(*EphemeralData_ByVersion)(nil), // 13: temporal.server.api.taskqueue.v1.EphemeralData.ByVersion
(*EphemeralData_ByPartition)(nil), // 14: temporal.server.api.taskqueue.v1.EphemeralData.ByPartition
(*emptypb.Empty)(nil), // 15: google.protobuf.Empty
(v1.VersioningBehavior)(0), // 16: temporal.api.enums.v1.VersioningBehavior
(*v11.Deployment)(nil), // 17: temporal.api.deployment.v1.Deployment
(*v12.WorkerDeploymentVersion)(nil), // 18: temporal.server.api.deployment.v1.WorkerDeploymentVersion
(*v13.TaskIdBlock)(nil), // 19: temporal.api.taskqueue.v1.TaskIdBlock
(*v13.PollerInfo)(nil), // 20: temporal.api.taskqueue.v1.PollerInfo
(*v13.TaskQueueStats)(nil), // 21: temporal.api.taskqueue.v1.TaskQueueStats
(v1.TaskQueueType)(0), // 22: temporal.api.enums.v1.TaskQueueType
(v14.TaskSource)(0), // 23: temporal.server.api.enums.v1.TaskSource
(*timestamppb.Timestamp)(nil), // 24: google.protobuf.Timestamp
(*PartitionScaleInfo)(nil), // 11: temporal.server.api.taskqueue.v1.PartitionScaleInfo
(*ClientPartitionCounts)(nil), // 12: temporal.server.api.taskqueue.v1.ClientPartitionCounts
nil, // 13: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.TaskQueueStatsByPriorityKeyEntry
(*EphemeralData_ByVersion)(nil), // 14: temporal.server.api.taskqueue.v1.EphemeralData.ByVersion
(*EphemeralData_ByPartition)(nil), // 15: temporal.server.api.taskqueue.v1.EphemeralData.ByPartition
(*emptypb.Empty)(nil), // 16: google.protobuf.Empty
(v1.VersioningBehavior)(0), // 17: temporal.api.enums.v1.VersioningBehavior
(*v11.Deployment)(nil), // 18: temporal.api.deployment.v1.Deployment
(*v12.WorkerDeploymentVersion)(nil), // 19: temporal.server.api.deployment.v1.WorkerDeploymentVersion
(*v13.TaskIdBlock)(nil), // 20: temporal.api.taskqueue.v1.TaskIdBlock
(*v13.PollerInfo)(nil), // 21: temporal.api.taskqueue.v1.PollerInfo
(*v13.TaskQueueStats)(nil), // 22: temporal.api.taskqueue.v1.TaskQueueStats
(v1.TaskQueueType)(0), // 23: temporal.api.enums.v1.TaskQueueType
(v14.TaskSource)(0), // 24: temporal.server.api.enums.v1.TaskSource
(*timestamppb.Timestamp)(nil), // 25: google.protobuf.Timestamp
}
var file_temporal_server_api_taskqueue_v1_message_proto_depIdxs = []int32{
15, // 0: temporal.server.api.taskqueue.v1.TaskVersionDirective.use_assignment_rules:type_name -> google.protobuf.Empty
16, // 1: temporal.server.api.taskqueue.v1.TaskVersionDirective.behavior:type_name -> temporal.api.enums.v1.VersioningBehavior
17, // 2: temporal.server.api.taskqueue.v1.TaskVersionDirective.deployment:type_name -> temporal.api.deployment.v1.Deployment
18, // 3: temporal.server.api.taskqueue.v1.TaskVersionDirective.deployment_version:type_name -> temporal.server.api.deployment.v1.WorkerDeploymentVersion
16, // 0: temporal.server.api.taskqueue.v1.TaskVersionDirective.use_assignment_rules:type_name -> google.protobuf.Empty
17, // 1: temporal.server.api.taskqueue.v1.TaskVersionDirective.behavior:type_name -> temporal.api.enums.v1.VersioningBehavior
18, // 2: temporal.server.api.taskqueue.v1.TaskVersionDirective.deployment:type_name -> temporal.api.deployment.v1.Deployment
19, // 3: temporal.server.api.taskqueue.v1.TaskVersionDirective.deployment_version:type_name -> temporal.server.api.deployment.v1.WorkerDeploymentVersion
1, // 4: temporal.server.api.taskqueue.v1.InternalTaskQueueStatus.fair_read_level:type_name -> temporal.server.api.taskqueue.v1.FairLevel
1, // 5: temporal.server.api.taskqueue.v1.InternalTaskQueueStatus.fair_ack_level:type_name -> temporal.server.api.taskqueue.v1.FairLevel
19, // 6: temporal.server.api.taskqueue.v1.InternalTaskQueueStatus.task_id_block:type_name -> temporal.api.taskqueue.v1.TaskIdBlock
20, // 6: temporal.server.api.taskqueue.v1.InternalTaskQueueStatus.task_id_block:type_name -> temporal.api.taskqueue.v1.TaskIdBlock
1, // 7: temporal.server.api.taskqueue.v1.InternalTaskQueueStatus.fair_max_read_level:type_name -> temporal.server.api.taskqueue.v1.FairLevel
4, // 8: temporal.server.api.taskqueue.v1.TaskQueueVersionInfoInternal.physical_task_queue_info:type_name -> temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo
20, // 9: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.pollers:type_name -> temporal.api.taskqueue.v1.PollerInfo
21, // 9: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.pollers:type_name -> temporal.api.taskqueue.v1.PollerInfo
2, // 10: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.internal_task_queue_status:type_name -> temporal.server.api.taskqueue.v1.InternalTaskQueueStatus
21, // 11: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.task_queue_stats:type_name -> temporal.api.taskqueue.v1.TaskQueueStats
12, // 12: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.task_queue_stats_by_priority_key:type_name -> temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.TaskQueueStatsByPriorityKeyEntry
22, // 13: temporal.server.api.taskqueue.v1.TaskQueuePartition.task_queue_type:type_name -> temporal.api.enums.v1.TaskQueueType
22, // 11: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.task_queue_stats:type_name -> temporal.api.taskqueue.v1.TaskQueueStats
13, // 12: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.task_queue_stats_by_priority_key:type_name -> temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.TaskQueueStatsByPriorityKeyEntry
23, // 13: temporal.server.api.taskqueue.v1.TaskQueuePartition.task_queue_type:type_name -> temporal.api.enums.v1.TaskQueueType
6, // 14: temporal.server.api.taskqueue.v1.TaskQueuePartition.worker_commands:type_name -> temporal.server.api.taskqueue.v1.WorkerCommandsPartitionId
23, // 15: temporal.server.api.taskqueue.v1.TaskForwardInfo.task_source:type_name -> temporal.server.api.enums.v1.TaskSource
24, // 16: temporal.server.api.taskqueue.v1.TaskForwardInfo.create_time:type_name -> google.protobuf.Timestamp
24, // 15: temporal.server.api.taskqueue.v1.TaskForwardInfo.task_source:type_name -> temporal.server.api.enums.v1.TaskSource
25, // 16: temporal.server.api.taskqueue.v1.TaskForwardInfo.create_time:type_name -> google.protobuf.Timestamp
7, // 17: temporal.server.api.taskqueue.v1.TaskForwardInfo.redirect_info:type_name -> temporal.server.api.taskqueue.v1.BuildIdRedirectInfo
14, // 18: temporal.server.api.taskqueue.v1.EphemeralData.partition:type_name -> temporal.server.api.taskqueue.v1.EphemeralData.ByPartition
9, // 19: temporal.server.api.taskqueue.v1.VersionedEphemeralData.data:type_name -> temporal.server.api.taskqueue.v1.EphemeralData
21, // 20: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.TaskQueueStatsByPriorityKeyEntry.value:type_name -> temporal.api.taskqueue.v1.TaskQueueStats
18, // 21: temporal.server.api.taskqueue.v1.EphemeralData.ByVersion.version:type_name -> temporal.server.api.deployment.v1.WorkerDeploymentVersion
13, // 22: temporal.server.api.taskqueue.v1.EphemeralData.ByPartition.version:type_name -> temporal.server.api.taskqueue.v1.EphemeralData.ByVersion
23, // [23:23] is the sub-list for method output_type
23, // [23:23] is the sub-list for method input_type
23, // [23:23] is the sub-list for extension type_name
23, // [23:23] is the sub-list for extension extendee
0, // [0:23] is the sub-list for field type_name
15, // 18: temporal.server.api.taskqueue.v1.EphemeralData.partition:type_name -> temporal.server.api.taskqueue.v1.EphemeralData.ByPartition
11, // 19: temporal.server.api.taskqueue.v1.EphemeralData.scale:type_name -> temporal.server.api.taskqueue.v1.PartitionScaleInfo
9, // 20: temporal.server.api.taskqueue.v1.VersionedEphemeralData.data:type_name -> temporal.server.api.taskqueue.v1.EphemeralData
22, // 21: temporal.server.api.taskqueue.v1.PhysicalTaskQueueInfo.TaskQueueStatsByPriorityKeyEntry.value:type_name -> temporal.api.taskqueue.v1.TaskQueueStats
19, // 22: temporal.server.api.taskqueue.v1.EphemeralData.ByVersion.version:type_name -> temporal.server.api.deployment.v1.WorkerDeploymentVersion
14, // 23: temporal.server.api.taskqueue.v1.EphemeralData.ByPartition.version:type_name -> temporal.server.api.taskqueue.v1.EphemeralData.ByVersion
24, // [24:24] is the sub-list for method output_type
24, // [24:24] is the sub-list for method input_type
24, // [24:24] is the sub-list for extension type_name
24, // [24:24] is the sub-list for extension extendee
0, // [0:24] is the sub-list for field type_name
}
func init() { file_temporal_server_api_taskqueue_v1_message_proto_init() }
@@ -1223,7 +1307,7 @@ func file_temporal_server_api_taskqueue_v1_message_proto_init() {
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
RawDescriptor: unsafe.Slice(unsafe.StringData(file_temporal_server_api_taskqueue_v1_message_proto_rawDesc), len(file_temporal_server_api_taskqueue_v1_message_proto_rawDesc)),
NumEnums: 0,
NumMessages: 15,
NumMessages: 16,
NumExtensions: 0,
NumServices: 0,
},

8
common/dynamicconfig/constants.go
View File

@@ -1518,6 +1518,14 @@ default as namespace cardinality can be high and this requires a metrics collect
false,
`MatchingAutoEnableV2 automatically enables fairness when a fairness or priority key is seen`,
)
MatchingPartitionScaleAllowedDrift = NewTaskQueueTypedSetting(
"matching.partitionScaleAllowedDrift",
PartitionScaleAllowedDrift{
Delta: 1,
Ratio: 1.5,
},
`How far off client partition scale values have to be to reject RPCs.`,
)
// Worker registry settings
MatchingWorkerRegistryNumBuckets = NewGlobalIntSetting(

11
common/dynamicconfig/shared_constants.go
View File

@@ -117,3 +117,14 @@ var DefaultHistoryCacheBackgroundEvictSettings = CacheBackgroundEvictSettings{
LoopInterval: 1 * time.Minute,
MaxEntryPerCall: 1024,
}
type PartitionScaleAllowedDrift struct {
// Delta and Ratio controls how far off client counts can be before we reject an RPC.
// If the client count is within the delta, it's allowed. Also, if the ratio of
// client's count / current count is within [1/ratio, ratio], then it's allowed.
// To always allow: set Delta to a very high number and Ratio to 1.0.
// To never allow except on exact match: set Delta to 0 and Ratio to 1.0.
// Allowing more means fewer retries, allowing less means more accurate load balancing.
Delta int32
Ratio float32
}

17
proto/internal/temporal/server/api/taskqueue/v1/message.proto
View File

@@ -163,6 +163,9 @@ message EphemeralData {
repeated ByVersion version = 2;
}
repeated ByPartition partition = 1;
// Current state of dynamic partition scaling
PartitionScaleInfo scale = 2;
}
message VersionedEphemeralData {
@@ -170,7 +173,21 @@ message VersionedEphemeralData {
int64 version = 2;
}
// PartitionScaleInfo is propagated among task queue partitions in ephemeral data.
message PartitionScaleInfo {
int32 read = 1;
int32 write = 2;
// version identifies a specific version of the scale state, which changes when the target
// number of partitions (i.e. the write count) changes. It may not change for other changes
// to scale state/info. This is used by the scale manager to know that a partition is
// operating with the latest scale info, to avoid ABA problems, i.e. to differentiate from a
// previous version whose read and write counts happen to be the same numbers.
sfixed64 version = 10;
}
// ClientPartitionCounts is propagated from the matching service to clients in grpc headers/trailers.
// It may be a subset of PartitionScaleInfo.
message ClientPartitionCounts {
int32 read = 1;
int32 write = 2;

12
service/matching/config.go
View File

@@ -134,7 +134,8 @@ type (
PollerScalingWaitTime dynamicconfig.DurationPropertyFnWithTaskQueueFilter
PollerScalingDecisionsPerSecond dynamicconfig.FloatPropertyFnWithTaskQueueFilter
FairnessCounter dynamicconfig.TypedPropertyFnWithTaskQueueFilter[counter.CounterParams]
FairnessCounter dynamicconfig.TypedPropertyFnWithTaskQueueFilter[counter.CounterParams]
PartitionScaleAllowedDrift dynamicconfig.TypedPropertyFnWithTaskQueueFilter[dynamicconfig.PartitionScaleAllowedDrift]
LogAllReqErrors dynamicconfig.BoolPropertyFnWithNamespaceFilter
}
@@ -222,7 +223,8 @@ type (
PollerScalingWaitTime func() time.Duration
PollerScalingDecisionsPerSecond func() float64
FairnessCounter func() counter.CounterParams
FairnessCounter func() counter.CounterParams
PartitionScaleAllowedDrift func() dynamicconfig.PartitionScaleAllowedDrift
loadCause loadCause
}
@@ -369,7 +371,8 @@ func NewConfig(
PollerScalingWaitTime: dynamicconfig.MatchingPollerScalingWaitTime.Get(dc),
PollerScalingDecisionsPerSecond: dynamicconfig.MatchingPollerScalingDecisionsPerSecond.Get(dc),
FairnessCounter: dynamicconfig.MatchingFairnessCounter.Get(dc),
FairnessCounter: dynamicconfig.MatchingFairnessCounter.Get(dc),
PartitionScaleAllowedDrift: dynamicconfig.MatchingPartitionScaleAllowedDrift.Get(dc),
LogAllReqErrors: dynamicconfig.LogAllReqErrors.Get(dc),
}
@@ -538,6 +541,9 @@ func newTaskQueueConfig(tq *tqid.TaskQueue, config *Config, ns namespace.Name) *
FairnessCounter: func() counter.CounterParams {
return config.FairnessCounter(ns.String(), taskQueueName, taskType)
},
PartitionScaleAllowedDrift: func() dynamicconfig.PartitionScaleAllowedDrift {
return config.PartitionScaleAllowedDrift(ns.String(), taskQueueName, taskType)
},
MaxVersionsInTaskQueue: func() int { return config.MaxVersionsInTaskQueue(ns.String()) },
}
}

129
service/matching/task_queue_partition_manager.go
View File

@@ -23,6 +23,7 @@ import (
"go.temporal.io/server/api/matchingservice/v1"
persistencespb "go.temporal.io/server/api/persistence/v1"
taskqueuespb "go.temporal.io/server/api/taskqueue/v1"
"go.temporal.io/server/client/matching"
"go.temporal.io/server/common/backoff"
"go.temporal.io/server/common/cache"
"go.temporal.io/server/common/dynamicconfig"
@@ -45,7 +46,12 @@ import (
"google.golang.org/protobuf/types/known/timestamppb"
)
var errDefaultQueueNotInit = serviceerror.NewInternal("defaultQueue is not initializaed")
var (
errDefaultQueueNotInit = serviceerror.NewInternal("defaultQueue is not initializaed")
errPartitionInvalid = serviceerrors.NewStalePartitionCounts("partition is invalid")
errPartitionDraining = serviceerrors.NewStalePartitionCounts("partition is draining")
errPartitionCountsStale = serviceerrors.NewStalePartitionCounts("counts are too far off")
)
const (
defaultTaskDispatchRPS = 100000.0
@@ -126,6 +132,7 @@ func newTaskQueuePartitionManager(
userDataManager,
tqConfig,
partition.TaskQueue().TaskType())
var taskHooks []hooks.TaskHook
for _, hookFactory := range e.taskHookFactories {
taskHook := hookFactory.Create(&hooks.TaskHookFactoryCreateDetails{
@@ -306,6 +313,99 @@ func (pm *taskQueuePartitionManagerImpl) Stop(unloadCause unloadCause) {
pm.goroGroup.Cancel()
}
func (pm *taskQueuePartitionManagerImpl) checkPartitionCounts(ctx context.Context, forWrite, forwarded bool) error {
normal, ok := pm.partition.(*tqid.NormalPartition)
if !ok {
return nil // only normal partitions do dynamic scaling
}
id := normal.PartitionId()
// userDataManager must be initialized here already so we can just ask it for scale info
scaleInfo := pm.userDataManager.PartitionScale()
if scaleInfo.GetRead() <= 0 || scaleInfo.GetWrite() <= 0 || scaleInfo.Write > scaleInfo.Read {
return nil // missing or invalid scale info
}
// always validate partition id based on read/write counts and scale info
if err := validatePartitionCounts(id, scaleInfo, forWrite); err != nil {
return err
}
// if client has sent its idea of counts, also validate drift
clientPC, err := matching.ParsePartitionCountsFromIncomingContext(ctx)
if err != nil {
pm.throttledLogger.Info("partition count header parse error", tag.Error(err))
return nil // just log and skip the check
} else if !clientPC.Valid() {
return nil // client didn't send anything or invalid, skip check
}
allowedDrift := pm.config.PartitionScaleAllowedDrift()
return validatePartitionScaleDrift(scaleInfo, forWrite, clientPC, allowedDrift)
}
// validatePartitionCounts checks whether a partition should accept an RPC based on the current
// scale info. It returns nil if the RPC should be accepted, or an error if it should be
// rejected. scaleInfo must be valid (positive values).
func validatePartitionCounts(
partitionID int,
scaleInfo *taskqueuespb.PartitionScaleInfo,
forWrite bool,
) error {
switch {
case partitionID < 0:
return serviceerror.NewInternal("negative partition id")
case partitionID >= int(scaleInfo.Read):
return errPartitionInvalid
case partitionID >= int(scaleInfo.Write) && forWrite:
return errPartitionDraining
default:
return nil
}
}
// validatePartitionScaleDrift checks whether a partition should accept an RPC based on
// the client's idea of partition counts. It returns nil if the RPC should be accepted, or an
// error if it should be rejected. scaleInfo and clientPC must both be valid.
func validatePartitionScaleDrift(
scaleInfo *taskqueuespb.PartitionScaleInfo,
forWrite bool,
clientPC matching.PartitionCounts,
allowedDrift dynamicconfig.PartitionScaleAllowedDrift,
) error {
var delta int32
var ratio float32
if forWrite {
delta = clientPC.Write - scaleInfo.Write
ratio = float32(clientPC.Write) / float32(scaleInfo.Write)
} else {
delta = clientPC.Read - scaleInfo.Read
ratio = float32(clientPC.Read) / float32(scaleInfo.Read)
}
effectiveRatio := max(1.001, allowedDrift.Ratio)
if delta >= -allowedDrift.Delta && delta <= allowedDrift.Delta ||
ratio >= 1/effectiveRatio && ratio <= effectiveRatio {
return nil
}
// otherwise reject to improve load balancing
return errPartitionCountsStale
}
func (pm *taskQueuePartitionManagerImpl) sendPartitionCountTrailer(ctx context.Context) {
// note this sends the trailer even if there is no scale info (i.e. dynamic partition
// scaling is not enabled). that will instruct clients to fall back to dynamic config.
scaleInfo := pm.userDataManager.PartitionScale()
err := matching.PartitionCounts{
Read: scaleInfo.GetRead(),
Write: scaleInfo.GetWrite(),
}.SetTrailer(ctx)
if err != nil {
// TODO: this is very noisy in unit tests, figure out how to log it only in non-test
pm.throttledLogger.Debug("error setting partition count trailer", tag.Error(err))
}
}
func (pm *taskQueuePartitionManagerImpl) GetRateLimitManager() *rateLimitManager {
return pm.rateLimitManager
}
@@ -394,6 +494,11 @@ func (pm *taskQueuePartitionManagerImpl) AddTask(
ctx context.Context,
params addTaskParams,
) (buildId string, syncMatched bool, err error) {
defer pm.sendPartitionCountTrailer(ctx)
if err := pm.checkPartitionCounts(ctx, true, params.forwardInfo != nil); err != nil {
return "", false, err
}
var spoolQueue, syncMatchQueue physicalTaskQueueManager
directive := params.taskInfo.GetVersionDirective()
@@ -550,6 +655,11 @@ func (pm *taskQueuePartitionManagerImpl) PollTask(
ctx context.Context,
pollMetadata *pollMetadata,
) (*internalTask, bool, error) {
defer pm.sendPartitionCountTrailer(ctx)
if err := pm.checkPartitionCounts(ctx, false, pollMetadata.forwardedFrom != ""); err != nil {
return nil, false, err
}
var err error
dbq := pm.defaultQueue()
if dbq == nil {
@@ -825,6 +935,12 @@ func (pm *taskQueuePartitionManagerImpl) DispatchQueryTask(
taskID string,
request *matchingservice.QueryWorkflowRequest,
) (*matchingservice.QueryWorkflowResponse, error) {
// query counts as "write" for partition load balancing
defer pm.sendPartitionCountTrailer(ctx)
if err := pm.checkPartitionCounts(ctx, true, request.ForwardInfo != nil); err != nil {
return nil, err
}
task := newInternalQueryTask(taskID, request)
pm.config.setDefaultPriority(task)
@@ -867,6 +983,12 @@ func (pm *taskQueuePartitionManagerImpl) DispatchNexusTask(
taskId string,
request *matchingservice.DispatchNexusTaskRequest,
) (*matchingservice.DispatchNexusTaskResponse, error) {
// nexus counts as "write" for partition load balancing
defer pm.sendPartitionCountTrailer(ctx)
if err := pm.checkPartitionCounts(ctx, true, request.ForwardInfo != nil); err != nil {
return nil, err
}
deadline, _ := ctx.Deadline() // If not set by user, our client will set a default.
var opDeadline time.Time
if header := nexus.Header(request.GetRequest().GetHeader()); header != nil {
@@ -1658,7 +1780,10 @@ func (pm *taskQueuePartitionManagerImpl) ForceLoadAllChildPartitions() {
return
}
partitions := int32(pm.config.NumReadPartitions())
partitions := pm.userDataManager.PartitionScale().GetRead()
if partitions == 0 {
partitions = int32(pm.config.NumReadPartitions())
}
if partitions <= 1 {
return
}

52
service/matching/task_queue_partition_manager_test.go
View File

@@ -25,6 +25,7 @@ import (
"go.temporal.io/server/common/metrics"
"go.temporal.io/server/common/metrics/metricstest"
"go.temporal.io/server/common/namespace"
"go.temporal.io/server/common/testing/await"
"go.temporal.io/server/common/testing/protorequire"
"go.temporal.io/server/common/testing/testlogger"
"go.temporal.io/server/common/tqid"
@@ -1117,15 +1118,17 @@ func (s *PartitionManagerTestSuite) TestHasAnyPollerAfter() {
// one unversioned poller
s.pollWithIdentity("uv", "", false, false)
s.True(s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-100 * time.Microsecond)))
time.Sleep(time.Millisecond)
s.False(s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-100 * time.Microsecond)))
s.True(s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-10 * time.Millisecond)))
await.RequireTrue(s.T(), func() bool {
return !s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-10 * time.Millisecond))
}, 20*time.Millisecond, time.Millisecond)
// one versioned poller
s.pollWithIdentity("v", "bid", true, false)
s.True(s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-100 * time.Microsecond)))
time.Sleep(time.Millisecond)
s.False(s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-100 * time.Microsecond)))
s.True(s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-10 * time.Millisecond)))
await.RequireTrue(s.T(), func() bool {
return !s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-10 * time.Millisecond))
}, 20*time.Millisecond, time.Millisecond)
}
func (s *PartitionManagerTestSuite) TestHasPollerAfter_Unversioned() {
@@ -1134,14 +1137,15 @@ func (s *PartitionManagerTestSuite) TestHasPollerAfter_Unversioned() {
// one unversioned poller
s.pollWithIdentity("uv", "", false, false)
s.True(s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-500 * time.Microsecond)))
s.True(s.partitionMgr.HasPollerAfter("", time.Now().Add(-500*time.Microsecond)))
time.Sleep(time.Millisecond)
s.False(s.partitionMgr.HasPollerAfter("", time.Now().Add(-100*time.Microsecond)))
s.True(s.partitionMgr.HasAnyPollerAfter(time.Now().Add(-10 * time.Millisecond)))
s.True(s.partitionMgr.HasPollerAfter("", time.Now().Add(-10*time.Millisecond)))
await.RequireTrue(s.T(), func() bool {
return !s.partitionMgr.HasPollerAfter("", time.Now().Add(-10*time.Millisecond))
}, 20*time.Millisecond, time.Millisecond)
// one versioned poller
s.pollWithIdentity("v", "bid", true, false)
s.False(s.partitionMgr.HasPollerAfter("", time.Now().Add(-100*time.Microsecond)))
s.False(s.partitionMgr.HasPollerAfter("", time.Now().Add(-10*time.Millisecond)))
}
func (s *PartitionManagerTestSuite) TestHasPollerAfter_Versioned() {
@@ -1151,13 +1155,14 @@ func (s *PartitionManagerTestSuite) TestHasPollerAfter_Versioned() {
// one version-set poller
bid := "bid"
s.pollWithIdentity("v", bid, true, false)
s.True(s.partitionMgr.HasPollerAfter(bid, time.Now().Add(-100*time.Microsecond)))
time.Sleep(time.Millisecond)
s.False(s.partitionMgr.HasPollerAfter(bid, time.Now().Add(-100*time.Microsecond)))
s.True(s.partitionMgr.HasPollerAfter(bid, time.Now().Add(-10*time.Millisecond)))
await.RequireTrue(s.T(), func() bool {
return !s.partitionMgr.HasPollerAfter(bid, time.Now().Add(-10*time.Millisecond))
}, 20*time.Millisecond, time.Millisecond)
// one unversioned poller
s.pollWithIdentity("uv", "", false, true)
s.False(s.partitionMgr.HasPollerAfter(bid, time.Now().Add(-100*time.Microsecond)))
s.False(s.partitionMgr.HasPollerAfter(bid, time.Now().Add(-10*time.Millisecond)))
}
func (s *PartitionManagerTestSuite) TestLegacyDescribeTaskQueue() {
@@ -1814,8 +1819,9 @@ func (s *PartitionManagerTestSuite) TestTaskAddHooks_MultipleHooksInvoked() {
type mockUserDataManager struct {
sync.Mutex
data *persistencespb.VersionedTaskQueueUserData
onChange UserDataOnChangeFunc
data *persistencespb.VersionedTaskQueueUserData
onChange UserDataOnChangeFunc
scaleInfo *taskqueuespb.PartitionScaleInfo
}
func (m *mockUserDataManager) Start() {
@@ -1863,6 +1869,18 @@ func (m *mockUserDataManager) LocalBacklogPriorityChanged(map[PhysicalTaskQueueV
panic("unused")
}
func (m *mockUserDataManager) SetPartitionScale(scaleInfo *taskqueuespb.PartitionScaleInfo) {
m.Lock()
defer m.Unlock()
m.scaleInfo = scaleInfo
}
func (m *mockUserDataManager) PartitionScale() *taskqueuespb.PartitionScaleInfo {
m.Lock()
defer m.Unlock()
return m.scaleInfo
}
func (m *mockUserDataManager) updateVersioningData(data *persistencespb.VersioningData) {
m.Lock()
defer m.Unlock()

27
service/matching/user_data_manager.go
View File

@@ -1,6 +1,7 @@
package matching
import (
"cmp"
"context"
"errors"
"slices"
@@ -58,6 +59,10 @@ type (
HandleGetUserDataRequest(ctx context.Context, req *matchingservice.GetTaskQueueUserDataRequest) (*matchingservice.GetTaskQueueUserDataResponse, error)
CheckTaskQueueUserDataPropagation(context.Context, int64, int, int) error
LocalBacklogPriorityChanged(map[PhysicalTaskQueueVersion]int64)
// SetPartitionScale is called on the root partition to propagate new scale info to child partitions.
SetPartitionScale(*taskqueuespb.PartitionScaleInfo)
// PartitionScale returns the current partition scale info from ephemeral data.
PartitionScale() *taskqueuespb.PartitionScaleInfo
}
UserDataUpdateOptions struct {
@@ -756,6 +761,23 @@ func (m *userDataManagerImpl) LocalBacklogPriorityChanged(backlogPriority map[Ph
})
}
// SetPartitionScale can only be called on a root partition.
func (m *userDataManagerImpl) SetPartitionScale(scaleInfo *taskqueuespb.PartitionScaleInfo) {
if !m.partition.IsRoot() {
return
}
m.updateEphemeralData(func(newData *taskqueuespb.EphemeralData) {
newData.Scale = scaleInfo
})
}
// PartitionScale gets the current partition scale state.
func (m *userDataManagerImpl) PartitionScale() *taskqueuespb.PartitionScaleInfo {
m.lock.Lock()
defer m.lock.Unlock()
return m.mergedEphemeralData.GetData().GetScale()
}
func (m *userDataManagerImpl) gotIncomingEphemeralData(eph *taskqueuespb.VersionedEphemeralData) {
if m.partition.IsRoot() {
// Root activity/nexus partition should not get ephemeral data from its fetch source
@@ -818,6 +840,11 @@ func (m *userDataManagerImpl) mergeEphemeralDataLocked() {
m.incomingEphemeralData.GetData().GetPartition(),
m.myEphemeralData.GetData().GetPartition(),
),
// scale info always comes from the root, so only one of these should be non-nil
Scale: cmp.Or(
m.incomingEphemeralData.GetData().GetScale(),
m.myEphemeralData.GetData().GetScale(),
),
},
Version: time.Now().UnixNano(),
}

176
service/matching/validate_partition_counts_test.go Normal file
View File

@@ -0,0 +1,176 @@
package matching
import (
"testing"
"github.com/stretchr/testify/require"
"go.temporal.io/api/serviceerror"
taskqueuespb "go.temporal.io/server/api/taskqueue/v1"
"go.temporal.io/server/client/matching"
"go.temporal.io/server/common/dynamicconfig"
serviceerrors "go.temporal.io/server/common/serviceerror"
)
func TestValidatePartitionCounts(t *testing.T) {
t.Parallel()
var stale *serviceerrors.StalePartitionCounts
var internal *serviceerror.Internal
tests := []struct {
name string
partitionID int
scaleInfo *taskqueuespb.PartitionScaleInfo
forWrite bool
expected any
}{
// partition id validity
{
name: "negative partition id",
partitionID: -1,
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 4},
forWrite: true,
expected: &internal,
},
{
name: "partition id == read",
partitionID: 8,
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 4},
forWrite: true,
expected: &stale,
},
{
name: "partition id > read",
partitionID: 10,
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 4},
forWrite: false,
expected: &stale,
},
{
name: "draining partition, write",
partitionID: 5,
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 4},
forWrite: true,
expected: &stale,
},
{
name: "draining partition, read",
partitionID: 5,
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 4},
forWrite: false,
expected: nil,
},
{
name: "active partition, write",
partitionID: 2,
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 4},
forWrite: true,
expected: nil,
},
{
name: "active partition, read",
partitionID: 2,
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 4},
forWrite: false,
expected: nil,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
err := validatePartitionCounts(tc.partitionID, tc.scaleInfo, tc.forWrite)
if tc.expected == nil {
require.NoError(t, err)
} else {
require.ErrorAs(t, err, tc.expected)
}
})
}
}
func TestValidatePartitionCountDifference(t *testing.T) {
t.Parallel()
var stale *serviceerrors.StalePartitionCounts
difference := dynamicconfig.PartitionScaleAllowedDrift{
Delta: 2,
Ratio: 1.5,
}
tests := []struct {
name string
scaleInfo *taskqueuespb.PartitionScaleInfo
clientPC matching.PartitionCounts
forWrite bool
expected any
}{
{
name: "client counts match",
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 8},
clientPC: matching.PartitionCounts{Read: 8, Write: 8},
forWrite: true,
expected: nil,
},
{
name: "client counts too far off",
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 8},
clientPC: matching.PartitionCounts{Read: 20, Write: 20},
forWrite: true,
expected: &stale,
},
{
name: "within delta",
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 8},
clientPC: matching.PartitionCounts{Read: 10, Write: 10},
forWrite: true,
expected: nil,
},
{
name: "within ratio (delta exceeds)",
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 100, Write: 100},
clientPC: matching.PartitionCounts{Read: 110, Write: 110},
forWrite: true,
expected: nil,
},
{
name: "both delta and ratio exceed",
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 4, Write: 4},
clientPC: matching.PartitionCounts{Read: 10, Write: 10},
forWrite: true,
expected: &stale,
},
{
name: "read path compares read counts",
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 4},
clientPC: matching.PartitionCounts{Read: 20, Write: 4},
forWrite: false,
expected: &stale,
},
{
name: "write path compares write counts",
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 8, Write: 4},
clientPC: matching.PartitionCounts{Read: 20, Write: 4},
forWrite: true,
expected: nil,
},
{
name: "client counts below server",
scaleInfo: &taskqueuespb.PartitionScaleInfo{Read: 20, Write: 20},
clientPC: matching.PartitionCounts{Read: 4, Write: 4},
forWrite: true,
expected: &stale,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
err := validatePartitionScaleDrift(tc.scaleInfo, tc.forWrite, tc.clientPC, difference)
if tc.expected == nil {
require.NoError(t, err)
} else {
require.ErrorAs(t, err, tc.expected)
}
})
}
}