212 lines
8.7 KiB
Go
212 lines
8.7 KiB
Go
package mysql
|
|
|
|
import (
|
|
"context"
|
|
"crypto/subtle"
|
|
"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/pulse/worker"
|
|
)
|
|
|
|
// ExecutionOptions 是服务端执行窗口，不由 Worker 任意指定。
|
|
type ExecutionOptions struct {
|
|
LeaseDuration time.Duration // LeaseDuration 每次领取/续租的最大窗口。
|
|
SessionTTL time.Duration // SessionTTL 会话心跳有效期，不代替执行租约。
|
|
}
|
|
|
|
// ExecutionStore 所有执行权以 MySQL 为准，节点内存不持有可绕过数据库的授权。
|
|
type ExecutionStore struct {
|
|
db *sql.DB
|
|
node core.ID
|
|
options ExecutionOptions
|
|
}
|
|
|
|
var _ worker.Store = (*ExecutionStore)(nil)
|
|
|
|
func (s *ExecutionStore) Authorize(ctx context.Context, token auth.Token) error {
|
|
return s.transaction(ctx, token, "workers:connect", func(*sql.Tx, auth.Permissions) error { return nil })
|
|
}
|
|
|
|
func NewExecutionStore(db *sql.DB, node core.ID, options ExecutionOptions) (*ExecutionStore, error) {
|
|
if options.LeaseDuration == 0 {
|
|
options.LeaseDuration = 30 * time.Second
|
|
}
|
|
if options.SessionTTL == 0 {
|
|
options.SessionTTL = 45 * time.Second
|
|
}
|
|
if node.Validate() != nil || options.LeaseDuration < time.Second || options.LeaseDuration > 2*time.Minute || options.SessionTTL < options.LeaseDuration || options.SessionTTL > 5*time.Minute {
|
|
return nil, fault.Invalid("invalid execution windows or node identity")
|
|
}
|
|
return &ExecutionStore{db: db, node: node, options: options}, nil
|
|
}
|
|
func (s *ExecutionStore) transaction(ctx context.Context, token auth.Token, action string, fn func(*sql.Tx, auth.Permissions) error) error {
|
|
return (&TaskStore{db: s.db}).transaction(ctx, token, action, fn)
|
|
}
|
|
|
|
type sessionState struct {
|
|
queue core.ID
|
|
capacity uint32
|
|
sequence uint64
|
|
draining bool
|
|
expires time.Time
|
|
capabilities []worker.Capability
|
|
}
|
|
|
|
func (s *ExecutionStore) welcome(b worker.Binding, state sessionState, now time.Time) worker.Welcome {
|
|
return worker.Welcome{SessionID: b.SessionID, Epoch: b.Epoch, ServerTime: now, ExpiresAt: state.expires, LeaseDurationMS: uint64(s.options.LeaseDuration / time.Millisecond), Draining: state.draining}
|
|
}
|
|
func (s *ExecutionStore) Connect(ctx context.Context, token auth.Token, in worker.ConnectInput) (worker.Binding, worker.Welcome, error) {
|
|
var binding worker.Binding
|
|
var out worker.Welcome
|
|
proof, err := worker.SecretDigest(in.InstanceSecret)
|
|
if err != nil {
|
|
return binding, out, err
|
|
}
|
|
registration, err := requestDigest(struct {
|
|
Queue core.ID
|
|
Capabilities []worker.Capability
|
|
}{in.QueueID, in.Capabilities})
|
|
if err != nil {
|
|
return binding, out, err
|
|
}
|
|
err = s.transaction(ctx, token, "workers:connect", func(tx *sql.Tx, p auth.Permissions) error {
|
|
if !p.AllowsQueue(in.QueueID) {
|
|
return fault.Forbidden
|
|
}
|
|
_, err := tx.ExecContext(ctx, `INSERT INTO pulse_worker_instances(namespace_id,id,credential_id,queue_id,secret_digest,registration_digest) VALUES (?,?,?,?,?,?) ON DUPLICATE KEY UPDATE id=pulse_worker_instances.id`, token.NamespaceID, in.InstanceID, token.CredentialID, in.QueueID, proof, registration)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var credential core.ID
|
|
var storedProof, storedRegistration core.Digest
|
|
var epoch uint64
|
|
var old *core.ID
|
|
var draining bool
|
|
err = tx.QueryRowContext(ctx, `SELECT credential_id,secret_digest,registration_digest,session_epoch,current_session_id,draining FROM pulse_worker_instances WHERE namespace_id=? AND id=? FOR UPDATE`, token.NamespaceID, in.InstanceID).Scan(&credential, &storedProof, &storedRegistration, &epoch, &old, &draining)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if credential != token.CredentialID || subtle.ConstantTimeCompare(proof[:], storedProof[:]) != 1 {
|
|
return fault.Forbidden
|
|
}
|
|
if registration != storedRegistration || epoch == ^uint64(0) {
|
|
return fault.Conflict
|
|
}
|
|
now, err := dbNow(ctx, tx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
session, err := core.NewID()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
binding = worker.Binding{InstanceID: in.InstanceID, SessionID: session, Epoch: epoch + 1, Proof: proof}
|
|
state := sessionState{queue: in.QueueID, capacity: in.Capacity, draining: draining, expires: now.Add(s.options.SessionTTL), capabilities: in.Capabilities}
|
|
caps, err := json.Marshal(in.Capabilities)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = tx.ExecContext(ctx, `INSERT INTO pulse_worker_sessions(namespace_id,id,worker_instance_id,credential_id,node_id,session_epoch,capabilities,capacity,draining,expires_at) VALUES (?,?,?,?,?,?,?,?,?,?)`, token.NamespaceID, session, in.InstanceID, token.CredentialID, s.node, binding.Epoch, string(caps), in.Capacity, draining, state.expires)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if old != nil {
|
|
if _, err := tx.ExecContext(ctx, `UPDATE pulse_worker_sessions SET expires_at=?,draining=TRUE,revision=revision+1,updated_at=? WHERE namespace_id=? AND id=?`, now, now, token.NamespaceID, *old); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
_, err = tx.ExecContext(ctx, `UPDATE pulse_worker_instances SET current_session_id=?,session_epoch=?,revision=revision+1,updated_at=? WHERE namespace_id=? AND id=?`, session, binding.Epoch, now, token.NamespaceID, in.InstanceID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := emitEvent(ctx, tx, token.NamespaceID, session, 1, "worker_session", "worker.connected", map[string]any{"instance_id": in.InstanceID, "node_id": s.node, "epoch": binding.Epoch}); err != nil {
|
|
return err
|
|
}
|
|
out = s.welcome(binding, state, now)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return worker.Binding{}, worker.Welcome{}, err
|
|
}
|
|
return binding, out, nil
|
|
}
|
|
|
|
// lockSession 的顺序为凭据 -> 进程头 -> 当前会话；新连接和旧连接命令在同一行上串行化。
|
|
func (s *ExecutionStore) lockSession(ctx context.Context, tx *sql.Tx, token auth.Token, b worker.Binding, p auth.Permissions) (sessionState, error) {
|
|
var state sessionState
|
|
var credential core.ID
|
|
var proof core.Digest
|
|
var current *core.ID
|
|
var epoch uint64
|
|
if b.InstanceID.Validate() != nil || b.SessionID.Validate() != nil || b.Epoch == 0 {
|
|
return state, worker.StaleSession
|
|
}
|
|
err := tx.QueryRowContext(ctx, `SELECT credential_id,queue_id,secret_digest,current_session_id,session_epoch,draining FROM pulse_worker_instances WHERE namespace_id=? AND id=? FOR UPDATE`, token.NamespaceID, b.InstanceID).Scan(&credential, &state.queue, &proof, ¤t, &epoch, &state.draining)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return state, worker.StaleSession
|
|
}
|
|
if err != nil {
|
|
return state, err
|
|
}
|
|
if credential != token.CredentialID || current == nil || *current != b.SessionID || epoch != b.Epoch || subtle.ConstantTimeCompare(proof[:], b.Proof[:]) != 1 {
|
|
return state, worker.StaleSession
|
|
}
|
|
if !p.AllowsQueue(state.queue) {
|
|
return state, fault.Forbidden
|
|
}
|
|
var node core.ID
|
|
var raw []byte
|
|
err = tx.QueryRowContext(ctx, `SELECT node_id,capacity,capacity_sequence,capabilities,expires_at FROM pulse_worker_sessions WHERE namespace_id=? AND id=? FOR UPDATE`, token.NamespaceID, b.SessionID).Scan(&node, &state.capacity, &state.sequence, &raw, &state.expires)
|
|
if err != nil {
|
|
return state, err
|
|
}
|
|
now, err := dbNow(ctx, tx)
|
|
if err != nil {
|
|
return state, err
|
|
}
|
|
if node != s.node || !state.expires.After(now) {
|
|
return state, worker.StaleSession
|
|
}
|
|
if err := json.Unmarshal(raw, &state.capabilities); err != nil {
|
|
return state, err
|
|
}
|
|
return state, nil
|
|
}
|
|
func (s *ExecutionStore) Capacity(ctx context.Context, token auth.Token, b worker.Binding, in worker.CapacityInput) (worker.Welcome, error) {
|
|
var out worker.Welcome
|
|
err := s.transaction(ctx, token, "workers:consume", func(tx *sql.Tx, p auth.Permissions) error {
|
|
state, err := s.lockSession(ctx, tx, token, b, p)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if in.Sequence < state.sequence || (in.Sequence == state.sequence && (in.Capacity != state.capacity || in.Drain != state.draining)) || (state.draining && !in.Drain) {
|
|
return fault.Conflict
|
|
}
|
|
now, err := dbNow(ctx, tx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
state.expires = now.Add(s.options.SessionTTL)
|
|
state.capacity = in.Capacity
|
|
state.draining = state.draining || in.Drain
|
|
_, err = tx.ExecContext(ctx, `UPDATE pulse_worker_sessions SET capacity=?,capacity_sequence=?,draining=?,expires_at=?,revision=revision+1,updated_at=? WHERE namespace_id=? AND id=?`, state.capacity, in.Sequence, state.draining, state.expires, now, token.NamespaceID, b.SessionID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if in.Drain {
|
|
if _, err := tx.ExecContext(ctx, `UPDATE pulse_worker_instances SET draining=TRUE,revision=revision+1,updated_at=? WHERE namespace_id=? AND id=? AND draining=FALSE`, now, token.NamespaceID, b.InstanceID); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
out = s.welcome(b, state, now)
|
|
return nil
|
|
})
|
|
return out, err
|
|
}
|