package redisactivity import ( "context" "encoding/json" "strconv" "time" "proxy-pool/internal/domain/activitypool" ) const ( maintenanceInventory = "inventory" maintenanceSweep = "sweep" maintenanceUnhealthy = "unhealthy" maintenanceDisabledUpstream = "disabled_upstream" maximumUnhealthySweepScan = 1024 ) var ( _ activitypool.InventoryReader = (*Adapter)(nil) _ activitypool.Maintainer = (*Adapter)(nil) _ activitypool.DisabledUpstreamDrainLister = (*Adapter)(nil) ) func (a *Adapter) Inventory(ctx context.Context, upstreamID string, now time.Time) (activitypool.Inventory, error) { result := activitypool.Inventory{UpstreamID: upstreamID} if ctx == nil { return result, activitypool.ErrInvalidInventory } if err := ctx.Err(); err != nil { return result, err } if a == nil || upstreamID == "" || now.IsZero() { return result, activitypool.ErrInvalidInventory } reply, err := a.runMaintenance(ctx, maintenanceInventory, now, a.options.CleanupLimit, upstreamID, "") if err != nil { return result, err } if reply.Status == scriptInvalid { return result, activitypool.ErrInvalidInventory } if reply.Status != scriptOK || reply.Count < 0 { return result, invalidScriptReply("unexpected inventory reply") } result.Managed = reply.Count return result, nil } func (a *Adapter) SweepExpired(ctx context.Context, now time.Time, limit int) (int, error) { if ctx == nil { return 0, activitypool.ErrInvalidMaintenance } if err := ctx.Err(); err != nil { return 0, err } if a == nil || now.IsZero() || limit <= 0 { return 0, activitypool.ErrInvalidMaintenance } reply, err := a.runMaintenance(ctx, maintenanceSweep, now, limit, "", "") if err != nil { return 0, err } if reply.Status == scriptInvalid { return 0, activitypool.ErrInvalidMaintenance } if reply.Status != scriptOK || reply.Count < 0 || reply.Count > limit { return 0, invalidScriptReply("unexpected expiry sweep reply") } return reply.Count, nil } func (a *Adapter) SweepUnhealthy( ctx context.Context, command activitypool.UnhealthySweepCommand, ) (activitypool.UnhealthySweepResult, error) { if ctx == nil { return activitypool.UnhealthySweepResult{}, activitypool.ErrInvalidMaintenance } if err := ctx.Err(); err != nil { return activitypool.UnhealthySweepResult{}, err } if a == nil || command.Now.IsZero() || command.Limit <= 0 || len(command.RemoveAfterByUpstream) == 0 { return activitypool.UnhealthySweepResult{}, activitypool.ErrInvalidMaintenance } policies := make(map[string]int64, len(command.RemoveAfterByUpstream)) for upstreamID, removeAfter := range command.RemoveAfterByUpstream { if upstreamID == "" || removeAfter <= 0 { return activitypool.UnhealthySweepResult{}, activitypool.ErrInvalidMaintenance } policies[upstreamID] = removeAfter.Milliseconds() if policies[upstreamID] <= 0 { return activitypool.UnhealthySweepResult{}, activitypool.ErrInvalidMaintenance } } payload, err := json.Marshal(policies) if err != nil { return activitypool.UnhealthySweepResult{}, err } limit := command.Limit if limit > maximumUnhealthySweepScan/4 { limit = maximumUnhealthySweepScan / 4 } reply, err := a.runMaintenance(ctx, maintenanceUnhealthy, command.Now, limit, "", string(payload)) if err != nil { return activitypool.UnhealthySweepResult{}, err } if reply.Status == scriptInvalid { return activitypool.UnhealthySweepResult{}, activitypool.ErrInvalidMaintenance } if reply.Status != scriptOK || reply.Count < 0 || reply.Count > limit || reply.DeferredOwned < 0 || reply.DeferredOwned > maximumUnhealthySweepScan || len(reply.DrainCandidates) > limit { return activitypool.UnhealthySweepResult{}, invalidScriptReply("unexpected unhealthy sweep reply") } result := activitypool.UnhealthySweepResult{Removed: reply.Count, DeferredOwned: reply.DeferredOwned} for _, candidate := range reply.DrainCandidates { if candidate.ProxyID == "" || candidate.WorkerID == "" || candidate.AssignmentEpoch == 0 || candidate.UnhealthySinceMS <= 0 { return activitypool.UnhealthySweepResult{}, invalidScriptReply("unhealthy sweep reply contained an invalid drain candidate") } result.DrainCandidates = append(result.DrainCandidates, activitypool.UnhealthyDrainCandidate{ ProxyID: candidate.ProxyID, WorkerID: candidate.WorkerID, AssignmentEpoch: candidate.AssignmentEpoch, UnhealthySince: time.UnixMilli(candidate.UnhealthySinceMS).UTC(), }) } return result, nil } func (a *Adapter) ListDisabledUpstreamDrainCandidates( ctx context.Context, now time.Time, policy activitypool.UpstreamDrainPolicy, limit int, ) ([]activitypool.DisabledUpstreamDrainCandidate, error) { if ctx == nil { return nil, activitypool.ErrInvalidMaintenance } if err := ctx.Err(); err != nil { return nil, err } if a == nil || now.IsZero() || policy.UpstreamID == "" || policy.Revision == 0 || policy.Enabled || limit <= 0 { return nil, activitypool.ErrInvalidMaintenance } if limit > maximumUnhealthySweepScan/4 { limit = maximumUnhealthySweepScan / 4 } reply, err := a.runMaintenance(ctx, maintenanceDisabledUpstream, now, limit, policy.UpstreamID, "") if err != nil { return nil, err } if reply.Status == scriptStale { return nil, nil } if reply.Status == scriptInvalid { return nil, activitypool.ErrInvalidMaintenance } if reply.Status != scriptOK || len(reply.DisabledDrainCandidates) > limit { return nil, invalidScriptReply("unexpected disabled upstream sweep reply") } wantRevision := strconv.FormatUint(policy.Revision, 10) result := make([]activitypool.DisabledUpstreamDrainCandidate, 0, len(reply.DisabledDrainCandidates)) for _, candidate := range reply.DisabledDrainCandidates { if candidate.UpstreamID != policy.UpstreamID || candidate.PolicyRevision != wantRevision || candidate.ProxyID == "" || candidate.WorkerID == "" || candidate.AssignmentEpoch == 0 { return nil, invalidScriptReply("disabled upstream sweep reply contained an invalid drain candidate") } result = append(result, activitypool.DisabledUpstreamDrainCandidate{ UpstreamID: candidate.UpstreamID, PolicyRevision: policy.Revision, ProxyID: candidate.ProxyID, WorkerID: candidate.WorkerID, AssignmentEpoch: candidate.AssignmentEpoch, }) } return result, nil } func (a *Adapter) runMaintenance( ctx context.Context, operation string, now time.Time, limit int, upstreamID string, policies string, ) (maintenanceScriptReply, error) { operationID, err := newOperationID() if err != nil { return maintenanceScriptReply{}, err } result, err := runScript(ctx, a.client, sweepScript, []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.operation(operationID), a.keys.healthDue, a.keys.healthQueued, a.keys.healthLeases, a.keys.healthTasks, a.keys.healthTaskExpiry, a.keys.healthRefTask, a.keys.healthUnhealthy, a.keys.owned(upstreamID), a.keys.upstreamDrainPolicies, }, operation, now.UnixMilli(), limit, upstreamID, operationTTLMillis(a.options.OperationTTL), policies) if err != nil { return maintenanceScriptReply{}, err } var reply maintenanceScriptReply if err := decodeScriptResult(result, &reply); err != nil { return maintenanceScriptReply{}, err } return reply, nil }