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

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
}