proxy-pool/internal/adapters/redisactivity/ownership.go

package redisactivity

import (
	"context"
	"errors"
	"fmt"
	"time"

	ownershipDomain "proxy-pool/internal/domain/ownership"
)

const (
	ownershipAssign           = "assign"
	ownershipRenew            = "renew"
	ownershipBeginDrain       = "begin_drain"
	ownershipAcknowledgeDrain = "acknowledge_drain"
	ownershipGet              = "get"
	ownershipExpire           = "expire"
)

var _ ownershipDomain.Repository = (*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) 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.owners, a.keys.ownerExpiry, a.keys.epoch, operationKey,
	}, operation, operationTTLMillis(a.options.OperationTTL), a.options.CleanupLimit,
		nowMS, proxyID, workerID, epoch, value, active, reserved)
	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
}