orbit/internal/storage/mysql/queue_resources.go

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_tasks<max_unfinished_tasks`, ns, queue)
	if err != nil {
		return err
	}
	n, err := r.RowsAffected()
	if err != nil {
		return err
	}
	if n != 1 {
		return fault.QueueFull
	}
	r, err = tx.ExecContext(ctx, `UPDATE orbit_application_resources SET unfinished_tasks=unfinished_tasks+1 WHERE application_id=? AND unfinished_tasks<max_unfinished_tasks`, v.ApplicationID)
	return exactlyOne(r, err)
}

func releaseQueueTask(ctx context.Context, tx *sql.Tx, ns, queue core.ID) error {
	r, err := tx.ExecContext(ctx, `UPDATE orbit_application_resources SET unfinished_tasks=unfinished_tasks-1 WHERE application_id=`+applicationForNamespace+` AND unfinished_tasks>0`, 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
}