feat: add bounded task priority and fair application dispatch

This commit is contained in:
2026-09-17 18:45:54 +08:00
parent 958b4a8cea
commit 309ed75666
42 changed files with 1277 additions and 61 deletions

196
internal/dispatch/pool.go Normal file
View 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)
}
}
}

View 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")
}
}

View File

@@ -65,6 +65,7 @@ func TestIntegrationApplicationResourcesUpgrade(t *testing.T) {
assertApplicationCounts(t, f, 4, 2, 1)
// 还原后续新增表和字段,得到真实 v9 结构;现有任务内容保持不变。
f.exec(`DROP TABLE orbit_application_resources`)
f.exec(`ALTER TABLE pulse_tasks DROP INDEX ix_task_priority_dispatch, DROP CHECK ck_task_priority, DROP COLUMN dispatch_at`)
f.exec(`ALTER TABLE pulse_tasks DROP CHECK ck_task_concurrency_key, DROP COLUMN concurrency_key, DROP COLUMN concurrency_key_digest`)
f.exec(`ALTER TABLE pulse_attempts DROP COLUMN concurrency_token`)
f.exec(`ALTER TABLE pulse_worker_sessions DROP COLUMN concurrency_keys`)

View File

@@ -0,0 +1,104 @@
package mysql
import (
"context"
"encoding/json"
"net/http/httptest"
"os"
"testing"
"time"
"git.yuntuo.hk/k2safe/orbit/internal/transport/httpserver"
"git.yuntuo.hk/k2safe/orbit/pulse/worker"
)
func TestIntegrationFairDispatchApplicationIsolation(t *testing.T) {
dsn := os.Getenv("ORBIT_TEST_MYSQL_DSN")
if dsn == "" {
t.Skip("requires disposable MySQL")
}
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
defer cancel()
db, err := Open(ctx, dsn, 4, 2)
if err != nil {
t.Fatal(err)
}
defer db.Close()
if err = Migrate(ctx, db); err != nil {
t.Fatal(err)
}
noisy := newExecutionFixture(t, ctx, db)
quiet := newExecutionFixture(t, ctx, db)
if noisy.binding.ApplicationID != noisy.resources.ApplicationID || quiet.binding.ApplicationID != quiet.resources.ApplicationID {
t.Fatal("dispatch group was not derived from authenticated namespace")
}
hub := httpserver.NewWorkerHub(ctx, noisy.a)
server := httpserver.New("", func(context.Context) error { return nil })
httpserver.RegisterWorkers(server, hub)
tls := httptest.NewTLSServer(server.Handler)
defer tls.Close()
defer hub.Shutdown(ctx)
var connections []*testWorkerSocket
for range 12 {
f := *noisy
f.hello.InstanceID = executionID(t)
c := connectTestWorker(t, ctx, tls, &f)
id := c.send("worker.capacity", worker.CapacityInput{Sequence: 1, Capacity: 0}, "")
if reply := c.read(); reply.ReplyTo != id || reply.Error != nil {
t.Fatal("could not disable initial claims")
}
connections = append(connections, c)
}
quiet.hello.InstanceID = executionID(t)
q := connectTestWorker(t, ctx, tls, quiet)
id := q.send("worker.capacity", worker.CapacityInput{Sequence: 1, Capacity: 0}, "")
if reply := q.read(); reply.ReplyTo != id || reply.Error != nil {
t.Fatal("quiet worker setup failed")
}
noisy.submit(1)
want := quiet.submit(1)
lock, err := db.BeginTx(ctx, nil)
if err != nil {
t.Fatal(err)
}
defer lock.Rollback()
var ignored string
if err = lock.QueryRowContext(ctx, `SELECT HEX(application_id) FROM orbit_application_resources WHERE application_id=? FOR UPDATE`, noisy.resources.ApplicationID).Scan(&ignored); err != nil {
t.Fatal(err)
}
for _, c := range connections {
id := c.send("worker.capacity", worker.CapacityInput{Sequence: 2, Capacity: 1}, "")
if reply := c.read(); reply.ReplyTo != id || reply.Error != nil {
t.Fatal("noisy worker control blocked by claims")
}
}
// 等到真的有领取事务在额度行上等待,再验证另一应用;不靠固定 sleep 猜测故障已发生。
for {
var waiting int
if err = db.QueryRowContext(ctx, `SELECT COUNT(*) FROM performance_schema.data_lock_waits`).Scan(&waiting); err != nil {
t.Fatal(err)
}
if waiting > 0 {
break
}
select {
case <-ctx.Done():
t.Fatal("held application did not receive a claim")
case <-time.After(20 * time.Millisecond):
}
}
id = q.send("worker.capacity", worker.CapacityInput{Sequence: 2, Capacity: 1}, "")
if reply := q.read(); reply.ReplyTo != id || reply.Error != nil {
t.Fatal("other application heartbeat starved")
}
reply := q.read()
var assignment worker.Assignment
if reply.Type != "task.assign" || json.Unmarshal(reply.Body, &assignment) != nil || assignment.Task.ID != want.ID {
t.Fatal("quiet application could not claim while noisy application was blocked")
}
// 此时锁仍未释放;静默应用的真实领取已提交,证明没有耗尽四连接数据库池。
if err := lock.Rollback(); err != nil {
t.Fatal(err)
}
t.Log("12 noisy connections blocked on one application; another application claimed with a four-connection database pool before lock release")
}

View File

@@ -309,11 +309,12 @@ func TestIntegrationKeyedConcurrency(t *testing.T) {
sf := newScheduleFixture(t, f)
def := sf.definition()
def.Task.ConcurrencyKey = "schedule/resource"
def.Task.Priority = 7
v := sf.create(def)
sf.due(v, 2*time.Second)
sf.tick()
a, ref := f.claim()
if a.Task.ConcurrencyKey != def.Task.ConcurrencyKey || a.ConcurrencyToken == 0 {
if a.Task.ConcurrencyKey != def.Task.ConcurrencyKey || a.ConcurrencyToken == 0 || a.Task.Priority != 7 {
t.Fatal("schedule omitted frozen key")
}
noKeyClaim(t, f)

View File

@@ -0,0 +1,7 @@
-- 任务优先级只影响已到期任务的排序,最多提前 270 秒的排序优势;不提前执行延时任务。
-- 生成列随受理、退避及恢复后的 available_at 自动变化,禁止独立修改。
ALTER TABLE pulse_tasks
MODIFY COLUMN priority SMALLINT NOT NULL DEFAULT 0 COMMENT '冻结任务优先级,0至9;每级带来30秒排序优势,不改变最早执行时间',
ADD CONSTRAINT ck_task_priority CHECK (priority BETWEEN 0 AND 9),
ADD COLUMN dispatch_at DATETIME(6) GENERATED ALWAYS AS (TIMESTAMPADD(SECOND, -30 * priority, available_at)) STORED COMMENT '派发排序时间,由最早执行时间减优先级乘30秒生成;不是执行授权或可用时间',
ADD KEY ix_task_priority_dispatch(namespace_id,queue_id,state,dispatch_at,available_at,id);

View File

@@ -88,7 +88,7 @@ func TestIntegrationNoForeignKeysUpgrade(t *testing.T) {
}
indexes := func() []string {
t.Helper()
rows, err := db.QueryContext(ctx, `SELECT CONCAT_WS(':',TABLE_NAME,INDEX_NAME,NON_UNIQUE,SEQ_IN_INDEX,COLUMN_NAME,COALESCE(SUB_PART,0),INDEX_TYPE,IS_VISIBLE) FROM information_schema.STATISTICS WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME<>'orbit_application_resources' ORDER BY TABLE_NAME,INDEX_NAME,SEQ_IN_INDEX`)
rows, err := db.QueryContext(ctx, `SELECT CONCAT_WS(':',TABLE_NAME,INDEX_NAME,NON_UNIQUE,SEQ_IN_INDEX,COLUMN_NAME,COALESCE(SUB_PART,0),INDEX_TYPE,IS_VISIBLE) FROM information_schema.STATISTICS WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME<>'orbit_application_resources' AND INDEX_NAME<>'ix_task_priority_dispatch' ORDER BY TABLE_NAME,INDEX_NAME,SEQ_IN_INDEX`)
if err != nil {
t.Fatal(err)
}
@@ -135,6 +135,14 @@ func TestIntegrationNoForeignKeysUpgrade(t *testing.T) {
}
result := make([]string, 0, len(cols))
for i, value := range values {
if cols[i] == "dispatch_at" {
for j, col := range cols {
if col == "available_at" && value != values[j] {
t.Fatal("legacy default priority changed ordering")
}
}
continue
}
if cols[i] == "concurrency_key" || cols[i] == "concurrency_key_digest" {
if value.Valid {
t.Fatal("legacy task unexpectedly enabled keyed concurrency")

View File

@@ -0,0 +1,126 @@
package mysql
import (
"context"
"encoding/json"
"os"
"testing"
"time"
pulse "git.yuntuo.hk/k2safe/orbit/pulse/model"
"git.yuntuo.hk/k2safe/orbit/pulse/service"
)
func TestIntegrationDispatchPriority(t *testing.T) {
dsn := os.Getenv("ORBIT_TEST_MYSQL_DSN")
if dsn == "" {
t.Skip("requires disposable MySQL")
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
db, err := Open(ctx, dsn, 16, 4)
if err != nil {
t.Fatal(err)
}
defer db.Close()
if err = Migrate(ctx, db); err != nil {
t.Fatal(err)
}
submit := func(f *executionFixture, p int16, max uint32) service.TaskView {
t.Helper()
out, err := f.tasks.Submit(ctx, f.token, string(executionID(t)), service.SubmitInput{QueueID: f.resources.QueueID, Handler: "object.prepare", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), Priority: p, Retry: pulse.RetryPolicy{MaxAttempts: max}})
if err != nil {
t.Fatal(err)
}
return out.Task
}
t.Run("priority_overtakes_recent_work_but_not_future_release", func(t *testing.T) {
f := newExecutionFixture(t, ctx, db)
low := submit(f, 0, 1)
high := submit(f, 9, 1)
later := submit(f, 9, 1)
f.exec(`UPDATE pulse_tasks SET available_at=DATE_SUB(UTC_TIMESTAMP(6),INTERVAL 5 SECOND) WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, low.ID)
f.exec(`UPDATE pulse_tasks SET available_at=DATE_ADD(UTC_TIMESTAMP(6),INTERVAL 1 HOUR) WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, later.ID)
a, _ := f.claim()
if a.Task.ID != high.ID || a.Task.Priority != 9 {
t.Fatal("priority not honored")
}
a, _ = f.claim()
if a.Task.ID != low.ID {
t.Fatal("priority bypassed actual availability")
}
noKeyClaim(t, f)
page, err := f.tasks.ListTasks(ctx, f.token, service.ListInput{QueueID: f.resources.QueueID, Limit: 20})
if err != nil {
t.Fatal(err)
}
found := false
for _, v := range page.Items {
if v.ID == high.ID {
found = v.Priority == 9
}
}
if !found {
t.Fatal("summary lost frozen priority")
}
})
t.Run("old_low_priority_cannot_be_overtaken_by_unbounded_new_high_work", func(t *testing.T) {
f := newExecutionFixture(t, ctx, db)
old := submit(f, 0, 1)
f.exec(`UPDATE pulse_tasks SET available_at=DATE_SUB(UTC_TIMESTAMP(6),INTERVAL 271 SECOND) WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, old.ID)
for range 24 {
submit(f, 9, 1)
}
a, _ := f.claim()
if a.Task.ID != old.ID {
t.Fatal("aged task starved behind fresh high priority prefix")
}
})
t.Run("retry_keeps_priority_and_recomputes_sort_time", func(t *testing.T) {
f := newExecutionFixture(t, ctx, db)
task := submit(f, 9, 2)
first := retryStart(t, f)
retryFinish(t, f, first, "FAILED", true)
var gap int
var priority int16
if err := db.QueryRowContext(ctx, `SELECT TIMESTAMPDIFF(SECOND,dispatch_at,available_at),priority FROM pulse_tasks WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, task.ID).Scan(&gap, &priority); err != nil {
t.Fatal(err)
}
if gap != 270 || priority != 9 {
t.Fatal("retry lost priority or sort derivation")
}
noKeyClaim(t, f)
f.exec(`UPDATE pulse_tasks SET available_at=UTC_TIMESTAMP(6) WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, task.ID)
second := retryStart(t, f)
retryFinish(t, f, second, "SUCCEEDED", false)
if f.get(task.ID).Priority != 9 {
t.Fatal("completed task lost priority")
}
})
t.Run("identity_validation_and_database_constraint", func(t *testing.T) {
f := newExecutionFixture(t, ctx, db)
in := service.SubmitInput{QueueID: f.resources.QueueID, Handler: "object.prepare", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), Priority: 0}
first, err := f.tasks.Submit(ctx, f.token, "same", in)
if err != nil {
t.Fatal(err)
}
dupe, err := f.tasks.Submit(ctx, f.token, "same", in)
if err != nil || !dupe.Deduplicated || dupe.Task.ID != first.Task.ID {
t.Fatal("default priority identity changed")
}
in.Priority = 9
_, err = f.tasks.Submit(ctx, f.token, "same", in)
expectWorkerCode(t, err, "IDEMPOTENCY_CONFLICT")
for _, p := range []int16{-1, 10} {
in.Priority = p
_, err = f.tasks.Submit(ctx, f.token, string(executionID(t)), in)
expectWorkerCode(t, err, "INVALID_ARGUMENT")
}
if _, err := db.ExecContext(ctx, `UPDATE pulse_tasks SET priority=10 WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, first.Task.ID); err == nil {
t.Fatal("database priority range unchecked")
}
if _, err := db.ExecContext(ctx, `UPDATE pulse_tasks SET dispatch_at=UTC_TIMESTAMP(6) WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, first.Task.ID); err == nil {
t.Fatal("generated dispatch order independently writable")
}
})
}

View File

@@ -224,7 +224,7 @@ func (s *ScheduleStore) saveOccurrence(ctx context.Context, tx *sql.Tx, v schedu
if err != nil {
return err
}
_, err = tx.ExecContext(ctx, `INSERT INTO pulse_tasks(namespace_id,id,queue_id,handler,handler_version,payload_version,payload,policy_snapshot,recovery_safety,business_key,concurrency_key,concurrency_key_digest,owner_kind,owner_id,available_at,deadline_at,attempt_timeout_ms,max_attempts,created_at,updated_at) VALUES (?,?,?,?,?,?,?,?,?,NULLIF(?,''),?,?,'schedule',?,?,?,?,?,?,?)`, v.NamespaceID, value, v.QueueID, in.Handler, in.HandlerVersion, in.PayloadVersion, string(in.Payload), string(policy), in.RecoverySafety, in.BusinessKey, concurrencyKeyValue(in.ConcurrencyKey), concurrencyDigest(in.ConcurrencyKey), v.ID, now, now.Add(time.Duration(in.TaskTimeoutMS)*time.Millisecond), in.Retry.AttemptTimeoutMS, in.Retry.MaxAttempts, now, now)
_, err = tx.ExecContext(ctx, `INSERT INTO pulse_tasks(namespace_id,id,queue_id,handler,handler_version,payload_version,payload,policy_snapshot,recovery_safety,business_key,concurrency_key,concurrency_key_digest,priority,owner_kind,owner_id,available_at,deadline_at,attempt_timeout_ms,max_attempts,created_at,updated_at) VALUES (?,?,?,?,?,?,?,?,?,NULLIF(?,''),?,?,?,'schedule',?,?,?,?,?,?,?)`, v.NamespaceID, value, v.QueueID, in.Handler, in.HandlerVersion, in.PayloadVersion, string(in.Payload), string(policy), in.RecoverySafety, in.BusinessKey, concurrencyKeyValue(in.ConcurrencyKey), concurrencyDigest(in.ConcurrencyKey), in.Priority, v.ID, now, now.Add(time.Duration(in.TaskTimeoutMS)*time.Millisecond), in.Retry.AttemptTimeoutMS, in.Retry.MaxAttempts, now, now)
if err != nil {
return err
}

View File

@@ -195,8 +195,8 @@ func (s *TaskStore) Submit(ctx context.Context, token auth.Token, key string, in
if err := reserveQueueTask(ctx, tx, token.NamespaceID, in.QueueID); err != nil {
return err
}
_, err = tx.ExecContext(ctx, `INSERT INTO pulse_tasks(namespace_id,id,queue_id,handler,handler_version,payload_version,payload,policy_snapshot,recovery_safety,business_key,concurrency_key,concurrency_key_digest,available_at,start_deadline_at,deadline_at,attempt_timeout_ms,max_attempts,created_at,updated_at)
VALUES (?,?,?,?,?,?,?,?,?,NULLIF(?,''),?,?,?,?,?,?,?,?,?)`, token.NamespaceID, taskID, in.QueueID, in.Handler, in.HandlerVersion, in.PayloadVersion, string(in.Payload), string(policy), in.RecoverySafety, in.BusinessKey, concurrencyKeyValue(in.ConcurrencyKey), concurrencyDigest(in.ConcurrencyKey), available, in.StartDeadlineAt, deadline, in.Retry.AttemptTimeoutMS, in.Retry.MaxAttempts, now, now)
_, err = tx.ExecContext(ctx, `INSERT INTO pulse_tasks(namespace_id,id,queue_id,handler,handler_version,payload_version,payload,policy_snapshot,recovery_safety,business_key,concurrency_key,concurrency_key_digest,priority,available_at,start_deadline_at,deadline_at,attempt_timeout_ms,max_attempts,created_at,updated_at)
VALUES (?,?,?,?,?,?,?,?,?,NULLIF(?,''),?,?,?,?,?,?,?,?,?,?)`, token.NamespaceID, taskID, in.QueueID, in.Handler, in.HandlerVersion, in.PayloadVersion, string(in.Payload), string(policy), in.RecoverySafety, in.BusinessKey, concurrencyKeyValue(in.ConcurrencyKey), concurrencyDigest(in.ConcurrencyKey), in.Priority, available, in.StartDeadlineAt, deadline, in.Retry.AttemptTimeoutMS, in.Retry.MaxAttempts, now, now)
if err != nil {
return err
}
@@ -215,14 +215,14 @@ func (s *TaskStore) Submit(ctx context.Context, token auth.Token, key string, in
return out, err
}
const taskColumns = `id,namespace_id,queue_id,state,revision,handler,handler_version,payload_version,payload,result,COALESCE(business_key,''),COALESCE(concurrency_key,''),available_at,start_deadline_at,deadline_at,recovery_safety,policy_snapshot,attempt_count,cancel_requested_at,completed_at,created_at,updated_at,effect_uncertain`
const taskColumns = `id,namespace_id,queue_id,state,revision,handler,handler_version,payload_version,payload,result,COALESCE(business_key,''),COALESCE(concurrency_key,''),priority,available_at,start_deadline_at,deadline_at,recovery_safety,policy_snapshot,attempt_count,cancel_requested_at,completed_at,created_at,updated_at,effect_uncertain`
type scanner interface{ Scan(...any) error }
func scanTask(row scanner) (service.TaskView, error) {
var v service.TaskView
var policy, payload, result []byte
err := row.Scan(&v.ID, &v.NamespaceID, &v.QueueID, &v.State, &v.Revision, &v.Handler, &v.HandlerVersion, &v.PayloadVersion, &payload, &result, &v.BusinessKey, &v.ConcurrencyKey, &v.AvailableAt, &v.StartDeadlineAt, &v.DeadlineAt, &v.RecoverySafety, &policy, &v.AttemptCount, &v.CancelRequestedAt, &v.CompletedAt, &v.CreatedAt, &v.UpdatedAt, &v.EffectUncertain)
err := row.Scan(&v.ID, &v.NamespaceID, &v.QueueID, &v.State, &v.Revision, &v.Handler, &v.HandlerVersion, &v.PayloadVersion, &payload, &result, &v.BusinessKey, &v.ConcurrencyKey, &v.Priority, &v.AvailableAt, &v.StartDeadlineAt, &v.DeadlineAt, &v.RecoverySafety, &policy, &v.AttemptCount, &v.CancelRequestedAt, &v.CompletedAt, &v.CreatedAt, &v.UpdatedAt, &v.EffectUncertain)
if err != nil {
return v, err
}
@@ -272,7 +272,7 @@ func (s *TaskStore) ListTasks(ctx context.Context, token auth.Token, in service.
if !p.AllowsQueue(in.QueueID) {
return fault.Forbidden
}
q := `SELECT id,queue_id,state,revision,handler,handler_version,COALESCE(business_key,''),COALESCE(concurrency_key,''),available_at,deadline_at,attempt_count,created_at,updated_at,effect_uncertain FROM pulse_tasks WHERE namespace_id=? AND queue_id=?`
q := `SELECT id,queue_id,state,revision,handler,handler_version,COALESCE(business_key,''),COALESCE(concurrency_key,''),priority,available_at,deadline_at,attempt_count,created_at,updated_at,effect_uncertain FROM pulse_tasks WHERE namespace_id=? AND queue_id=?`
args := []any{token.NamespaceID, in.QueueID}
if in.State != "" {
q += " AND state=?"
@@ -296,7 +296,7 @@ func (s *TaskStore) ListTasks(ctx context.Context, token auth.Token, in service.
defer rows.Close()
for rows.Next() {
var v service.TaskSummary
if err := rows.Scan(&v.ID, &v.QueueID, &v.State, &v.Revision, &v.Handler, &v.HandlerVersion, &v.BusinessKey, &v.ConcurrencyKey, &v.AvailableAt, &v.DeadlineAt, &v.AttemptCount, &v.CreatedAt, &v.UpdatedAt, &v.EffectUncertain); err != nil {
if err := rows.Scan(&v.ID, &v.QueueID, &v.State, &v.Revision, &v.Handler, &v.HandlerVersion, &v.BusinessKey, &v.ConcurrencyKey, &v.Priority, &v.AvailableAt, &v.DeadlineAt, &v.AttemptCount, &v.CreatedAt, &v.UpdatedAt, &v.EffectUncertain); err != nil {
return err
}
out.Items = append(out.Items, v)

View File

@@ -152,7 +152,7 @@ func (s *ExecutionStore) Claim(ctx context.Context, token auth.Token, b worker.B
if !state.concurrencyKeys {
keyFilter = "concurrency_key_digest IS NULL"
}
rows, err := tx.QueryContext(ctx, `SELECT id,gate_id FROM pulse_tasks WHERE namespace_id=? AND queue_id=? AND state='WAITING' AND available_at<=UTC_TIMESTAMP(6) AND deadline_at>UTC_TIMESTAMP(6) AND (start_deadline_at IS NULL OR start_deadline_at>UTC_TIMESTAMP(6)) AND attempt_count<max_attempts AND (gate_id IS NULL OR EXISTS (SELECT 1 FROM pulse_execution_gates g WHERE g.namespace_id=pulse_tasks.namespace_id AND g.id=pulse_tasks.gate_id AND g.paused=FALSE)) AND (`+strings.Join(matches, " OR ")+`) AND `+keyFilter+` ORDER BY available_at,priority DESC,id LIMIT 16`, args...)
rows, err := tx.QueryContext(ctx, `SELECT id,gate_id FROM pulse_tasks WHERE namespace_id=? AND queue_id=? AND state='WAITING' AND available_at<=UTC_TIMESTAMP(6) AND deadline_at>UTC_TIMESTAMP(6) AND (start_deadline_at IS NULL OR start_deadline_at>UTC_TIMESTAMP(6)) AND attempt_count<max_attempts AND (gate_id IS NULL OR EXISTS (SELECT 1 FROM pulse_execution_gates g WHERE g.namespace_id=pulse_tasks.namespace_id AND g.id=pulse_tasks.gate_id AND g.paused=FALSE)) AND (`+strings.Join(matches, " OR ")+`) AND `+keyFilter+` ORDER BY dispatch_at,available_at,id LIMIT 16`, args...)
if err != nil {
return err
}

View File

@@ -118,6 +118,9 @@ func (s *ExecutionStore) Connect(ctx context.Context, token auth.Token, in worke
return err
}
binding = worker.Binding{InstanceID: in.InstanceID, SessionID: session, Epoch: epoch + 1, Proof: proof}
if err := tx.QueryRowContext(ctx, `SELECT application_id FROM orbit_namespaces WHERE id=?`, token.NamespaceID).Scan(&binding.ApplicationID); err != nil {
return err
}
state := sessionState{queue: in.QueueID, capacity: in.Capacity, draining: draining, expires: now.Add(s.options.SessionTTL), capabilities: in.Capabilities, concurrencyKeys: in.ConcurrencyKeys}
caps, err := json.Marshal(in.Capabilities)
if err != nil {

View File

@@ -12,6 +12,7 @@ import (
"sync"
"time"
"git.yuntuo.hk/k2safe/orbit/internal/dispatch"
"git.yuntuo.hk/k2safe/orbit/internal/fault"
"git.yuntuo.hk/k2safe/orbit/internal/jsonvalue"
"git.yuntuo.hk/k2safe/orbit/internal/jsonwire"
@@ -25,19 +26,25 @@ const WorkerSubprotocol = "pulse.worker.v1"
// WorkerHub 管理节点内连接生命周期;任务归属和容量仍由 MySQL 决定。
type WorkerHub struct {
service *worker.Service
ctx context.Context
cancel context.CancelFunc
slots chan struct{}
mu sync.Mutex
closing bool
wg sync.WaitGroup
poll time.Duration
service *worker.Service
ctx context.Context
cancel context.CancelFunc
slots chan struct{}
mu sync.Mutex
closing bool
wg sync.WaitGroup
poll time.Duration
dispatch *dispatch.Pool
emptyPolls map[[32]byte]time.Time
}
func NewWorkerHub(parent context.Context, service *worker.Service) *WorkerHub {
ctx, cancel := context.WithCancel(parent)
return &WorkerHub{service: service, ctx: ctx, cancel: cancel, slots: make(chan struct{}, 128), poll: time.Second}
h := &WorkerHub{service: service, ctx: ctx, cancel: cancel, slots: make(chan struct{}, 128), poll: time.Second, emptyPolls: make(map[[32]byte]time.Time)}
h.dispatch = dispatch.New(ctx, 4, 128)
h.wg.Add(1)
go func() { defer h.wg.Done(); <-h.dispatch.Done() }()
return h
}
func RegisterWorkers(server *http.Server, hub *WorkerHub) {
server.Handler.(*http.ServeMux).Handle("GET /v1/worker/connect", hub)
@@ -249,12 +256,53 @@ func (h *WorkerHub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
eligible := false
leases := map[core.ID]time.Time{}
cancelSent := map[core.ID]bool{}
results := make(chan claimResult, 1)
pending, burstLeft := false, 0
pollKey := workerPollKey(token.NamespaceID, input)
canRequest := func() bool {
return eligible && !draining && capacity > 0 && len(leases) < int(capacity) && len(outgoing) < 4
}
requestClaim := func() {
if pending || burstLeft == 0 || !canRequest() {
return
}
pending = h.dispatch.Submit(dispatch.Job{Context: ctx, Application: string(binding.ApplicationID), Queue: string(token.NamespaceID) + "/" + string(input.QueueID), Run: func(jobCtx context.Context) {
results <- h.claimFair(jobCtx, token, binding, pollKey)
}})
if pending {
burstLeft--
}
}
ticker := time.NewTicker(h.poll)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case result := <-results:
pending = false
if result.err != nil {
if fault.Code(result.err) == "STALE_SESSION" || fault.Code(result.err) == "UNAUTHENTICATED" || fault.Code(result.err) == "PERMISSION_DENIED" {
send(workerFrame{Type: "error", Error: publicWorkerError(result.err)}, true)
return
}
burstLeft = 0
continue
}
if result.assignment == nil {
burstLeft = 0
continue
}
assignment := result.assignment
raw, err := jsonwire.Marshal(assignment)
if err != nil {
return
}
leases[assignment.AttemptID] = assignment.LeaseUntil
if !send(workerFrame{Type: "task.assign", Body: raw}, false) {
return
}
requestClaim()
case frame := <-incoming:
call, stop := context.WithTimeout(ctx, 8*time.Second)
var body any
@@ -370,31 +418,9 @@ func (h *WorkerHub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
delete(leases, attempt)
}
}
if !eligible || draining || capacity == 0 || len(leases) >= int(capacity) || len(outgoing) >= 4 {
continue
}
for burst := 0; burst < 4 && len(leases) < int(capacity) && len(outgoing) < 4; burst++ {
call, stop := context.WithTimeout(ctx, 8*time.Second)
assignment, err := h.service.Claim(call, token, binding)
stop()
if err != nil {
if fault.Code(err) == "STALE_SESSION" || fault.Code(err) == "UNAUTHENTICATED" || fault.Code(err) == "PERMISSION_DENIED" {
send(workerFrame{Type: "error", Error: publicWorkerError(err)}, true)
return
}
break
}
if assignment == nil {
break
}
raw, err := jsonwire.Marshal(assignment)
if err != nil {
return
}
leases[assignment.AttemptID] = assignment.LeaseUntil
if !send(workerFrame{Type: "task.assign", Body: raw}, false) {
return
}
if !pending {
burstLeft = 4
requestClaim()
}
}
}

View File

@@ -194,6 +194,7 @@ func TestWorkerPublishedContract(t *testing.T) {
assignment := worker.Assignment{Task: task, AttemptID: id, Epoch: 1, LeaseSecret: secret, LeaseUntil: lease.LeaseUntil, ServerTime: now}
keyed := assignment
keyed.Task.ConcurrencyKey = "object/共享资源"
keyed.Task.Priority = 9
keyed.ConcurrencyToken = ^uint64(0) - 1
for _, tc := range []struct {
kind string
@@ -215,6 +216,8 @@ func TestWorkerPublishedContract(t *testing.T) {
keyedBody, _ := jsonwire.Marshal(keyed)
keyedFrame, _ := jsonwire.Marshal(workerFrame{Version: "1", Type: "task.assign", MessageID: id, SessionID: id, Body: keyedBody})
for _, raw := range []string{
strings.Replace(string(keyedFrame), `"priority":9`, `"priority":10`, 1),
strings.Replace(string(keyedFrame), `"priority":9`, `"priority":-1`, 1),
strings.Replace(string(keyedFrame), `"concurrency_token":"18446744073709551614",`, "", 1),
strings.Replace(string(keyedFrame), `"concurrency_key":"object/共享资源",`, "", 1),
strings.Replace(string(keyedFrame), `"concurrency_token":"18446744073709551614"`, `"concurrency_token":0`, 1),

View File

@@ -0,0 +1,60 @@
package httpserver
import (
"context"
"crypto/sha256"
"encoding/json"
"time"
"git.yuntuo.hk/k2safe/orbit/internal/fault"
core "git.yuntuo.hk/k2safe/orbit/model"
"git.yuntuo.hk/k2safe/orbit/platform/auth"
"git.yuntuo.hk/k2safe/orbit/pulse/worker"
)
type claimResult struct {
assignment *worker.Assignment
err error
}
func workerPollKey(ns core.ID, in worker.ConnectInput) [32]byte {
// 正式注册已校验;再次规范化仅取得稳定能力集合,不包含任何秘密。
in, _ = worker.NormalizeConnect(in)
raw, _ := json.Marshal(struct {
Namespace, Queue core.ID
Capabilities []worker.Capability
Keys bool
}{ns, in.QueueID, in.Capabilities, in.ConcurrencyKeys})
return sha256.Sum256(raw)
}
func (h *WorkerHub) claimFair(ctx context.Context, t auth.Token, b worker.Binding, key [32]byte) (out claimResult) {
// HTTP handler 原本具有 panic 隔离;异步领取同样不能因一次异常杀死其他连接。
defer func() {
if recover() != nil {
out = claimResult{err: fault.Unavailable}
}
}()
h.mu.Lock()
until := h.emptyPolls[key]
h.mu.Unlock()
if until.After(time.Now()) {
return out
}
call, stop := context.WithTimeout(ctx, 8*time.Second)
defer stop()
out.assignment, out.err = h.service.Claim(call, t, b)
if out.assignment == nil && out.err == nil {
h.mu.Lock()
now := time.Now()
for k, v := range h.emptyPolls {
if !v.After(now) {
delete(h.emptyPolls, k)
}
}
if len(h.emptyPolls) < 128 {
h.emptyPolls[key] = now.Add(h.poll)
}
h.mu.Unlock()
}
return out
}

View File

@@ -0,0 +1,75 @@
package httpserver
import (
"context"
"strings"
"testing"
"time"
core "git.yuntuo.hk/k2safe/orbit/model"
"git.yuntuo.hk/k2safe/orbit/platform/auth"
"git.yuntuo.hk/k2safe/orbit/pulse/worker"
)
type emptyClaimStore struct {
worker.Store
calls int
panicNext bool
}
func (s *emptyClaimStore) Claim(context.Context, auth.Token, worker.Binding) (*worker.Assignment, error) {
s.calls++
if s.panicNext {
s.panicNext = false
panic("private database details")
}
return nil, nil
}
func TestEmptyPollCoalescesExactCapabilitiesAndBoundsCache(t *testing.T) {
store := &emptyClaimStore{}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
h := NewWorkerHub(ctx, worker.New(store))
defer h.Shutdown(context.Background())
in := worker.ConnectInput{InstanceID: core.ID(strings.Repeat("1", 32)), InstanceSecret: strings.Repeat("A", 43), QueueID: core.ID(strings.Repeat("2", 32)), Capacity: 1, Capabilities: []worker.Capability{{Handler: "a", HandlerVersion: "v1", PayloadVersion: "v1"}, {Handler: "b", HandlerVersion: "v1", PayloadVersion: "v1"}}}
ns := core.ID(strings.Repeat("3", 32))
key := workerPollKey(ns, in)
for range 20 {
if out := h.claimFair(ctx, auth.Token{}, worker.Binding{}, key); out.err != nil || out.assignment != nil {
t.Fatal("unexpected poll result")
}
}
if store.calls != 1 {
t.Fatal("identical idle connections scanned separately")
}
in.Capabilities[0], in.Capabilities[1] = in.Capabilities[1], in.Capabilities[0]
if workerPollKey(ns, in) != key {
t.Fatal("capability order changed coalescing identity")
}
in.ConcurrencyKeys = true
keyed := workerPollKey(ns, in)
if keyed == key {
t.Fatal("legacy and keyed workers shared an eligibility cache")
}
h.claimFair(ctx, auth.Token{}, worker.Binding{}, keyed)
if store.calls != 2 {
t.Fatal("different capabilities hidden by cache")
}
h.emptyPolls[key] = time.Now().Add(-time.Second)
h.claimFair(ctx, auth.Token{}, worker.Binding{}, key)
if store.calls != 3 {
t.Fatal("expired cache prevented rediscovery")
}
for i := 0; i < 256; i++ {
k := [32]byte{byte(i), 1}
h.claimFair(ctx, auth.Token{}, worker.Binding{}, k)
}
if len(h.emptyPolls) > 128 {
t.Fatal("unbounded empty poll cache")
}
store.panicNext = true
out := h.claimFair(ctx, auth.Token{}, worker.Binding{}, [32]byte{1, 2, 3})
if out.err == nil || strings.Contains(out.err.Error(), "private") {
t.Fatal("claim panic escaped isolation")
}
}