Files
tasks/service/history/queues/scheduler_quotas.go
zeekay 5a601d081b rebrand: go.temporal.io/server -> github.com/hanzoai/tasks
Full module path rebrand from Temporal to Hanzo Tasks:
- go.mod module: github.com/hanzoai/tasks
- All internal Go imports updated (1906 files)
- cmd/server renamed to cmd/tasksd
- Binary name: tasksd (was temporal-server)
- Docker images: hanzoai/tasks (was temporalio/server)
- README rewritten with Hanzo Tasks branding
- LLM.md created with integration notes
- Makefile, goreleaser, CI workflows updated
- go.temporal.io/sdk and go.temporal.io/api imports preserved (separate repos)

Fork of Temporal.io (MIT license).
2026-03-19 10:36:40 -07:00

88 lines
2.3 KiB
Go

package queues
import (
"github.com/hanzoai/tasks/common/quotas"
"github.com/hanzoai/tasks/common/tasks"
)
type SchedulerRateLimiter quotas.RequestRateLimiter
func NewPrioritySchedulerRateLimiter(
namespaceRateFn quotas.NamespaceRateFn,
hostRateFn quotas.RateFn,
persistenceNamespaceRateFn quotas.NamespaceRateFn,
persistenceHostRateFn quotas.RateFn,
) (SchedulerRateLimiter, error) {
namespaceRateFnWithFallback := func(namespace string) float64 {
if rate := namespaceRateFn(namespace); rate > 0 {
return rate
}
return persistenceNamespaceRateFn(namespace)
}
hostRateFnWithFallback := func() float64 {
if rate := hostRateFn(); rate > 0 {
return rate
}
return persistenceHostRateFn()
}
requestPriorityFn := func(req quotas.Request) int {
// NOTE: task scheduler will use the string format for task priority as the caller type.
// see channelQuotaRequestFn in scheduler.go
// TODO: we don't need this hack when requestRateLimiter uses generics
if priority, ok := tasks.CallerTypeToPriority[req.CallerType]; ok {
return int(priority)
}
// default to low priority
return int(tasks.PriorityPreemptable)
}
priorityToRateLimiters := make(map[int]quotas.RequestRateLimiter, len(tasks.PriorityName))
for priority := range tasks.PriorityName {
priorityToRateLimiters[int(priority)] = newTaskRequestRateLimiter(
namespaceRateFnWithFallback,
hostRateFnWithFallback,
)
}
priorityLimiter := quotas.NewPriorityRateLimiter(requestPriorityFn, priorityToRateLimiters)
return priorityLimiter, nil
}
func newTaskRequestRateLimiter(
namespaceRateFn quotas.NamespaceRateFn,
hostRateFn quotas.RateFn,
) quotas.RequestRateLimiter {
hostRequestRateLimiter := quotas.NewRequestRateLimiterAdapter(
quotas.NewDefaultIncomingRateLimiter(hostRateFn),
)
namespaceRequestRateLimiterFn := func(req quotas.Request) quotas.RequestRateLimiter {
if len(req.Caller) == 0 {
return quotas.NoopRequestRateLimiter
}
return quotas.NewRequestRateLimiterAdapter(
quotas.NewDefaultIncomingRateLimiter(
func() float64 {
if rate := namespaceRateFn(req.Caller); rate > 0 {
return rate
}
return hostRateFn()
},
),
)
}
return quotas.NewMultiRequestRateLimiter(
quotas.NewNamespaceRequestRateLimiter(namespaceRequestRateLimiterFn),
hostRequestRateLimiter,
)
}