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

package redisactivity

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	healthDomain "proxy-pool/internal/domain/health"
	proxyDomain "proxy-pool/internal/domain/proxy"
)

const healthTaskRecordVersion = 1

var _ healthDomain.TaskBroker = (*Adapter)(nil)

type healthTaskRecord struct {
	Version          int    `json:"version"`
	TaskID           string `json:"taskId"`
	Reference        string `json:"reference"`
	DueMember        string `json:"dueMember"`
	ProxyID          string `json:"proxyId"`
	UpstreamID       string `json:"upstreamId"`
	UpstreamTasksKey string `json:"upstreamTasksKey"`
	Level            string `json:"level"`
	Priority         int    `json:"priority"`
	DeadlineMS       int64  `json:"deadlineMs"`
	NextDueMS        int64  `json:"nextDueMs"`
	Attempts         int    `json:"attempts"`
	MaxInFlight      int    `json:"maxInFlight"`
	State            string `json:"state"`
	RoutingName      string `json:"routingName,omitempty"`
	TargetURL        string `json:"targetUrl,omitempty"`
	LeaseCheckerID   string `json:"leaseCheckerId,omitempty"`
	LeaseInstanceID  string `json:"leaseInstanceId,omitempty"`
	LeaseToken       string `json:"leaseToken,omitempty"`
	LeaseExpiresAtMS int64  `json:"leaseExpiresAtMs,omitempty"`
	CheckerLeaseKey  string `json:"checkerLeaseKey,omitempty"`
}

type healthTaskRequest struct {
	Limit            int                  `json:"limit"`
	Tasks            []healthTaskRecord   `json:"tasks,omitempty"`
	CheckerID        string               `json:"checkerId,omitempty"`
	InstanceID       string               `json:"instanceId,omitempty"`
	MaxInFlight      int                  `json:"maxInFlight,omitempty"`
	LeaseTTLMS       int64                `json:"leaseTTLMS,omitempty"`
	Levels           []healthDomain.Level `json:"levels,omitempty"`
	UpstreamID       string               `json:"upstreamId,omitempty"`
	Level            string               `json:"level,omitempty"`
	RoutingName      string               `json:"routingName,omitempty"`
	TargetURL        string               `json:"targetUrl,omitempty"`
	UpstreamTasksKey string               `json:"upstreamTasksKey,omitempty"`
	ScanLimit        int                  `json:"scanLimit,omitempty"`
	Tokens           []string             `json:"tokens,omitempty"`
	Fact             *healthTaskFact      `json:"fact,omitempty"`
}

type healthTaskFact struct {
	TaskID      string `json:"taskId"`
	ProxyID     string `json:"proxyId"`
	Level       string `json:"level"`
	RoutingName string `json:"routingName,omitempty"`
	TargetURL   string `json:"targetUrl,omitempty"`
	CheckerID   string `json:"checkerId"`
	LeaseToken  string `json:"leaseToken"`
}

// InFlight returns the total queued and leased BASIC tasks after bounded
// maintenance. The Scheduler uses it only outside the Gateway hot path.
func (a *Adapter) InFlight(ctx context.Context, now time.Time) (int, error) {
	if a == nil || ctx == nil || now.IsZero() {
		return 0, healthDomain.ErrInvalidTaskBroker
	}
	reply, err := a.runHealthTaskScript(ctx, "inflight", now, healthTaskRequest{Limit: a.options.MaxCheckTasks}, "runtime")
	if err != nil {
		return 0, err
	}
	if reply.Status != scriptOK || reply.Count < 0 {
		return 0, invalidScriptReply("invalid health task inflight reply")
	}
	return reply.Count, nil
}

// DueCandidates returns a bounded BASIC due batch. Stale references are
// discarded inside the Lua script before they can reach the Controller.
func (a *Adapter) DueCandidates(ctx context.Context, now time.Time, limit int) ([]healthDomain.Candidate, error) {
	return a.dueCandidates(ctx, "", healthDomain.LevelBasic, "", "", now, limit)
}

func (a *Adapter) DueCandidatesForUpstream(ctx context.Context, upstreamID string, now time.Time, limit int) ([]healthDomain.Candidate, error) {
	if !validHealthTaskIdentifier(upstreamID) {
		return nil, healthDomain.ErrInvalidTaskBroker
	}
	return a.dueCandidates(ctx, upstreamID, healthDomain.LevelBasic, "", "", now, limit)
}

// DueEgressCandidatesForUpstream lazily initializes bounded EGRESS references
// from live BASIC references. The independent reference lets BASIC and EGRESS
// checks for one proxy be leased at the same time.
func (a *Adapter) DueEgressCandidatesForUpstream(ctx context.Context, upstreamID, targetURL string, now time.Time, limit int) ([]healthDomain.Candidate, error) {
	if !validHealthTaskIdentifier(upstreamID) {
		return nil, healthDomain.ErrInvalidTaskBroker
	}
	normalizedTarget, err := healthDomain.NormalizeEgressTarget(targetURL)
	if err != nil || normalizedTarget != targetURL {
		return nil, healthDomain.ErrInvalidTaskBroker
	}
	return a.dueCandidates(ctx, upstreamID, healthDomain.LevelEgress, targetURL, "", now, limit)
}

// DueTargetCandidatesForUpstream lazily initializes bounded TARGET references
// from live BASIC due items. TARGET health remains isolated by the configured
// Routing name and URL, so one target failure cannot change global proxy state.
func (a *Adapter) DueTargetCandidatesForUpstream(
	ctx context.Context,
	upstreamID string,
	routingName string,
	targetURL string,
	now time.Time,
	limit int,
) ([]healthDomain.Candidate, error) {
	if !validHealthTaskIdentifier(upstreamID) {
		return nil, healthDomain.ErrInvalidTaskBroker
	}
	profile, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{
		RoutingName: routingName, TargetURL: targetURL,
	})
	if err != nil || profile.RoutingName != routingName || profile.TargetURL != targetURL {
		return nil, healthDomain.ErrInvalidTaskBroker
	}
	return a.dueCandidates(ctx, upstreamID, healthDomain.LevelTarget, targetURL, routingName, now, limit)
}

func (a *Adapter) dueCandidates(
	ctx context.Context,
	upstreamID string,
	level healthDomain.Level,
	targetURL string,
	routingName string,
	now time.Time,
	limit int,
) ([]healthDomain.Candidate, error) {
	if a == nil || ctx == nil || now.IsZero() || limit <= 0 || limit > a.options.MaxCheckTasks {
		return nil, healthDomain.ErrInvalidTaskBroker
	}
	if level != healthDomain.LevelBasic && level != healthDomain.LevelEgress && level != healthDomain.LevelTarget {
		return nil, healthDomain.ErrInvalidTaskBroker
	}
	reply, err := a.runHealthTaskScript(ctx, "due", now, healthTaskRequest{
		Limit: limit, ScanLimit: a.options.MaxCandidateScan, UpstreamID: upstreamID, Level: string(level),
		RoutingName: routingName, TargetURL: targetURL,
	}, "runtime")
	if err != nil {
		return nil, err
	}
	if reply.Status != scriptOK || reply.CandidatesJSON == "" {
		return nil, invalidScriptReply("invalid health task due reply")
	}
	var candidates []healthTaskCandidateWire
	if err := decodeJSON(reply.CandidatesJSON, &candidates); err != nil || len(candidates) > limit {
		return nil, invalidScriptReply("invalid health task due candidate payload")
	}
	result := make([]healthDomain.Candidate, 0, len(candidates))
	for _, candidate := range candidates {
		state := proxyDomain.State(candidate.State)
		switch state {
		case proxyDomain.StateFetched, proxyDomain.StateChecking, proxyDomain.StateAvailable, proxyDomain.StateSuspect, proxyDomain.StateUnhealthy:
		default:
			return nil, invalidScriptReply("health task due reply contains invalid state")
		}
		if !validHealthTaskIdentifier(candidate.ProxyID) || !validHealthTaskIdentifier(candidate.UpstreamID) || candidate.DueAtMS <= 0 ||
			healthDomain.Level(candidate.Level) != level {
			return nil, invalidScriptReply("health task due reply contains invalid candidate")
		}
		switch level {
		case healthDomain.LevelEgress:
			normalizedTarget, err := healthDomain.NormalizeEgressTarget(candidate.TargetURL)
			if err != nil || normalizedTarget != targetURL {
				return nil, invalidScriptReply("health task due reply contains invalid EGRESS target")
			}
		case healthDomain.LevelTarget:
			profile, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{
				RoutingName: candidate.RoutingName, TargetURL: candidate.TargetURL,
			})
			if err != nil || profile.RoutingName != routingName || profile.TargetURL != targetURL {
				return nil, invalidScriptReply("health task due reply contains invalid TARGET profile")
			}
		}
		result = append(result, healthDomain.Candidate{
			ProxyID: candidate.ProxyID, UpstreamID: candidate.UpstreamID, State: state, Level: level,
			RoutingName: candidate.RoutingName, TargetURL: candidate.TargetURL,
			DueAt: time.UnixMilli(candidate.DueAtMS).UTC(),
		})
	}
	return result, nil
}

func (a *Adapter) InFlightForUpstream(ctx context.Context, upstreamID string, now time.Time) (int, error) {
	if a == nil || ctx == nil || now.IsZero() || !validHealthTaskIdentifier(upstreamID) {
		return 0, healthDomain.ErrInvalidTaskBroker
	}
	reply, err := a.runHealthTaskScript(ctx, "upstream_inflight", now, healthTaskRequest{
		Limit: a.options.MaxCheckTasks, UpstreamID: upstreamID, UpstreamTasksKey: a.keys.upstreamTasks(upstreamID),
	}, "runtime")
	if err != nil {
		return 0, err
	}
	if reply.Status != scriptOK || reply.Count < 0 {
		return 0, invalidScriptReply("invalid upstream health task inflight reply")
	}
	return reply.Count, nil
}

// Offer atomically turns due BASIC plans into shared queued tasks. A plan that
// loses the race leaves no duplicate task and no in-memory queue behind.
func (a *Adapter) Offer(ctx context.Context, plans []healthDomain.PlannedTask) (int, error) {
	if a == nil || ctx == nil {
		return 0, healthDomain.ErrInvalidTaskBroker
	}
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	if len(plans) == 0 {
		return 0, nil
	}
	if len(plans) > a.options.MaxCheckTasks {
		return 0, healthDomain.ErrInvalidLeasedTask
	}
	now := time.Now().UTC()
	tasks := make([]healthTaskRecord, len(plans))
	for index, plan := range plans {
		if err := validateRedisPlannedTask(plan, now); err != nil {
			return 0, err
		}
		tasks[index] = healthTaskRecord{
			Version: healthTaskRecordVersion, TaskID: healthDomain.TaskIDFor(plan), Reference: redisHealthTaskReference(plan.Candidate),
			DueMember: redisHealthDueMember(plan.Candidate), ProxyID: plan.Candidate.ProxyID,
			UpstreamID: plan.Candidate.UpstreamID, UpstreamTasksKey: a.keys.upstreamTasks(plan.Candidate.UpstreamID),
			Level: string(plan.Candidate.Level), Priority: int(plan.Priority), DeadlineMS: plan.Deadline.UnixMilli(),
			NextDueMS: plan.NextDue.UnixMilli(), Attempts: plan.Attempts, MaxInFlight: plan.MaxInFlight, State: "QUEUED",
			RoutingName: plan.Candidate.RoutingName, TargetURL: plan.Candidate.TargetURL,
		}
	}
	reply, err := a.runHealthTaskScript(ctx, "offer", now, healthTaskRequest{Limit: a.options.MaxCheckTasks, Tasks: tasks}, "runtime")
	if err != nil {
		return 0, err
	}
	if reply.Status != scriptOK || reply.Count < 0 || reply.Count > len(plans) {
		return 0, invalidScriptReply("invalid health task offer reply")
	}
	return reply.Count, nil
}

// Claim leases a bounded task batch. Proxy credentials are read only from the
// live record at lease time and are never copied into Redis task state.
func (a *Adapter) Claim(ctx context.Context, claim healthDomain.TaskClaim) ([]healthDomain.LeasedTask, error) {
	if a == nil || ctx == nil {
		return nil, healthDomain.ErrInvalidTaskBroker
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if err := validateRedisTaskClaim(claim, a.options.MaxCheckTasks); err != nil {
		return nil, err
	}
	tokens := make([]string, claim.MaxInFlight)
	for index := range tokens {
		token, err := newRedisLeaseToken()
		if err != nil {
			return nil, err
		}
		tokens[index] = token
	}
	now := time.Now().UTC()
	reply, err := a.runHealthTaskScript(ctx, "claim", now, healthTaskRequest{
		Limit: a.options.MaxCheckTasks, CheckerID: claim.CheckerID, InstanceID: claim.InstanceID,
		MaxInFlight: claim.MaxInFlight, LeaseTTLMS: a.options.CheckLeaseTTL.Milliseconds(),
		Levels: claim.SupportedLevels, Tokens: tokens,
	}, claim.CheckerID)
	if err != nil {
		return nil, err
	}
	if reply.Status != scriptOK || len(reply.Tasks) > claim.MaxInFlight {
		return nil, invalidScriptReply("invalid health task claim reply")
	}
	result := make([]healthDomain.LeasedTask, 0, len(reply.Tasks))
	for _, item := range reply.Tasks {
		task, err := decodeHealthTaskRecord(item.Task)
		if err != nil {
			return nil, invalidScriptReply("health task claim reply contains invalid task")
		}
		record, err := decodeProxyRecord(item.Record)
		if err != nil {
			return nil, invalidScriptReply("health task claim reply contains invalid record")
		}
		if task.State != "LEASED" || task.LeaseCheckerID != claim.CheckerID || task.LeaseToken == "" ||
			task.ProxyID != record.ID || task.DeadlineMS <= now.UnixMilli() || !validRedisTaskExecution(task) {
			return nil, invalidScriptReply("health task claim reply violates lease contract")
		}
		result = append(result, healthDomain.LeasedTask{
			TaskID: task.TaskID, LeaseToken: task.LeaseToken, ProxyID: task.ProxyID,
			Protocol: proxyDomain.Scheme(record.Scheme), Host: record.Host, Port: uint16(record.Port),
			Username: record.Username, Password: record.Password, Level: healthDomain.Level(task.Level),
			RoutingName: task.RoutingName, TargetURL: task.TargetURL,
			Deadline: time.UnixMilli(task.DeadlineMS).UTC(), Attempts: task.Attempts,
		})
	}
	return result, nil
}

func (a *Adapter) AuthorizeObservation(
	ctx context.Context,
	checkerID string,
	leaseToken string,
	observation healthDomain.Observation,
	now time.Time,
) error {
	return a.authorizeHealthTask(ctx, "authorize", checkerID, leaseToken, observation, now)
}

func (a *Adapter) CompleteObservation(
	ctx context.Context,
	checkerID string,
	leaseToken string,
	observation healthDomain.Observation,
	now time.Time,
) error {
	return a.authorizeHealthTask(ctx, "complete", checkerID, leaseToken, observation, now)
}

func (a *Adapter) authorizeHealthTask(
	ctx context.Context,
	operation string,
	checkerID string,
	leaseToken string,
	observation healthDomain.Observation,
	now time.Time,
) error {
	if a == nil || ctx == nil || !validHealthTaskIdentifier(checkerID) || !validHealthTaskIdentifier(leaseToken) || now.IsZero() {
		return healthDomain.ErrInvalidTaskBroker
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	normalized, err := healthDomain.NormalizeObservation(observation)
	if err != nil {
		return err
	}
	reply, err := a.runHealthTaskScript(ctx, operation, now, healthTaskRequest{
		Limit: a.options.MaxCheckTasks,
		Fact: &healthTaskFact{TaskID: normalized.TaskID, ProxyID: normalized.ProxyID, Level: string(normalized.Level),
			RoutingName: normalized.RoutingName, TargetURL: normalized.TargetURL, CheckerID: checkerID, LeaseToken: leaseToken},
	}, checkerID)
	if err != nil {
		return err
	}
	switch reply.Status {
	case scriptOK:
		return nil
	case scriptNotFound:
		return healthDomain.ErrTaskNotFound
	case scriptInvalid:
		return healthDomain.ErrInvalidLeasedTask
	case scriptStatus("lease_expired"):
		return healthDomain.ErrTaskLeaseExpired
	case scriptStatus("not_owned"):
		return healthDomain.ErrTaskLeaseNotOwned
	case scriptStatus("observation"):
		return healthDomain.ErrTaskObservation
	default:
		return invalidScriptReply("unexpected health task authorization status")
	}
}

func (a *Adapter) runHealthTaskScript(
	ctx context.Context,
	operation string,
	now time.Time,
	request healthTaskRequest,
	checkerID string,
) (healthTaskScriptReply, error) {
	if a == nil || ctx == nil || now.IsZero() || request.Limit <= 0 || strings.TrimSpace(operation) != operation || operation == "" {
		return healthTaskScriptReply{}, healthDomain.ErrInvalidTaskBroker
	}
	payload, err := json.Marshal(request)
	if err != nil {
		return healthTaskScriptReply{}, fmt.Errorf("encode health task request: %w", err)
	}
	result, err := runScript(ctx, a.client, healthTasksScript, []string{
		a.keys.records, a.keys.expiry, a.keys.healthDue, a.keys.healthQueued, a.keys.healthLeases,
		a.keys.healthTasks, a.keys.healthTaskExpiry, a.keys.healthRefTask, a.keys.checkerLeases(checkerID),
		a.keys.stateInventory, a.keys.healthEgressDue, a.keys.healthTargetDue,
	}, operation, now.UTC().UnixMilli(), string(payload))
	if err != nil {
		return healthTaskScriptReply{}, err
	}
	var reply healthTaskScriptReply
	if err := decodeScriptResult(result, &reply); err != nil {
		return healthTaskScriptReply{}, err
	}
	return reply, nil
}

func validateRedisPlannedTask(task healthDomain.PlannedTask, now time.Time) error {
	if !validHealthTaskIdentifier(task.Candidate.ProxyID) || !validHealthTaskIdentifier(task.Candidate.UpstreamID) || task.Candidate.DueAt.IsZero() ||
		task.Deadline.IsZero() || !task.Deadline.After(now) || task.NextDue.IsZero() || !task.NextDue.After(now) ||
		task.Attempts <= 0 || task.MaxInFlight <= 0 {
		return healthDomain.ErrInvalidLeasedTask
	}
	switch task.Candidate.Level {
	case healthDomain.LevelBasic:
		if task.Candidate.RoutingName != "" || task.Candidate.TargetURL != "" {
			return healthDomain.ErrInvalidLeasedTask
		}
	case healthDomain.LevelEgress:
		normalizedTarget, err := healthDomain.NormalizeEgressTarget(task.Candidate.TargetURL)
		if task.Candidate.RoutingName != "" || err != nil || normalizedTarget != task.Candidate.TargetURL {
			return healthDomain.ErrInvalidLeasedTask
		}
	case healthDomain.LevelTarget:
		profile, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{
			RoutingName: task.Candidate.RoutingName, TargetURL: task.Candidate.TargetURL,
		})
		if err != nil || profile.RoutingName != task.Candidate.RoutingName || profile.TargetURL != task.Candidate.TargetURL {
			return healthDomain.ErrInvalidLeasedTask
		}
	default:
		return healthDomain.ErrInvalidLeasedTask
	}
	switch task.Candidate.State {
	case proxyDomain.StateFetched, proxyDomain.StateAvailable, proxyDomain.StateSuspect, proxyDomain.StateUnhealthy:
		return nil
	default:
		return healthDomain.ErrInvalidLeasedTask
	}
}

func validateRedisTaskClaim(claim healthDomain.TaskClaim, maximum int) error {
	if !validHealthTaskIdentifier(claim.CheckerID) || !validHealthTaskIdentifier(claim.InstanceID) || claim.MaxInFlight <= 0 ||
		claim.MaxInFlight > maximum || len(claim.SupportedLevels) == 0 {
		return healthDomain.ErrInvalidTaskClaim
	}
	levels := make(map[healthDomain.Level]struct{}, len(claim.SupportedLevels))
	for _, level := range claim.SupportedLevels {
		switch level {
		case healthDomain.LevelBasic, healthDomain.LevelEgress, healthDomain.LevelTarget:
		default:
			return healthDomain.ErrInvalidTaskClaim
		}
		if _, exists := levels[level]; exists {
			return healthDomain.ErrInvalidTaskClaim
		}
		levels[level] = struct{}{}
	}
	return nil
}

func decodeHealthTaskRecord(payload string) (healthTaskRecord, error) {
	var record healthTaskRecord
	if err := decodeJSON(payload, &record); err != nil {
		return healthTaskRecord{}, err
	}
	if record.Level == string(healthDomain.LevelBasic) {
		if record.Reference == "" {
			record.Reference = record.ProxyID
		}
		if record.DueMember == "" {
			record.DueMember = record.ProxyID
		}
	}
	if record.Version != healthTaskRecordVersion || !validHealthTaskIdentifier(record.TaskID) ||
		!validHealthTaskIdentifier(record.ProxyID) || !validHealthTaskIdentifier(record.UpstreamID) || record.UpstreamTasksKey == "" ||
		!validHealthTaskReference(record.Reference) || !validHealthTaskReference(record.DueMember) ||
		record.Priority < int(healthDomain.PriorityFetched) || record.Priority > int(healthDomain.PriorityAvailable) ||
		record.DeadlineMS <= 0 || record.NextDueMS <= 0 || record.Attempts <= 0 || record.MaxInFlight <= 0 ||
		(record.State != "QUEUED" && record.State != "LEASED" && record.State != "DONE") {
		return healthTaskRecord{}, errors.New("invalid health task record")
	}
	return record, nil
}

func redisHealthTaskReference(candidate healthDomain.Candidate) string {
	if candidate.Level == healthDomain.LevelBasic {
		return candidate.ProxyID
	}
	return healthDomain.CandidateIdentity(candidate)
}

func redisHealthDueMember(candidate healthDomain.Candidate) string {
	switch candidate.Level {
	case healthDomain.LevelBasic:
		return candidate.ProxyID
	case healthDomain.LevelEgress:
		return candidate.TargetURL + "\x00" + candidate.ProxyID
	case healthDomain.LevelTarget:
		return candidate.RoutingName + "\x00" + candidate.TargetURL + "\x00" + candidate.ProxyID
	default:
		return ""
	}
}

func validRedisTaskExecution(task healthTaskRecord) bool {
	switch healthDomain.Level(task.Level) {
	case healthDomain.LevelBasic:
		return task.RoutingName == "" && task.TargetURL == ""
	case healthDomain.LevelEgress:
		targetURL, err := healthDomain.NormalizeEgressTarget(task.TargetURL)
		return task.RoutingName == "" && err == nil && targetURL == task.TargetURL
	case healthDomain.LevelTarget:
		profile, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{
			RoutingName: task.RoutingName, TargetURL: task.TargetURL,
		})
		return err == nil && profile.RoutingName == task.RoutingName && profile.TargetURL == task.TargetURL
	default:
		return false
	}
}

func validHealthTaskReference(value string) bool {
	return value != "" && len(value) <= 2048 && strings.TrimSpace(value) == value && !strings.ContainsRune(value, '\x7f')
}

func newRedisLeaseToken() (string, error) {
	var entropy [24]byte
	if _, err := rand.Read(entropy[:]); err != nil {
		return "", err
	}
	return "lease_" + hex.EncodeToString(entropy[:]), nil
}

func validHealthTaskIdentifier(value string) bool {
	if value == "" || len(value) > 256 || strings.TrimSpace(value) != value {
		return false
	}
	for _, character := range value {
		if character <= ' ' || character == '\x7f' {
			return false
		}
	}
	return true
}