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

package redisactivity

import (
	"context"
	"crypto/rand"
	_ "embed"
	"encoding/hex"
	"errors"
	"fmt"
	"time"

	"github.com/redis/go-redis/v9"

	extractionDomain "proxy-pool/internal/domain/extraction"
)

type scriptStatus string

const (
	scriptOK                   scriptStatus = "ok"
	scriptInvalid              scriptStatus = "invalid"
	scriptNotFound             scriptStatus = "not_found"
	scriptConflict             scriptStatus = "conflict"
	scriptStale                scriptStatus = "stale"
	scriptUnavailable          scriptStatus = "unavailable"
	scriptInsufficient         scriptStatus = "insufficient"
	scriptAlreadyOwned         scriptStatus = "already_owned"
	scriptNotDraining          scriptStatus = "not_draining"
	scriptDrainNotReady        scriptStatus = "drain_not_ready"
	scriptSnapshotMismatch     scriptStatus = "snapshot_mismatch"
	scriptStaleAcknowledgement scriptStatus = "stale_acknowledgement"
)

type upsertScriptReply struct {
	Status    scriptStatus `json:"status"`
	Accepted  int          `json:"accepted"`
	Inserted  int          `json:"inserted"`
	Refreshed int          `json:"refreshed"`
	Dropped   int          `json:"dropped"`
}

type healthScriptReply struct {
	Status scriptStatus `json:"status"`
	Record string       `json:"record,omitempty"`
}

type healthTaskScriptReply struct {
	Status         scriptStatus          `json:"status"`
	Count          int                   `json:"count"`
	CandidatesJSON string                `json:"candidatesJSON"`
	Tasks          []healthTaskClaimWire `json:"tasks"`
}

type healthTaskCandidateWire struct {
	ProxyID     string `json:"proxyId"`
	UpstreamID  string `json:"upstreamId"`
	State       string `json:"state"`
	Level       string `json:"level"`
	RoutingName string `json:"routingName,omitempty"`
	TargetURL   string `json:"targetUrl,omitempty"`
	DueAtMS     int64  `json:"dueAtMs"`
}

type healthTaskClaimWire struct {
	Task   string `json:"task"`
	Record string `json:"record"`
}

type targetHealthScriptReply struct {
	Status scriptStatus `json:"status"`
	Target string       `json:"target,omitempty"`
}

type upstreamLookupScriptReply struct {
	Status   scriptStatus `json:"status"`
	Upstream string       `json:"upstream,omitempty"`
}

type extractScriptReply struct {
	Status        scriptStatus `json:"status"`
	RequestDigest string       `json:"requestDigest"`
	Record        string       `json:"record,omitempty"`
}

type ownershipScriptReply struct {
	Status  scriptStatus `json:"status"`
	Record  string       `json:"record,omitempty"`
	Started bool         `json:"started"`
}

type drainTicketsScriptReply struct {
	Status  scriptStatus `json:"status"`
	Tickets []string     `json:"tickets"`
}

type maintenanceScriptReply struct {
	Status          scriptStatus                         `json:"status"`
	Count           int                                  `json:"count"`
	DeferredOwned   int                                  `json:"deferredOwned"`
	DrainCandidates []maintenanceDrainCandidateScriptRow `json:"drainCandidates"`
}

type maintenanceDrainCandidateScriptRow struct {
	ProxyID          string `json:"proxyId"`
	WorkerID         string `json:"workerId"`
	AssignmentEpoch  uint64 `json:"assignmentEpoch"`
	UnhealthySinceMS int64  `json:"unhealthySinceMs"`
}

type statusScriptReply struct {
	Status      scriptStatus            `json:"status"`
	Inventories []statusScriptInventory `json:"inventories"`
}

type statusScriptInventory struct {
	UpstreamID string `json:"upstreamId"`
	Fetched    int64  `json:"fetched"`
	Checking   int64  `json:"checking"`
	Available  int64  `json:"available"`
	Suspect    int64  `json:"suspect"`
	Draining   int64  `json:"draining"`
	Unhealthy  int64  `json:"unhealthy"`
	Extracted  int64  `json:"extracted"`
}

type runtimeScriptReply struct {
	Status    scriptStatus          `json:"status"`
	Snapshots []runtimeSnapshotWire `json:"snapshots"`
	Record    string                `json:"record,omitempty"`
}

type capacityScriptReply struct {
	Status         scriptStatus `json:"status"`
	Managed        int          `json:"managed"`
	AvailableSlots int64        `json:"availableSlots,string"`
}

type workerSnapshotScriptReply struct {
	Status  scriptStatus              `json:"status"`
	Proxies []workerSnapshotProxyWire `json:"proxies"`
}

type workerSnapshotProxyWire struct {
	Record           string `json:"record"`
	OwnershipEpoch   string `json:"ownershipEpoch"`
	LeaseExpiresAtMS int64  `json:"leaseExpiresAtMs"`
}

//go:embed scripts/upsert.lua
var upsertSource string

//go:embed scripts/health.lua
var healthSource string

//go:embed scripts/health_tasks.lua
var healthTasksSource string

//go:embed scripts/target_health.lua
var targetHealthSource string

//go:embed scripts/upstream_lookup.lua
var upstreamLookupSource string

//go:embed scripts/extract.lua
var extractSource string

//go:embed scripts/ownership.lua
var ownershipSource string

//go:embed scripts/drain_tickets.lua
var drainTicketsSource string

//go:embed scripts/bind_drain_ticket.lua
var bindDrainTicketSource string

//go:embed scripts/sweep.lua
var sweepSource string

//go:embed scripts/status.lua
var statusSource string

//go:embed scripts/runtime.lua
var runtimeSource string

//go:embed scripts/capacity.lua
var capacitySource string

//go:embed scripts/worker_snapshot.lua
var workerSnapshotSource string

var (
	upsertScript          = redis.NewScript(upsertSource)
	healthScript          = redis.NewScript(healthSource)
	healthTasksScript     = redis.NewScript(healthTasksSource)
	targetHealthScript    = redis.NewScript(targetHealthSource)
	upstreamLookupScript  = redis.NewScript(upstreamLookupSource)
	extractScript         = redis.NewScript(extractSource)
	ownershipScript       = redis.NewScript(ownershipSource)
	drainTicketsScript    = redis.NewScript(drainTicketsSource)
	bindDrainTicketScript = redis.NewScript(bindDrainTicketSource)
	sweepScript           = redis.NewScript(sweepSource)
	statusScript          = redis.NewScript(statusSource)
	runtimeScript         = redis.NewScript(runtimeSource)
	capacityScript        = redis.NewScript(capacitySource)
	workerSnapshotScript  = redis.NewScript(workerSnapshotSource)
)

func runScript(ctx context.Context, client redis.Scripter, script *redis.Script, keys []string, args ...any) (any, error) {
	result, err := script.Run(ctx, client, keys, args...).Result()
	if err != nil {
		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
			return nil, err
		}
		return nil, errors.Join(
			extractionDomain.ErrStoreUnavailable,
			fmt.Errorf("run redis activity script: %w", err),
		)
	}
	return result, nil
}

func newOperationID() (string, error) {
	var value [16]byte
	if _, err := rand.Read(value[:]); err != nil {
		return "", fmt.Errorf("create redis activity operation ID: %w", err)
	}
	return hex.EncodeToString(value[:]), nil
}

func operationTTLMillis(ttl time.Duration) int64 {
	milliseconds := ttl / time.Millisecond
	if ttl%time.Millisecond != 0 {
		milliseconds++
	}
	return int64(milliseconds)
}

func decodeScriptResult(result any, destination any) error {
	var payload string
	switch value := result.(type) {
	case string:
		payload = value
	case []byte:
		payload = string(value)
	default:
		return errors.Join(extractionDomain.ErrStoreUnavailable, errors.New("invalid Redis script reply type"))
	}
	if err := decodeJSON(payload, destination); err != nil {
		return errors.Join(extractionDomain.ErrStoreUnavailable, fmt.Errorf("decode Redis script reply: %w", err))
	}
	return nil
}