143 lines
6.5 KiB
Go
143 lines
6.5 KiB
Go
package mysql
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"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/platform/resources"
|
|
"git.yuntuo.hk/k2safe/orbit/pulse/worker"
|
|
)
|
|
|
|
// 归属不允许在线修改；认证已经锁定身份。恢复也需释放被停用应用的额度。
|
|
const applicationForNamespace = `(SELECT application_id FROM orbit_namespaces WHERE id=?)`
|
|
|
|
// 统一顺序：认证 → Worker 会话（如有）→ 应用额度 → 队列 → 任务/尝试 → 许可。
|
|
// 与身份行分离，避免认证共享锁升级造成的互相等待。
|
|
func lockApplicationResources(ctx context.Context, tx *sql.Tx, ns core.ID, skip bool) (bool, error) {
|
|
q := `SELECT application_id FROM orbit_application_resources WHERE application_id=` + applicationForNamespace + ` FOR UPDATE`
|
|
if skip {
|
|
q += ` SKIP LOCKED`
|
|
}
|
|
var id core.ID
|
|
err := tx.QueryRowContext(ctx, q, ns).Scan(&id)
|
|
if errors.Is(err, sql.ErrNoRows) && skip {
|
|
return false, nil
|
|
}
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return false, errors.New("application resources missing")
|
|
}
|
|
return err == nil, err
|
|
}
|
|
func readApplicationResources(ctx context.Context, tx *sql.Tx, ns core.ID) (resources.View, error) {
|
|
var v resources.View
|
|
err := tx.QueryRowContext(ctx, `SELECT application_id,revision,max_unfinished_tasks,unfinished_tasks,max_concurrency,active_attempts,pending_starts,rate_per_second FROM orbit_application_resources WHERE application_id=`+applicationForNamespace, ns).Scan(&v.ApplicationID, &v.Revision, &v.MaxUnfinishedTasks, &v.UnfinishedTasks, &v.MaxConcurrency, &v.ActiveAttempts, &v.PendingStarts, &v.RatePerSecond)
|
|
return v, err
|
|
}
|
|
func applicationStartBudget(ctx context.Context, tx *sql.Tx, ns core.ID, now time.Time) (rate, used uint32, pending, active uint64, max uint32, until *time.Time, err error) {
|
|
err = tx.QueryRowContext(ctx, `SELECT rate_per_second,starts_in_window,pending_starts,active_attempts,max_concurrency,start_window_until FROM orbit_application_resources WHERE application_id=`+applicationForNamespace, ns).Scan(&rate, &used, &pending, &active, &max, &until)
|
|
if until == nil || !until.After(now) {
|
|
used, until = 0, nil
|
|
}
|
|
return
|
|
}
|
|
func applicationCanClaim(ctx context.Context, tx *sql.Tx, ns core.ID, now time.Time) (bool, error) {
|
|
rate, used, pending, active, max, _, err := applicationStartBudget(ctx, tx, ns, now)
|
|
return err == nil && active < uint64(max) && used < rate && pending < uint64(rate-used), err
|
|
}
|
|
func reserveApplicationAttempt(ctx context.Context, tx *sql.Tx, ns core.ID) error {
|
|
r, err := tx.ExecContext(ctx, `UPDATE orbit_application_resources SET active_attempts=active_attempts+1,pending_starts=pending_starts+1 WHERE application_id=`+applicationForNamespace+` AND active_attempts<max_concurrency`, ns)
|
|
return exactlyOne(r, err)
|
|
}
|
|
func consumeApplicationStart(ctx context.Context, tx *sql.Tx, ns core.ID, now time.Time) error {
|
|
rate, used, _, _, _, until, err := applicationStartBudget(ctx, tx, ns, now)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if used >= rate {
|
|
return worker.RateLimited
|
|
}
|
|
if until == nil {
|
|
t := now.Add(time.Second)
|
|
until = &t
|
|
}
|
|
r, err := tx.ExecContext(ctx, `UPDATE orbit_application_resources SET starts_in_window=?,start_window_until=?,pending_starts=pending_starts-1 WHERE application_id=`+applicationForNamespace+` AND pending_starts>0`, used+1, until, ns)
|
|
return exactlyOne(r, err)
|
|
}
|
|
func releaseApplicationAttempt(ctx context.Context, tx *sql.Tx, ns, attempt core.ID) error {
|
|
var unstarted bool
|
|
if err := tx.QueryRowContext(ctx, `SELECT authorized_at IS NULL FROM pulse_attempts WHERE namespace_id=? AND id=?`, ns, attempt).Scan(&unstarted); err != nil {
|
|
return err
|
|
}
|
|
var pending uint64
|
|
if unstarted {
|
|
pending = 1
|
|
}
|
|
r, err := tx.ExecContext(ctx, `UPDATE orbit_application_resources SET active_attempts=active_attempts-1,pending_starts=pending_starts-? WHERE application_id=`+applicationForNamespace+` AND active_attempts>0 AND pending_starts>=?`, pending, ns, pending)
|
|
return exactlyOne(r, err)
|
|
}
|
|
|
|
var _ resources.Store = (*TaskStore)(nil)
|
|
|
|
func (s *TaskStore) GetApplicationResources(ctx context.Context, token auth.Token) (resources.View, error) {
|
|
var out resources.View
|
|
err := s.transaction(ctx, token, "application:resources:read", func(tx *sql.Tx, p auth.Permissions) error {
|
|
var err error
|
|
out, err = readApplicationResources(ctx, tx, token.NamespaceID)
|
|
return err
|
|
})
|
|
return out, err
|
|
}
|
|
func (s *TaskStore) ConfigureApplicationResources(ctx context.Context, token auth.Token, key string, in resources.ConfigureInput) (resources.Result, error) {
|
|
var out resources.Result
|
|
err := s.transaction(ctx, token, "application:resources:manage", func(tx *sql.Tx, p auth.Permissions) error {
|
|
// 应用级动作显式授权；不从 queues:* 推导跨命名空间权限。
|
|
v, err := readApplicationResources(ctx, tx, token.NamespaceID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
identity, cached, err := reserveIdentity(ctx, tx, token.NamespaceID, "application.resources.configure", string(v.ApplicationID), key, "application", v.ApplicationID, 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 := lockApplicationResources(ctx, tx, token.NamespaceID, false); err != nil {
|
|
return err
|
|
}
|
|
before, err := readApplicationResources(ctx, tx, token.NamespaceID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if before.Revision != in.ExpectedRevision {
|
|
return fault.Conflict
|
|
}
|
|
r, err := tx.ExecContext(ctx, `UPDATE orbit_application_resources SET max_unfinished_tasks=?,max_concurrency=?,rate_per_second=?,revision=revision+1,updated_at=UTC_TIMESTAMP(6) WHERE application_id=? AND revision=?`, in.MaxUnfinishedTasks, in.MaxConcurrency, in.RatePerSecond, before.ApplicationID, in.ExpectedRevision)
|
|
if err := exactlyOne(r, err); err != nil {
|
|
return err
|
|
}
|
|
out.Resources, err = readApplicationResources(ctx, tx, token.NamespaceID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := appendAudit(ctx, tx, token, identity, before.ApplicationID, "application", "application.resources.configure", in.Reason, before, out.Resources); err != nil {
|
|
return err
|
|
}
|
|
if err := emitEvent(ctx, tx, token.NamespaceID, before.ApplicationID, out.Resources.Revision, "application", "application.resources.configured", out.Resources); err != nil {
|
|
return err
|
|
}
|
|
return saveIdentity(ctx, tx, token.NamespaceID, identity, out)
|
|
})
|
|
return out, err
|
|
}
|