402 lines
13 KiB
Go
402 lines
13 KiB
Go
package redisactivity
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"proxy-pool/internal/domain/activitypool"
|
|
ownershipDomain "proxy-pool/internal/domain/ownership"
|
|
)
|
|
|
|
const (
|
|
ownershipAssign = "assign"
|
|
ownershipRenew = "renew"
|
|
ownershipBeginDrain = "begin_drain"
|
|
ownershipBeginUnhealthyDrain = "begin_unhealthy_drain"
|
|
ownershipAcknowledgeDrain = "acknowledge_drain"
|
|
ownershipGet = "get"
|
|
ownershipExpire = "expire"
|
|
)
|
|
|
|
var _ ownershipDomain.Repository = (*Adapter)(nil)
|
|
var _ ownershipDomain.DrainTicketStore = (*Adapter)(nil)
|
|
var _ activitypool.UnhealthyDrainStarter = (*Adapter)(nil)
|
|
|
|
func (a *Adapter) Assign(
|
|
ctx context.Context,
|
|
now time.Time,
|
|
proxyID string,
|
|
workerID string,
|
|
ttl time.Duration,
|
|
) (ownershipDomain.Assignment, error) {
|
|
if err := validateOwnershipCall(ctx, a); err != nil {
|
|
return ownershipDomain.Assignment{}, err
|
|
}
|
|
if now.IsZero() || proxyID == "" || workerID == "" || ttl <= 0 {
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
|
}
|
|
reply, err := a.runOwnership(ctx, ownershipAssign, true, now.UnixMilli(), proxyID, workerID, 0, durationMillis(ttl), 0, 0)
|
|
if err != nil {
|
|
return ownershipDomain.Assignment{}, err
|
|
}
|
|
switch reply.Status {
|
|
case scriptOK:
|
|
return decodeAssignmentReply(reply)
|
|
case scriptAlreadyOwned:
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrAlreadyOwned
|
|
case scriptUnavailable, scriptNotFound:
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrOwnershipUnavailable
|
|
case scriptInvalid:
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
|
default:
|
|
return ownershipDomain.Assignment{}, invalidScriptReply("unexpected ownership assign status")
|
|
}
|
|
}
|
|
|
|
func (a *Adapter) Renew(
|
|
ctx context.Context,
|
|
now time.Time,
|
|
proxyID string,
|
|
workerID string,
|
|
epoch uint64,
|
|
ttl time.Duration,
|
|
) (ownershipDomain.Assignment, error) {
|
|
if err := validateOwnershipCall(ctx, a); err != nil {
|
|
return ownershipDomain.Assignment{}, err
|
|
}
|
|
if now.IsZero() || proxyID == "" || workerID == "" || epoch == 0 || ttl <= 0 {
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
|
}
|
|
reply, err := a.runOwnership(ctx, ownershipRenew, true, now.UnixMilli(), proxyID, workerID, epoch, durationMillis(ttl), 0, 0)
|
|
if err != nil {
|
|
return ownershipDomain.Assignment{}, err
|
|
}
|
|
switch reply.Status {
|
|
case scriptOK:
|
|
return decodeAssignmentReply(reply)
|
|
case scriptStale, scriptNotFound, scriptUnavailable:
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
|
|
case scriptInvalid:
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
|
default:
|
|
return ownershipDomain.Assignment{}, invalidScriptReply("unexpected ownership renew status")
|
|
}
|
|
}
|
|
|
|
func (a *Adapter) BeginDrain(
|
|
ctx context.Context,
|
|
proxyID string,
|
|
workerID string,
|
|
epoch uint64,
|
|
) (ownershipDomain.Assignment, error) {
|
|
if err := validateOwnershipCall(ctx, a); err != nil {
|
|
return ownershipDomain.Assignment{}, err
|
|
}
|
|
if proxyID == "" || workerID == "" || epoch == 0 {
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
|
}
|
|
reply, err := a.runOwnership(ctx, ownershipBeginDrain, true, 0, proxyID, workerID, epoch, 0, 0, 0)
|
|
if err != nil {
|
|
return ownershipDomain.Assignment{}, err
|
|
}
|
|
switch reply.Status {
|
|
case scriptOK:
|
|
return decodeAssignmentReply(reply)
|
|
case scriptStale, scriptNotFound:
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
|
|
case scriptInvalid:
|
|
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
|
default:
|
|
return ownershipDomain.Assignment{}, invalidScriptReply("unexpected ownership drain status")
|
|
}
|
|
}
|
|
|
|
func (a *Adapter) BeginUnhealthyDrain(
|
|
ctx context.Context,
|
|
now time.Time,
|
|
candidate activitypool.UnhealthyDrainCandidate,
|
|
) (bool, error) {
|
|
if err := validateOwnershipCall(ctx, a); err != nil {
|
|
return false, err
|
|
}
|
|
if now.IsZero() || candidate.ProxyID == "" || candidate.WorkerID == "" || candidate.AssignmentEpoch == 0 || candidate.UnhealthySince.IsZero() {
|
|
return false, activitypool.ErrInvalidMaintenance
|
|
}
|
|
reply, err := a.runOwnership(
|
|
ctx, ownershipBeginUnhealthyDrain, true, now.UTC().UnixMilli(), candidate.ProxyID, candidate.WorkerID,
|
|
candidate.AssignmentEpoch, candidate.UnhealthySince.UTC().UnixMilli(), 0, 0,
|
|
)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
switch reply.Status {
|
|
case scriptOK:
|
|
if _, err := decodeAssignmentReply(reply); err != nil {
|
|
return false, err
|
|
}
|
|
return reply.Started, nil
|
|
case scriptStale, scriptNotFound:
|
|
return false, nil
|
|
case scriptInvalid:
|
|
return false, activitypool.ErrInvalidMaintenance
|
|
default:
|
|
return false, invalidScriptReply("unexpected unhealthy ownership drain status")
|
|
}
|
|
}
|
|
|
|
func (a *Adapter) PendingDrains(ctx context.Context, workerID string, limit int) ([]ownershipDomain.DrainTicket, error) {
|
|
if err := validateOwnershipCall(ctx, a); err != nil {
|
|
return nil, err
|
|
}
|
|
if workerID == "" || limit <= 0 {
|
|
return nil, ownershipDomain.ErrInvalidDrainTicket
|
|
}
|
|
result, err := runScript(ctx, a.client, drainTicketsScript, []string{
|
|
a.keys.drainTickets, a.keys.workerDraining(workerID), a.keys.owners,
|
|
}, workerID, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var reply drainTicketsScriptReply
|
|
if err := decodeScriptResult(result, &reply); err != nil {
|
|
return nil, err
|
|
}
|
|
if reply.Status == scriptInvalid {
|
|
return nil, ownershipDomain.ErrInvalidDrainTicket
|
|
}
|
|
if reply.Status != scriptOK {
|
|
return nil, invalidScriptReply("unexpected pending drain tickets reply")
|
|
}
|
|
tickets := make([]ownershipDomain.DrainTicket, 0, len(reply.Tickets))
|
|
for _, raw := range reply.Tickets {
|
|
var record drainTicketRecord
|
|
if err := decodeJSON(raw, &record); err != nil || validateDrainTicketRecord(record) != nil {
|
|
return nil, invalidScriptReply("pending drain tickets reply contained an invalid ticket")
|
|
}
|
|
ticket := ownershipDomain.DrainTicket{
|
|
ProxyID: record.ProxyID, WorkerID: record.WorkerID, AssignmentEpoch: record.AssignmentEpoch,
|
|
RequiredSnapshotEpoch: record.RequiredSnapshotEpoch,
|
|
}
|
|
if record.SnapshotVersion != 0 {
|
|
checksum, err := hex.DecodeString(record.SnapshotChecksum)
|
|
if err != nil || len(checksum) != sha256.Size {
|
|
return nil, invalidScriptReply("pending drain ticket barrier checksum is invalid")
|
|
}
|
|
copy(ticket.Barrier.Checksum[:], checksum)
|
|
ticket.Barrier = ownershipDomain.SnapshotBarrier{
|
|
SessionID: record.SessionID, Version: record.SnapshotVersion,
|
|
OwnershipEpoch: record.SnapshotOwnershipEpoch, Checksum: ticket.Barrier.Checksum,
|
|
}
|
|
}
|
|
tickets = append(tickets, ticket)
|
|
}
|
|
return tickets, nil
|
|
}
|
|
|
|
func (a *Adapter) BindDrainBarrier(ctx context.Context, ticket ownershipDomain.DrainTicket) error {
|
|
if err := validateOwnershipCall(ctx, a); err != nil {
|
|
return err
|
|
}
|
|
if !validDrainTicket(ticket) || !validDrainBarrier(ticket.Barrier) ||
|
|
ticket.Barrier.OwnershipEpoch < ticket.RequiredSnapshotEpoch {
|
|
return ownershipDomain.ErrInvalidDrainTicket
|
|
}
|
|
result, err := runScript(ctx, a.client, bindDrainTicketScript, []string{
|
|
a.keys.drainTickets, a.keys.owners, a.keys.workerDraining(ticket.WorkerID),
|
|
}, ticket.ProxyID, ticket.WorkerID, ticket.AssignmentEpoch, ticket.RequiredSnapshotEpoch,
|
|
ticket.Barrier.SessionID, ticket.Barrier.Version, ticket.Barrier.OwnershipEpoch, hex.EncodeToString(ticket.Barrier.Checksum[:]))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var reply ownershipScriptReply
|
|
if err := decodeScriptResult(result, &reply); err != nil {
|
|
return err
|
|
}
|
|
switch reply.Status {
|
|
case scriptOK:
|
|
return nil
|
|
case scriptStale, scriptNotFound:
|
|
return ownershipDomain.ErrStaleAssignment
|
|
case scriptInvalid:
|
|
return ownershipDomain.ErrInvalidDrainTicket
|
|
default:
|
|
return invalidScriptReply("unexpected drain ticket barrier reply")
|
|
}
|
|
}
|
|
|
|
func (a *Adapter) AcknowledgeDrain(
|
|
ctx context.Context,
|
|
proxyID string,
|
|
workerID string,
|
|
epoch uint64,
|
|
active int64,
|
|
reserved int64,
|
|
) error {
|
|
if err := validateOwnershipCall(ctx, a); err != nil {
|
|
return err
|
|
}
|
|
if proxyID == "" || workerID == "" || epoch == 0 || active < 0 || reserved < 0 {
|
|
return ownershipDomain.ErrInvalidOwnership
|
|
}
|
|
reply, err := a.runOwnership(ctx, ownershipAcknowledgeDrain, true, 0, proxyID, workerID, epoch, 0, active, reserved)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
switch reply.Status {
|
|
case scriptOK:
|
|
return nil
|
|
case scriptStale, scriptNotFound:
|
|
return ownershipDomain.ErrStaleAssignment
|
|
case scriptNotDraining:
|
|
return ownershipDomain.ErrNotDraining
|
|
case scriptDrainNotReady:
|
|
return ownershipDomain.ErrDrainNotReady
|
|
case scriptInvalid:
|
|
return ownershipDomain.ErrInvalidOwnership
|
|
default:
|
|
return invalidScriptReply("unexpected ownership acknowledge status")
|
|
}
|
|
}
|
|
|
|
func (a *Adapter) Get(ctx context.Context, proxyID string) (ownershipDomain.Assignment, bool, error) {
|
|
if err := validateOwnershipCall(ctx, a); err != nil {
|
|
return ownershipDomain.Assignment{}, false, err
|
|
}
|
|
if proxyID == "" {
|
|
return ownershipDomain.Assignment{}, false, ownershipDomain.ErrInvalidOwnership
|
|
}
|
|
reply, err := a.runOwnership(ctx, ownershipGet, false, 0, proxyID, "", 0, 0, 0, 0)
|
|
if err != nil {
|
|
return ownershipDomain.Assignment{}, false, err
|
|
}
|
|
switch reply.Status {
|
|
case scriptNotFound:
|
|
return ownershipDomain.Assignment{}, false, nil
|
|
case scriptOK:
|
|
assignment, err := decodeAssignmentReply(reply)
|
|
return assignment, err == nil, err
|
|
case scriptInvalid:
|
|
return ownershipDomain.Assignment{}, false, ownershipDomain.ErrInvalidOwnership
|
|
default:
|
|
return ownershipDomain.Assignment{}, false, invalidScriptReply("unexpected ownership get status")
|
|
}
|
|
}
|
|
|
|
func (a *Adapter) Expire(ctx context.Context, now time.Time, limit int) ([]ownershipDomain.Assignment, error) {
|
|
if err := validateOwnershipCall(ctx, a); err != nil {
|
|
return nil, err
|
|
}
|
|
if now.IsZero() || limit <= 0 {
|
|
return nil, ownershipDomain.ErrInvalidOwnership
|
|
}
|
|
reply, err := a.runOwnership(ctx, ownershipExpire, true, now.UnixMilli(), "", "", 0, int64(limit), 0, 0)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if reply.Status == scriptInvalid {
|
|
return nil, ownershipDomain.ErrInvalidOwnership
|
|
}
|
|
if reply.Status != scriptOK || reply.Record == "" {
|
|
return nil, invalidScriptReply("unexpected ownership expire reply")
|
|
}
|
|
var records []ownershipRecord
|
|
if err := decodeJSON(reply.Record, &records); err != nil {
|
|
return nil, errors.Join(invalidScriptReply("ownership expire reply contained invalid records"), err)
|
|
}
|
|
assignments := make([]ownershipDomain.Assignment, 0, len(records))
|
|
for _, record := range records {
|
|
if err := validateOwnershipRecord(record); err != nil {
|
|
return nil, errors.Join(invalidScriptReply("ownership expire reply contained invalid records"), err)
|
|
}
|
|
assignments = append(assignments, assignmentFromRecord(record))
|
|
}
|
|
if len(assignments) > limit {
|
|
return nil, invalidScriptReply("ownership expire reply exceeded limit")
|
|
}
|
|
return assignments, nil
|
|
}
|
|
|
|
func (a *Adapter) runOwnership(
|
|
ctx context.Context,
|
|
operation string,
|
|
mutating bool,
|
|
nowMS int64,
|
|
proxyID string,
|
|
workerID string,
|
|
epoch uint64,
|
|
value int64,
|
|
active int64,
|
|
reserved int64,
|
|
) (ownershipScriptReply, error) {
|
|
operationKey := a.keys.epoch
|
|
if mutating {
|
|
operationID, err := newOperationID()
|
|
if err != nil {
|
|
return ownershipScriptReply{}, err
|
|
}
|
|
operationKey = a.keys.operation(operationID)
|
|
}
|
|
result, err := runScript(ctx, a.client, ownershipScript, []string{
|
|
a.keys.records, a.keys.unique, a.keys.idkeys, a.keys.expiry, a.keys.available,
|
|
a.keys.inventory, a.keys.stateInventory, a.keys.owners, a.keys.ownerExpiry,
|
|
a.keys.epoch, operationKey, a.keys.healthDue, a.keys.healthQueued, a.keys.healthLeases,
|
|
a.keys.healthTasks, a.keys.healthTaskExpiry, a.keys.healthRefTask, a.keys.healthUnhealthy,
|
|
a.keys.drainTickets,
|
|
}, operation, operationTTLMillis(a.options.OperationTTL), a.options.CleanupLimit,
|
|
nowMS, proxyID, workerID, epoch, value, active, reserved, a.keys.workerOwned(workerID), a.keys.workerDraining(workerID))
|
|
if err != nil {
|
|
return ownershipScriptReply{}, err
|
|
}
|
|
var reply ownershipScriptReply
|
|
if err := decodeScriptResult(result, &reply); err != nil {
|
|
return ownershipScriptReply{}, err
|
|
}
|
|
return reply, nil
|
|
}
|
|
|
|
func decodeAssignmentReply(reply ownershipScriptReply) (ownershipDomain.Assignment, error) {
|
|
if reply.Record == "" {
|
|
return ownershipDomain.Assignment{}, invalidScriptReply("ownership reply omitted assignment")
|
|
}
|
|
record, err := decodeOwnershipRecord(reply.Record)
|
|
if err != nil {
|
|
return ownershipDomain.Assignment{}, errors.Join(
|
|
invalidScriptReply("ownership reply contained an invalid assignment"),
|
|
fmt.Errorf("decode ownership assignment: %w", err),
|
|
)
|
|
}
|
|
return assignmentFromRecord(record), nil
|
|
}
|
|
|
|
func assignmentFromRecord(record ownershipRecord) ownershipDomain.Assignment {
|
|
return ownershipDomain.Assignment{
|
|
ProxyID: record.ProxyID, WorkerID: record.WorkerID, Epoch: record.Epoch,
|
|
Version: record.AssignmentVersion, ExpiresAt: time.UnixMilli(record.ExpiresAtMS).UTC(),
|
|
Draining: record.Draining,
|
|
}
|
|
}
|
|
|
|
func validateOwnershipCall(ctx context.Context, adapter *Adapter) error {
|
|
if ctx == nil || adapter == nil {
|
|
return ownershipDomain.ErrInvalidOwnership
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validDrainTicket(ticket ownershipDomain.DrainTicket) bool {
|
|
return ticket.ProxyID != "" && ticket.WorkerID != "" && ticket.AssignmentEpoch > 0 &&
|
|
ticket.RequiredSnapshotEpoch > ticket.AssignmentEpoch
|
|
}
|
|
|
|
func validDrainBarrier(barrier ownershipDomain.SnapshotBarrier) bool {
|
|
return barrier.SessionID != "" && barrier.Version > 0 && barrier.OwnershipEpoch > 0 &&
|
|
barrier.Checksum != [32]byte{}
|
|
}
|