feat: add bounded task priority and fair application dispatch
This commit is contained in:
196
internal/dispatch/pool.go
Normal file
196
internal/dispatch/pool.go
Normal file
@@ -0,0 +1,196 @@
|
||||
// Package dispatch 对节点内领取机会按应用、队列轮转；数据库仍决定真实授权与配额。
|
||||
package dispatch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Job 只执行一个有界领取操作；不在此运行外部业务处理器。
|
||||
type Job struct {
|
||||
Context context.Context
|
||||
Application string
|
||||
Queue string
|
||||
Run func(context.Context)
|
||||
release func()
|
||||
}
|
||||
|
||||
// Pool 的内存队列可随节点重建，不保存任务状态。每个应用最多一个在途领取。
|
||||
type Pool struct {
|
||||
mu sync.Mutex
|
||||
closed bool
|
||||
ctx context.Context
|
||||
input chan Job
|
||||
done chan struct{}
|
||||
workers int
|
||||
slots chan struct{}
|
||||
}
|
||||
|
||||
func New(ctx context.Context, workers, capacity int) *Pool {
|
||||
if workers < 1 || capacity < 1 {
|
||||
panic("dispatch limits must be positive")
|
||||
}
|
||||
p := &Pool{ctx: ctx, input: make(chan Job, capacity), done: make(chan struct{}), workers: workers, slots: make(chan struct{}, capacity)}
|
||||
go p.run()
|
||||
return p
|
||||
}
|
||||
func (p *Pool) Done() <-chan struct{} { return p.done }
|
||||
|
||||
// Submit 非阻塞；调用者每个连接最多保存一个未完成 Job，满额后等待下一次扫描。
|
||||
func (p *Pool) Submit(job Job) bool {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.closed {
|
||||
return false
|
||||
}
|
||||
if job.Context == nil || job.Run == nil || job.Application == "" || job.Queue == "" || p.ctx.Err() != nil {
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case p.slots <- struct{}{}:
|
||||
default:
|
||||
return false
|
||||
}
|
||||
job.release = func() { <-p.slots }
|
||||
select {
|
||||
case p.input <- job:
|
||||
return true
|
||||
default:
|
||||
job.release()
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
type queue struct {
|
||||
name string
|
||||
jobs []Job
|
||||
}
|
||||
type application struct {
|
||||
name string
|
||||
busy bool
|
||||
queues []*queue
|
||||
}
|
||||
type roster struct{ apps []*application }
|
||||
|
||||
func (r *roster) add(j Job) {
|
||||
var a *application
|
||||
for _, v := range r.apps {
|
||||
if v.name == j.Application {
|
||||
a = v
|
||||
break
|
||||
}
|
||||
}
|
||||
if a == nil {
|
||||
a = &application{name: j.Application}
|
||||
r.apps = append(r.apps, a)
|
||||
}
|
||||
for _, q := range a.queues {
|
||||
if q.name == j.Queue {
|
||||
q.jobs = append(q.jobs, j)
|
||||
return
|
||||
}
|
||||
}
|
||||
a.queues = append(a.queues, &queue{name: j.Queue, jobs: []Job{j}})
|
||||
}
|
||||
func (r *roster) next() (Job, bool) {
|
||||
for range len(r.apps) {
|
||||
a := r.apps[0]
|
||||
r.apps = append(r.apps[1:], a)
|
||||
if a.busy {
|
||||
continue
|
||||
}
|
||||
for len(a.queues) > 0 {
|
||||
q := a.queues[0]
|
||||
j := q.jobs[0]
|
||||
q.jobs[0] = Job{}
|
||||
q.jobs = q.jobs[1:]
|
||||
a.queues = a.queues[1:]
|
||||
if len(q.jobs) > 0 {
|
||||
a.queues = append(a.queues, q)
|
||||
}
|
||||
if j.Context.Err() != nil {
|
||||
if j.release != nil {
|
||||
j.release()
|
||||
}
|
||||
continue
|
||||
}
|
||||
a.busy = true
|
||||
return j, true
|
||||
}
|
||||
}
|
||||
return Job{}, false
|
||||
}
|
||||
func (r *roster) finish(name string) {
|
||||
for _, a := range r.apps {
|
||||
if a.name == name {
|
||||
a.busy = false
|
||||
break
|
||||
}
|
||||
}
|
||||
r.prune()
|
||||
}
|
||||
func (r *roster) prune() {
|
||||
kept := r.apps[:0]
|
||||
for _, a := range r.apps {
|
||||
if a.busy || len(a.queues) > 0 {
|
||||
kept = append(kept, a)
|
||||
}
|
||||
}
|
||||
clear(r.apps[len(kept):])
|
||||
r.apps = kept
|
||||
}
|
||||
func (p *Pool) run() {
|
||||
defer close(p.done)
|
||||
var r roster
|
||||
finished := make(chan string, p.workers)
|
||||
running := 0
|
||||
for {
|
||||
if p.ctx.Err() != nil {
|
||||
p.mu.Lock()
|
||||
p.closed = true
|
||||
for len(p.input) > 0 {
|
||||
j := <-p.input
|
||||
j.release()
|
||||
}
|
||||
p.mu.Unlock()
|
||||
for _, a := range r.apps {
|
||||
for _, q := range a.queues {
|
||||
for _, j := range q.jobs {
|
||||
j.release()
|
||||
}
|
||||
}
|
||||
}
|
||||
for running > 0 {
|
||||
<-finished
|
||||
running--
|
||||
}
|
||||
return
|
||||
}
|
||||
for running < p.workers {
|
||||
j, ok := r.next()
|
||||
if !ok {
|
||||
r.prune()
|
||||
break
|
||||
}
|
||||
running++
|
||||
go func() {
|
||||
defer func() { j.release(); finished <- j.Application }()
|
||||
ctx, cancel := context.WithCancel(j.Context)
|
||||
stop := context.AfterFunc(p.ctx, cancel)
|
||||
defer cancel()
|
||||
defer stop()
|
||||
if ctx.Err() == nil {
|
||||
j.Run(ctx)
|
||||
}
|
||||
}()
|
||||
}
|
||||
select {
|
||||
case <-p.ctx.Done():
|
||||
case j := <-p.input:
|
||||
r.add(j)
|
||||
case name := <-finished:
|
||||
running--
|
||||
r.finish(name)
|
||||
}
|
||||
}
|
||||
}
|
||||
127
internal/dispatch/pool_test.go
Normal file
127
internal/dispatch/pool_test.go
Normal file
@@ -0,0 +1,127 @@
|
||||
package dispatch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRoundRobinApplicationsAndQueues(t *testing.T) {
|
||||
var r roster
|
||||
for i := 0; i < 40; i++ {
|
||||
r.add(Job{Context: context.Background(), Application: "noisy", Queue: "many-workers"})
|
||||
}
|
||||
for i := 0; i < 4; i++ {
|
||||
r.add(Job{Context: context.Background(), Application: "noisy", Queue: "other-queue"})
|
||||
r.add(Job{Context: context.Background(), Application: "quiet", Queue: "only-worker"})
|
||||
}
|
||||
want := []string{"noisy/many-workers", "quiet/only-worker", "noisy/other-queue", "quiet/only-worker", "noisy/many-workers", "quiet/only-worker", "noisy/other-queue", "quiet/only-worker"}
|
||||
for _, expected := range want {
|
||||
j, ok := r.next()
|
||||
if !ok || j.Application+"/"+j.Queue != expected {
|
||||
t.Fatalf("dispatch=%s/%s want=%s", j.Application, j.Queue, expected)
|
||||
}
|
||||
r.finish(j.Application)
|
||||
}
|
||||
}
|
||||
func TestPoolBoundsCancellationAndIsolation(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
p := New(ctx, 2, 4)
|
||||
started := make(chan string, 4)
|
||||
var active, same atomic.Int32
|
||||
block := func(app string) Job {
|
||||
return Job{Context: ctx, Application: app, Queue: "queue", Run: func(ctx context.Context) {
|
||||
if active.Add(1) > 2 {
|
||||
t.Error("global dispatch concurrency exceeded")
|
||||
}
|
||||
defer active.Add(-1)
|
||||
if app == "noisy" {
|
||||
if same.Add(1) > 1 {
|
||||
t.Error("application ran concurrent claims")
|
||||
}
|
||||
defer same.Add(-1)
|
||||
}
|
||||
started <- app
|
||||
<-ctx.Done()
|
||||
}}
|
||||
}
|
||||
for _, app := range []string{"noisy", "noisy", "noisy", "quiet"} {
|
||||
if !p.Submit(block(app)) {
|
||||
t.Fatal("unexpected rejection")
|
||||
}
|
||||
}
|
||||
if p.Submit(block("overflow")) {
|
||||
t.Fatal("unbounded pending jobs")
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
for range 2 {
|
||||
select {
|
||||
case app := <-started:
|
||||
seen[app] = true
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("one application blocked another")
|
||||
}
|
||||
}
|
||||
if !seen["quiet"] || !seen["noisy"] {
|
||||
t.Fatal("missing isolated application")
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case <-p.Done():
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("shutdown did not cancel claims")
|
||||
}
|
||||
if len(p.slots) != 0 || active.Load() != 0 || p.Submit(block("late")) {
|
||||
t.Fatal("shutdown leaked slots or accepted more work")
|
||||
}
|
||||
}
|
||||
func TestPoolConcurrentSubmitAndShutdown(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
p := New(ctx, 4, 32)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 20; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < 100; j++ {
|
||||
p.Submit(Job{Context: ctx, Application: fmt.Sprint(i % 5), Queue: fmt.Sprint(i % 3), Run: func(ctx context.Context) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-time.After(time.Millisecond):
|
||||
}
|
||||
}})
|
||||
}
|
||||
}()
|
||||
}
|
||||
cancel()
|
||||
wg.Wait()
|
||||
select {
|
||||
case <-p.Done():
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("shutdown stuck")
|
||||
}
|
||||
if len(p.slots) != 0 {
|
||||
t.Fatal("pending jobs leaked")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelledJobsDoNotUseQueueTurnOrCapacity(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
var r roster
|
||||
released := 0
|
||||
r.add(Job{Context: ctx, Application: "stale", Queue: "gone", release: func() { released++ }})
|
||||
r.add(Job{Context: context.Background(), Application: "ready", Queue: "ready"})
|
||||
j, ok := r.next()
|
||||
if !ok || j.Application != "ready" || released != 1 {
|
||||
t.Fatal("cancelled job blocked live work")
|
||||
}
|
||||
r.finish("ready")
|
||||
if len(r.apps) != 0 {
|
||||
t.Fatal("empty applications retained")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user