tasks/common/tqid/task_queue_id.go

package tqid

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	enumspb "go.temporal.io/api/enums/v1"
	"go.temporal.io/api/serviceerror"
	taskqueuepb "go.temporal.io/api/taskqueue/v1"
	taskqueuespb "go.temporal.io/server/api/taskqueue/v1"
)

const (
	// nonRootPartitionPrefix is the prefix for all mangled task queue names.
	nonRootPartitionPrefix = "/_sys/"
	partitionDelimiter     = "/"
)

type (
	// TaskQueueFamily represents the high-level "task queue" that user creates by explicitly providing a task queue name
	// when starting a worker or a workflow. A task queue family consists of separate TaskQueues for different types of
	// task (e.g. Workflow, Activity).
	TaskQueueFamily struct {
		namespaceId string
		// this can be any string as long as it does not start with /_sys/.
		name string
	}

	// TaskQueue represents a logical task queue for a type of tasks (e.g. Activity or Workflow). Under the hood,
	// a TaskQueue can be broken down to multiple sticky or normal partitions.
	TaskQueue struct {
		family   TaskQueueFamily
		taskType enumspb.TaskQueueType
	}

	// Partition is a sticky or normal partition of a TaskQueue.
	// Each Partition has a distinct task queue partition manager in memory in Matching service.
	// Normal partition with `partitionId=0` is called the "root". Sticky queues are not considered root.
	Partition interface {
		NamespaceId() string
		TaskQueue() *TaskQueue
		TaskType() enumspb.TaskQueueType
		// IsRoot always returns false for Sticky partitions
		IsRoot() bool
		Kind() enumspb.TaskQueueKind
		IsChild() bool
		// PersistenceTTL returns the TTL for the partition's persistence metadata.
		// A zero value means no TTL (metadata persists indefinitely).
		PersistenceTTL() time.Duration
		SupportsFairness() bool
		SupportsVersioning() bool
		// SupportsPartitions returns true for partition types that can be split into
		// multiple partitions for load distribution. When false, partition count must be 1.
		SupportsPartitions() bool
		// MetricTag returns a label for this partition kind to be used in metrics.
		// When partitionIDBreakdown is true, partitions that support multiple partitions
		// return their partition ID instead of the kind label.
		MetricTag(partitionIDBreakdown bool) string

		// RpcName returns the mangled name of the task queue partition, to be used in RPCs.
		//
		// RPC names look like this:
		//
		//  sticky partition:			<sticky name>
		//	root normal partition: 		<task queue name>
		//	non-root normal partition: 	/_sys/<task queue name>/<partition id>
		//
		// This scheme lets users use anything they like for a base name, except for strings
		// starting with "/_sys/", without ambiguity.
		//
		// For backward compatibility, unversioned low-level task queues with partition 0 do not
		// use mangled names, they use the bare base name.
		RpcName() string
		Key() PartitionKey
		// RoutingKey returns the key and index within the key that should be used to find the
		// owner of a task queue partition.
		RoutingKey(int) (string, int)
		// GradualChangeKey returns an identifier that can be used with gradual changes.
		GradualChangeKey() []byte
	}

	// NormalPartition is used to distribute load of a TaskQueue in multiple Matching instances. A normal partition is
	// identified by `partitionId`. The partition with ID 0 is called a root partition.
	NormalPartition struct {
		taskQueue   *TaskQueue
		partitionId int
	}

	// StickyPartition is made by SDK for a single workflow worker to keep workflow tasks of the same execution
	// in the same worker for caching benefits. Each sticky partition is identified by a unique `stickyName`
	// generated by SDK. A StickyPartition can only have workflow task type.
	StickyPartition struct {
		stickyName string
		taskQueue  *TaskQueue
	}

	// WorkerCommandsPartition is used for server-to-worker communication (e.g. activity cancellations).
	// These queues are per-worker-process and only exist for the lifetime of the worker process. The SDK sets
	// Kind=TASK_QUEUE_KIND_WORKER_COMMANDS when polling on these queues.
	WorkerCommandsPartition struct {
		taskQueue *TaskQueue
	}

	// PartitionKey uniquely identifies a task queue partition, to be used in maps.
	// Note that task queue kind (sticky vs normal) and normal name for sticky task queues are not
	// part of the task queue partition identity.
	PartitionKey struct {
		namespaceId string
		name        string
		partitionId int
		taskType    enumspb.TaskQueueType
	}
)

var _ Partition = (*NormalPartition)(nil)
var _ Partition = (*StickyPartition)(nil)
var _ Partition = (*WorkerCommandsPartition)(nil)

var (
	ErrNoParent      = errors.New("root task queue partition has no parent")
	ErrInvalidDegree = errors.New("invalid task queue partition branching degree")
	ErrNonZeroSticky = errors.New("only sticky partitions can not have non-zero partition ID")
)

// NewTaskQueueFamily takes a user-provided task queue name (aka family name) and returns a TaskQueueFamily. Returns an
// error if name looks like a mangled name.
func NewTaskQueueFamily(namespaceId string, name string) (*TaskQueueFamily, error) {
	if strings.HasPrefix(name, nonRootPartitionPrefix) {
		return nil, serviceerror.NewInvalidArgument("task queue family name cannot have prefix /_sys/ " + name)
	}
	return &TaskQueueFamily{
		namespaceId: namespaceId,
		name:        name,
	}, nil
}

// UnsafeTaskQueueFamily returns a TaskQueueFamily object without validating the task queue name.
// This method should only be used in logs/metrics, not in the server logic (use NewTaskQueueFamily instead).
func UnsafeTaskQueueFamily(namespaceId string, name string) *TaskQueueFamily {
	return &TaskQueueFamily{namespaceId, name}
}

// UnsafePartitionFromProto tries parsing proto using PartitionFromProto but if it fails still returns a Partition
// object using the raw values in the proto.
// This method should only be used in logs/metrics, not in the server logic.
func UnsafePartitionFromProto(proto *taskqueuepb.TaskQueue, namespaceId string, taskType enumspb.TaskQueueType) Partition {
	p, err := PartitionFromProto(proto, namespaceId, taskType)
	if err == nil {
		return p
	}
	switch proto.GetKind() { //nolint:exhaustive
	case enumspb.TASK_QUEUE_KIND_STICKY:
		tq := &TaskQueue{TaskQueueFamily{namespaceId, proto.GetNormalName()}, taskType}
		return tq.StickyPartition(proto.GetName())
	case enumspb.TASK_QUEUE_KIND_WORKER_COMMANDS:
		tq := &TaskQueue{TaskQueueFamily{namespaceId, proto.GetName()}, taskType}
		return tq.WorkerCommandsPartition()
	default:
		tq := &TaskQueue{TaskQueueFamily{namespaceId, proto.GetName()}, taskType}
		return tq.RootPartition()
	}
}

func PartitionFromProto(proto *taskqueuepb.TaskQueue, namespaceId string, taskType enumspb.TaskQueueType) (Partition, error) {
	baseName, partition, err := parseRpcName(proto.GetName())
	if err != nil {
		return nil, err
	}

	kind := proto.GetKind()
	normalName := proto.GetNormalName()
	if normalName != "" && kind != enumspb.TASK_QUEUE_KIND_STICKY {
		return nil, serviceerror.NewInvalidArgumentf("only sticky queues can have normal name. tq: %s, normal name: %s", baseName, normalName)
	}

	switch kind {
	case enumspb.TASK_QUEUE_KIND_STICKY:
		if partition != 0 {
			return nil, fmt.Errorf("%w. base name: %s, normal name: %s", ErrNonZeroSticky, baseName, normalName)
		}
		tq := &TaskQueue{TaskQueueFamily{namespaceId, normalName}, taskType}
		return tq.StickyPartition(baseName), nil
	case enumspb.TASK_QUEUE_KIND_WORKER_COMMANDS:
		if partition != 0 {
			return nil, serviceerror.NewInvalidArgumentf("worker-commands partitions cannot have non-zero partition ID. base name: %s", baseName)
		}
		if taskType != enumspb.TASK_QUEUE_TYPE_NEXUS {
			return nil, serviceerror.NewInvalidArgumentf("worker-commands partitions must have nexus task type, got: %v. base name: %s", taskType, baseName)
		}
		tq := &TaskQueue{TaskQueueFamily{namespaceId, baseName}, taskType}
		return tq.WorkerCommandsPartition(), nil
	default:
		tq := &TaskQueue{TaskQueueFamily{namespaceId, baseName}, taskType}
		return tq.NormalPartition(partition), nil
	}
}

func PartitionFromPartitionProto(proto *taskqueuespb.TaskQueuePartition, namespaceId string) Partition {
	tq := &TaskQueue{TaskQueueFamily{namespaceId, proto.GetTaskQueue()}, proto.GetTaskQueueType()}
	switch proto.GetPartitionId().(type) {
	case *taskqueuespb.TaskQueuePartition_StickyName:
		return tq.StickyPartition(proto.GetStickyName())
	case *taskqueuespb.TaskQueuePartition_WorkerCommands:
		return tq.WorkerCommandsPartition()
	default:
		return tq.NormalPartition(int(proto.GetNormalPartitionId()))
	}
}

func NormalPartitionFromRpcName(rpcName string, namespaceId string, taskType enumspb.TaskQueueType) (*NormalPartition, error) {
	baseName, partition, err := parseRpcName(rpcName)
	if err != nil {
		return nil, err
	}
	tq := &TaskQueue{TaskQueueFamily{namespaceId, baseName}, taskType}
	return tq.NormalPartition(partition), nil
}

func MustNormalPartitionFromRpcName(rpcName string, namespaceId string, taskType enumspb.TaskQueueType) *NormalPartition {
	p, err := NormalPartitionFromRpcName(rpcName, namespaceId, taskType)
	if err != nil {
		panic(err)
	}
	return p
}

func (n *TaskQueueFamily) Name() string {
	return n.name
}

func (n *TaskQueueFamily) NamespaceId() string {
	return n.namespaceId
}

func (n *TaskQueueFamily) TaskQueue(taskType enumspb.TaskQueueType) *TaskQueue {
	return &TaskQueue{
		family:   *n,
		taskType: taskType,
	}
}

func (n *TaskQueue) Name() string {
	return n.family.Name()
}

func (n *TaskQueue) Family() *TaskQueueFamily {
	return &n.family
}

func (n *TaskQueue) NamespaceId() string {
	return n.family.NamespaceId()
}

func (n *TaskQueue) TaskType() enumspb.TaskQueueType {
	return n.taskType
}

func (n *TaskQueue) NormalPartition(partitionId int) *NormalPartition {
	return &NormalPartition{
		taskQueue:   n,
		partitionId: partitionId,
	}
}

func (n *TaskQueue) StickyPartition(stickyName string) *StickyPartition {
	return &StickyPartition{stickyName, n}
}

func (n *TaskQueue) WorkerCommandsPartition() *WorkerCommandsPartition {
	return &WorkerCommandsPartition{n}
}

func (n *TaskQueue) RootPartition() *NormalPartition {
	return n.NormalPartition(0)
}

func (s *StickyPartition) StickyName() string {
	return s.stickyName
}

func (s *StickyPartition) TaskType() enumspb.TaskQueueType {
	return s.taskQueue.TaskType()
}

func (s *StickyPartition) Kind() enumspb.TaskQueueKind {
	return enumspb.TASK_QUEUE_KIND_STICKY
}

func (s *StickyPartition) NamespaceId() string {
	return s.taskQueue.family.NamespaceId()
}

func (s *StickyPartition) RootPartition() Partition {
	return s
}

func (s *StickyPartition) TaskQueue() *TaskQueue {
	return s.taskQueue
}

func (s *StickyPartition) IsRoot() bool {
	return false
}

func (s *StickyPartition) IsChild() bool {
	return false
}

func (s *StickyPartition) PersistenceTTL() time.Duration { return 24 * time.Hour }
func (s *StickyPartition) SupportsFairness() bool        { return false }
func (s *StickyPartition) SupportsVersioning() bool      { return false }
func (s *StickyPartition) SupportsPartitions() bool      { return false }
func (s *StickyPartition) MetricTag(bool) string         { return "__sticky__" }

func (s *StickyPartition) RpcName() string {
	return s.stickyName
}

func (s *StickyPartition) Key() PartitionKey {
	return PartitionKey{
		namespaceId: s.NamespaceId(),
		name:        s.StickyName(),
		taskType:    s.TaskType(),
	}
}

func (s *StickyPartition) RoutingKey(int) (string, int) {
	return fmt.Sprintf("%s:%s:%d", s.NamespaceId(), s.RpcName(), s.TaskType()), 0
}

func (s *StickyPartition) GradualChangeKey() []byte {
	key := fmt.Sprintf("%s:%s:%d", s.NamespaceId(), s.RpcName(), s.TaskType())
	return []byte(key)
}

func (w *WorkerCommandsPartition) TaskType() enumspb.TaskQueueType {
	return w.taskQueue.TaskType()
}

func (w *WorkerCommandsPartition) Kind() enumspb.TaskQueueKind {
	return enumspb.TASK_QUEUE_KIND_WORKER_COMMANDS
}

func (w *WorkerCommandsPartition) NamespaceId() string { //nolint:stylecheck,staticcheck // matches Partition interface
	return w.taskQueue.family.NamespaceId()
}

func (w *WorkerCommandsPartition) TaskQueue() *TaskQueue {
	return w.taskQueue
}

func (w *WorkerCommandsPartition) IsRoot() bool {
	return false
}

func (w *WorkerCommandsPartition) IsChild() bool {
	return false
}

func (w *WorkerCommandsPartition) PersistenceTTL() time.Duration { return 24 * time.Hour }
func (w *WorkerCommandsPartition) SupportsFairness() bool        { return false }
func (w *WorkerCommandsPartition) SupportsVersioning() bool      { return false }
func (w *WorkerCommandsPartition) SupportsPartitions() bool      { return false }
func (w *WorkerCommandsPartition) MetricTag(bool) string         { return "__worker_commands__" }

func (w *WorkerCommandsPartition) RpcName() string { //nolint:stylecheck,staticcheck // matches Partition interface
	return w.taskQueue.Name()
}

func (w *WorkerCommandsPartition) Key() PartitionKey {
	return PartitionKey{
		namespaceId: w.NamespaceId(),
		name:        w.taskQueue.Name(),
		taskType:    w.TaskType(),
	}
}

func (w *WorkerCommandsPartition) RoutingKey(int) (string, int) {
	return fmt.Sprintf("%s:%s:%d", w.NamespaceId(), w.RpcName(), w.TaskType()), 0
}

func (w *WorkerCommandsPartition) GradualChangeKey() []byte {
	key := fmt.Sprintf("%s:%s:%d", w.NamespaceId(), w.RpcName(), w.TaskType())
	return []byte(key)
}

func (p *NormalPartition) TaskQueue() *TaskQueue {
	return p.taskQueue
}

func (p *NormalPartition) IsRoot() bool {
	return p.partitionId == 0
}

func (p *NormalPartition) IsChild() bool {
	return !p.IsRoot()
}

func (p *NormalPartition) PersistenceTTL() time.Duration { return 0 }
func (p *NormalPartition) SupportsFairness() bool        { return true }
func (p *NormalPartition) SupportsVersioning() bool      { return true }
func (p *NormalPartition) SupportsPartitions() bool      { return true }
func (p *NormalPartition) MetricTag(partitionIDBreakdown bool) string {
	if partitionIDBreakdown {
		return strconv.Itoa(p.partitionId)
	}
	return "__normal__"
}

func (p *NormalPartition) Kind() enumspb.TaskQueueKind {
	return enumspb.TASK_QUEUE_KIND_NORMAL
}

func (p *NormalPartition) PartitionId() int {
	return p.partitionId
}

func (p *NormalPartition) NamespaceId() string {
	return p.taskQueue.family.namespaceId
}

func (p *NormalPartition) TaskType() enumspb.TaskQueueType {
	return p.taskQueue.taskType
}

// ParentPartition returns a NormalPartition for the parent partition, using the given branching degree.
func (p *NormalPartition) ParentPartition(degree int) (*NormalPartition, error) {
	if p.IsRoot() {
		return nil, ErrNoParent
	} else if degree < 1 {
		return nil, ErrInvalidDegree
	}
	parent := (p.partitionId+degree-1)/degree - 1
	return p.taskQueue.NormalPartition(parent), nil
}

func (p *NormalPartition) RpcName() string {
	if p.IsRoot() {
		return p.TaskQueue().family.Name()
	}
	return nonRootPartitionPrefix + p.TaskQueue().Name() + partitionDelimiter + strconv.Itoa(p.partitionId)
}

func (p *NormalPartition) Key() PartitionKey {
	return PartitionKey{
		namespaceId: p.NamespaceId(),
		name:        p.TaskQueue().Name(),
		partitionId: p.partitionId,
		taskType:    p.TaskType(),
	}
}

func (p *NormalPartition) RoutingKey(batchSize int) (string, int) {
	if batchSize == 0 {
		return fmt.Sprintf("%s:%s:%d", p.NamespaceId(), p.RpcName(), p.TaskType()), 0
	}
	// We want to use LookupN to spread partitions across available nodes, but LookupN takes O(n)
	// time and space, so we should limit the n that we pass to it. Reduce the partition id by some
	// factor and move that factor into the name.
	batch := p.partitionId / batchSize
	if batch == 0 {
		// For the first batch, omit the batch number as if batchSize==0 so that the root
		// partition (id 0, batch number 0) doesn't move if we turn on/off spreading.
		return fmt.Sprintf("%s:%s:%d", p.NamespaceId(), p.TaskQueue().Name(), p.TaskType()), p.partitionId
	}
	// For subsequent batches, put the batch number in the key to hash.
	return fmt.Sprintf("%s:%s:%d:%d", p.NamespaceId(), p.TaskQueue().Name(), batch, p.TaskType()),
		p.partitionId % batchSize
}

func (p *NormalPartition) GradualChangeKey() []byte {
	key := fmt.Sprintf("%s:%s:%d", p.NamespaceId(), p.RpcName(), p.TaskType())
	return []byte(key)
}

// parseRpcName takes the rpc name of a task queue partition and returns a ParseTaskQueuePartition.
// Returns an error if the given name is not a valid rpc name.
func parseRpcName(rpcName string) (string, int, error) {
	baseName := rpcName
	partition := 0

	if strings.HasPrefix(rpcName, nonRootPartitionPrefix) {
		suffixOff := strings.LastIndex(rpcName, partitionDelimiter)
		if suffixOff <= len(nonRootPartitionPrefix) {
			return "", 0, serviceerror.NewInvalidArgument("invalid task queue partition name " + rpcName)
		}
		baseName = rpcName[len(nonRootPartitionPrefix):suffixOff]
		suffix := rpcName[suffixOff+1:]

		var err error
		partition, err = strconv.Atoi(suffix)
		if err != nil || partition <= 0 {
			return "", 0, serviceerror.NewInvalidArgument("invalid task queue partition name " + rpcName)
		}
	}

	if strings.HasPrefix(baseName, nonRootPartitionPrefix) {
		return "", 0, serviceerror.NewInvalidArgument("task queue family name cannot have prefix /_sys/ " + baseName)
	}
	return baseName, partition, nil
}