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.stateInventory, 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, a.keys.workerOwned(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 }