orbit/internal/dispatch/pool.go

// 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)
		}
	}
}