package mysql import ( "context" "database/sql" "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/service" "git.yuntuo.hk/k2safe/orbit/pulse/worker" ) // 以下额度操作要求调用者先取得应用额度锁,再取得队列排他锁,且在同一事务中写入任务/尝试。 func reserveQueueTask(ctx context.Context, tx *sql.Tx, ns, queue core.ID) error { v, err := readApplicationResources(ctx, tx, ns) if err != nil { return err } if v.UnfinishedTasks >= uint64(v.MaxUnfinishedTasks) { return fault.ApplicationFull } r, err := tx.ExecContext(ctx, `UPDATE pulse_queues SET unfinished_tasks=unfinished_tasks+1 WHERE namespace_id=? AND id=? AND unfinished_tasks0`, ns) if err := exactlyOne(r, err); err != nil { return err } r, err = tx.ExecContext(ctx, `UPDATE pulse_queues SET unfinished_tasks=unfinished_tasks-1 WHERE namespace_id=? AND id=? AND unfinished_tasks>0`, ns, queue) return exactlyOne(r, err) } // queueStartBudget 只读取当前窗口,不把空轮询写成新的窗口。 // 窗口由第一次成功 Start 建立;以数据库时间计算,节点重启不清零。 func queueStartBudget(ctx context.Context, tx *sql.Tx, ns, queue core.ID, now time.Time) (uint32, uint32, *time.Time, error) { var rate, used uint32 var until *time.Time err := tx.QueryRowContext(ctx, `SELECT rate_per_second,starts_in_window,start_window_until FROM pulse_queues WHERE namespace_id=? AND id=?`, ns, queue).Scan(&rate, &used, &until) if err != nil { return 0, 0, nil, err } if until == nil || !until.After(now) { used, until = 0, nil } return rate, used, until, nil } // queueCanClaim 把尚未 Start 的预占也算入当前窗口可用空间,避免限流时空耗尝试预算。 func queueCanClaim(ctx context.Context, tx *sql.Tx, ns, queue core.ID, now time.Time) (bool, error) { if ok, err := applicationCanClaim(ctx, tx, ns, now); err != nil || !ok { return false, err } rate, used, _, err := queueStartBudget(ctx, tx, ns, queue, now) if err != nil || used >= rate { return false, err } var reserved uint64 err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM pulse_tasks WHERE namespace_id=? AND queue_id=? AND state='LEASED'`, ns, queue).Scan(&reserved) return err == nil && reserved < uint64(rate-used), err } func consumeQueueStart(ctx context.Context, tx *sql.Tx, ns, queue core.ID, now time.Time) error { rate, used, until, err := queueStartBudget(ctx, tx, ns, queue, now) if err != nil { return err } if used >= rate { return worker.RateLimited } if until == nil { end := now.Add(time.Second) until = &end } if err := consumeApplicationStart(ctx, tx, ns, now); err != nil { return err } r, err := tx.ExecContext(ctx, `UPDATE pulse_queues SET start_window_until=?,starts_in_window=? WHERE namespace_id=? AND id=?`, until, used+1, ns, queue) return exactlyOne(r, err) } // ConfigureQueue 完整替换控制参数,运行计数不参与配置版本,也不因配置变更重置。 func (s *TaskStore) ConfigureQueue(ctx context.Context, token auth.Token, id core.ID, key string, in service.QueueConfigInput) (service.QueueResult, error) { var out service.QueueResult err := s.transaction(ctx, token, "queues:manage", func(tx *sql.Tx, p auth.Permissions) error { if !p.AllowsQueue(id) { return fault.NotFound } identity, cached, err := reserveIdentity(ctx, tx, token.NamespaceID, "queue.configure", string(id), key, "queue", id, in) if err != nil { return err } if cached != nil { if err := json.Unmarshal(cached, &out); err != nil { return err } out.Deduplicated = true return nil } if _, _, err := lockQueue(ctx, tx, token.NamespaceID, id); err != nil { return err } before, err := readQueue(ctx, tx, token.NamespaceID, id) if err != nil { return err } if before.Revision != in.ExpectedRevision { return fault.Conflict } now, err := dbNow(ctx, tx) if err != nil { return err } r, err := tx.ExecContext(ctx, `UPDATE pulse_queues SET paused=?,max_concurrency=?,rate_per_second=?,max_unfinished_tasks=?,revision=revision+1,updated_at=? WHERE namespace_id=? AND id=? AND revision=?`, in.Paused, in.MaxConcurrency, in.RatePerSecond, in.MaxUnfinishedTasks, now, token.NamespaceID, id, in.ExpectedRevision) if err := exactlyOne(r, err); err != nil { return err } if _, _, err := queueLimit(ctx, tx, token.NamespaceID, id, in.MaxConcurrency); err != nil { return err } out.Queue, err = readQueue(ctx, tx, token.NamespaceID, id) if err != nil { return err } if err := appendAudit(ctx, tx, token, identity, id, "queue", "queue.configure", in.Reason, before, out.Queue); err != nil { return err } if err := emitEvent(ctx, tx, token.NamespaceID, id, out.Queue.Revision, "queue", "queue.configured", out.Queue); err != nil { return err } return saveIdentity(ctx, tx, token.NamespaceID, identity, out) }) return out, err }