72 lines
2.6 KiB
Go
72 lines
2.6 KiB
Go
package mysql
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"database/sql"
|
||
"errors"
|
||
|
||
"git.yuntuo.hk/k2safe/orbit/internal/fault"
|
||
core "git.yuntuo.hk/k2safe/orbit/model"
|
||
)
|
||
|
||
func concurrencyDigest(key string) *core.Digest {
|
||
if key == "" {
|
||
return nil
|
||
}
|
||
digest := core.Digest(sha256.Sum256(append([]byte("orbit-concurrency-key-v1\x00"), []byte(key)...)))
|
||
return &digest
|
||
}
|
||
|
||
// concurrencyKeyValue 在 Go 中判断空字节串，避免 MySQL 文本排序规则将某些字符视为空。
|
||
func concurrencyKeyValue(key string) *string {
|
||
if key == "" {
|
||
return nil
|
||
}
|
||
return &key
|
||
}
|
||
|
||
// keyEligibility 在 LIMIT 前排除已占用资源，避免前 16 个同键任务挡住其他资源。
|
||
// 调用者已锁应用额度与队列；命名空间中的所有领取都遵守此顺序。
|
||
const keyEligibility = `(concurrency_key_digest IS NULL OR NOT EXISTS (
|
||
SELECT 1 FROM pulse_concurrency_limits k WHERE k.namespace_id=pulse_tasks.namespace_id
|
||
AND k.scope_key='key' AND k.key_digest=pulse_tasks.concurrency_key_digest AND k.used_permits>=k.max_permits))`
|
||
|
||
// reserveConcurrencyKey 为一个显式资源键预占单个许可。零 token 表示未启用或暂不可用。
|
||
// 序号使用资源行单调 revision；释放也递增，序号允许有间隔，禁止删除或重置资源行。
|
||
func reserveConcurrencyKey(ctx context.Context, tx *sql.Tx, ns core.ID, key string) (core.ID, uint64, error) {
|
||
if key == "" {
|
||
return "", 0, nil
|
||
}
|
||
id, err := core.NewID()
|
||
if err != nil {
|
||
return "", 0, err
|
||
}
|
||
digest := concurrencyDigest(key)
|
||
_, err = tx.ExecContext(ctx, `INSERT INTO pulse_concurrency_limits(namespace_id,id,scope_key,key_digest,max_permits) VALUES (?,?,'key',?,1) ON DUPLICATE KEY UPDATE id=pulse_concurrency_limits.id`, ns, id, digest)
|
||
if err != nil {
|
||
return "", 0, err
|
||
}
|
||
var used, max uint32
|
||
var revision uint64
|
||
err = tx.QueryRowContext(ctx, `SELECT id,used_permits,max_permits,revision FROM pulse_concurrency_limits WHERE namespace_id=? AND scope_key='key' AND key_digest=? FOR UPDATE`, ns, digest).Scan(&id, &used, &max, &revision)
|
||
if err != nil {
|
||
return "", 0, err
|
||
}
|
||
if max != 1 {
|
||
return "", 0, errors.New("key permit limit must be one")
|
||
}
|
||
if used != 0 {
|
||
return id, 0, nil
|
||
}
|
||
// 为当前授权及后续释放各保留一次递增空间，不能溢出后重用旧序号。
|
||
if revision > ^uint64(0)-2 {
|
||
return "", 0, fault.Conflict
|
||
}
|
||
r, err := tx.ExecContext(ctx, `UPDATE pulse_concurrency_limits SET used_permits=1,revision=revision+1,updated_at=UTC_TIMESTAMP(6) WHERE namespace_id=? AND id=? AND used_permits=0 AND revision=?`, ns, id, revision)
|
||
if err := exactlyOne(r, err); err != nil {
|
||
return "", 0, err
|
||
}
|
||
return id, revision + 1, nil
|
||
}
|