lislgosms/gateway/internal/inbound/recovery.go

package inbound

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

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

const (
	defaultRecoveryLockTTL     = 30 * time.Second
	defaultRecoveryBackoffBase = 30 * time.Second
	defaultRecoveryBackoffMax  = 10 * time.Minute
)

var ErrRecoveryLockLost = errors.New("recovery lock lost")

type RecoveryStore interface {
	StartAccountRecovery(ctx context.Context, account string, instanceID string) (RecoveryStartDecision, error)
	CompleteAccountRecovery(ctx context.Context, status DownstreamRecoveryStatus) error
	GetAccountRecoveryStatus(ctx context.Context, account string) (DownstreamRecoveryStatus, error)
	ListRecoveryStatuses(ctx context.Context) ([]DownstreamRecoveryStatus, error)
}

type RecoveryStartDecision struct {
	Allowed    bool
	SkipReason string
	Status     DownstreamRecoveryStatus
}

type DownstreamRecoveryStatus struct {
	Account           string    `json:"account"`
	GatewayInstanceID string    `json:"gatewayInstanceId,omitempty"`
	State             string    `json:"state"`
	LockOwner         string    `json:"lockOwner,omitempty"`
	LockToken         string    `json:"lockToken,omitempty"`
	LockAcquiredAt    time.Time `json:"lockAcquiredAt,omitempty"`
	LockExpiresAt     time.Time `json:"lockExpiresAt,omitempty"`
	LastAttemptAt     time.Time `json:"lastAttemptAt,omitempty"`
	LastSuccessAt     time.Time `json:"lastSuccessAt,omitempty"`
	LastFailureAt     time.Time `json:"lastFailureAt,omitempty"`
	NextRetryAt       time.Time `json:"nextRetryAt,omitempty"`
	AttemptCount      int       `json:"attemptCount"`
	FailureCategory   string    `json:"failureCategory,omitempty"`
	LastError         string    `json:"lastError,omitempty"`
	LastSkipReason    string    `json:"lastSkipReason,omitempty"`
}

type RedisRecoveryStore struct {
	Client *redis.Client
	Prefix string
}

func NewRedisRecoveryStore(redisURL string) (*RedisRecoveryStore, error) {
	if redisURL == "" {
		redisURL = "redis://127.0.0.1:6379"
	}
	options, err := redis.ParseURL(redisURL)
	if err != nil {
		return nil, err
	}
	return &RedisRecoveryStore{
		Client: redis.NewClient(options),
		Prefix: "gateway:downstream:recovery",
	}, nil
}

func (s *RedisRecoveryStore) StartAccountRecovery(ctx context.Context, account string, instanceID string) (RecoveryStartDecision, error) {
	if s == nil || s.Client == nil || account == "" {
		return RecoveryStartDecision{Allowed: true}, nil
	}
	status, err := s.getStatus(ctx, account)
	if err != nil {
		return RecoveryStartDecision{}, err
	}
	now := time.Now().UTC()
	if !status.NextRetryAt.IsZero() && status.NextRetryAt.After(now) {
		status.LastSkipReason = "backoff"
		status.FailureCategory = "backoff"
		return RecoveryStartDecision{Allowed: false, SkipReason: "backoff", Status: status}, nil
	}
	lockToken := newRecoveryLockToken(instanceID)
	acquired, err := s.Client.SetNX(ctx, s.lockKey(account), lockToken, defaultRecoveryLockTTL).Result()
	if err != nil {
		return RecoveryStartDecision{}, err
	}
	if !acquired {
		status.LockToken, _ = s.Client.Get(ctx, s.lockKey(account)).Result()
		status.LockOwner = recoveryLockOwner(status.LockToken)
		status.LockExpiresAt = lockExpiresAt(ctx, s.Client, s.lockKey(account), now)
		status.LastSkipReason = "locked"
		status.FailureCategory = "lock_contended"
		return RecoveryStartDecision{Allowed: false, SkipReason: "locked", Status: status}, nil
	}
	status.Account = account
	status.GatewayInstanceID = instanceID
	status.State = "running"
	status.LockOwner = instanceID
	status.LockToken = lockToken
	status.LockAcquiredAt = now
	status.LockExpiresAt = now.Add(defaultRecoveryLockTTL)
	status.LastAttemptAt = now
	status.LastSkipReason = ""
	status.FailureCategory = ""
	if err := s.saveStatus(ctx, status); err != nil {
		_ = s.deleteLockIfOwned(ctx, account, lockToken)
		return RecoveryStartDecision{}, err
	}
	return RecoveryStartDecision{Allowed: true, Status: status}, nil
}

func (s *RedisRecoveryStore) CompleteAccountRecovery(ctx context.Context, status DownstreamRecoveryStatus) error {
	if s == nil || s.Client == nil || status.Account == "" {
		return nil
	}
	now := time.Now().UTC()
	if status.LastAttemptAt.IsZero() {
		status.LastAttemptAt = now
	}
	if status.LockOwner == "" {
		status.LockOwner = status.GatewayInstanceID
	}
	switch status.State {
	case "success":
		status.LastSuccessAt = now
		status.AttemptCount = 0
		status.NextRetryAt = time.Time{}
		status.LastError = ""
		status.FailureCategory = ""
	case "waiting_connection", "failed", "partial":
		status.AttemptCount++
		if status.State == "failed" {
			status.LastFailureAt = now
		}
		if status.FailureCategory == "" {
			status.FailureCategory = classifyRecoveryFailure(status)
		}
		status.NextRetryAt = now.Add(recoveryBackoffDelay(status.AttemptCount))
	default:
		status.State = "unknown"
	}
	payload, err := json.Marshal(status)
	if err != nil {
		return err
	}
	result, err := s.Client.Eval(ctx, completeRecoveryScript, []string{s.lockKey(status.Account), s.statusKey(status.Account), s.accountsKey()}, status.LockToken, payload, int64((24*time.Hour)/time.Millisecond), status.Account).Int()
	if err != nil {
		return err
	}
	if result == 0 {
		return ErrRecoveryLockLost
	}
	return nil
}

func classifyRecoveryFailure(status DownstreamRecoveryStatus) string {
	switch status.State {
	case "waiting_connection":
		return "client_disconnected"
	case "partial":
		return "partial_delivery_failed"
	case "failed":
		if status.LastSkipReason == "backoff" {
			return "backoff"
		}
		if status.LastSkipReason == "locked" {
			return "lock_contended"
		}
		if status.LastError != "" {
			return "flush_failed"
		}
	}
	return "unknown"
}

func (s *RedisRecoveryStore) ListRecoveryStatuses(ctx context.Context) ([]DownstreamRecoveryStatus, error) {
	if s == nil || s.Client == nil {
		return nil, nil
	}
	accounts, err := s.Client.SMembers(ctx, s.accountsKey()).Result()
	if err != nil {
		return nil, err
	}
	result := make([]DownstreamRecoveryStatus, 0, len(accounts))
	for _, account := range accounts {
		status, getErr := s.getStatus(ctx, account)
		if getErr == redis.Nil {
			_ = s.Client.SRem(ctx, s.accountsKey(), account).Err()
			continue
		}
		if getErr != nil {
			return nil, getErr
		}
		result = append(result, status)
	}
	return result, nil
}

func (s *RedisRecoveryStore) GetAccountRecoveryStatus(ctx context.Context, account string) (DownstreamRecoveryStatus, error) {
	if s == nil || s.Client == nil || account == "" {
		return DownstreamRecoveryStatus{Account: account}, nil
	}
	return s.getStatus(ctx, account)
}

func (s *RedisRecoveryStore) getStatus(ctx context.Context, account string) (DownstreamRecoveryStatus, error) {
	payload, err := s.Client.Get(ctx, s.statusKey(account)).Bytes()
	if err != nil {
		if err == redis.Nil {
			return DownstreamRecoveryStatus{Account: account}, nil
		}
		return DownstreamRecoveryStatus{}, err
	}
	var status DownstreamRecoveryStatus
	if err := json.Unmarshal(payload, &status); err != nil {
		return DownstreamRecoveryStatus{}, fmt.Errorf("decode recovery status %s: %w", account, err)
	}
	if status.Account == "" {
		status.Account = account
	}
	return status, nil
}

func (s *RedisRecoveryStore) saveStatus(ctx context.Context, status DownstreamRecoveryStatus) error {
	payload, err := json.Marshal(status)
	if err != nil {
		return err
	}
	pipe := s.Client.TxPipeline()
	pipe.Set(ctx, s.statusKey(status.Account), payload, 24*time.Hour)
	pipe.SAdd(ctx, s.accountsKey(), status.Account)
	pipe.Expire(ctx, s.accountsKey(), 24*time.Hour)
	_, err = pipe.Exec(ctx)
	return err
}

func (s *RedisRecoveryStore) deleteLockIfOwned(ctx context.Context, account string, token string) error {
	if token == "" {
		return nil
	}
	_, err := s.Client.Eval(ctx, deleteLockIfOwnedScript, []string{s.lockKey(account)}, token).Result()
	return err
}

func (s *RedisRecoveryStore) statusKey(account string) string {
	return fmt.Sprintf("%s:status:%s", s.prefix(), account)
}

func (s *RedisRecoveryStore) lockKey(account string) string {
	return fmt.Sprintf("%s:lock:%s", s.prefix(), account)
}

func (s *RedisRecoveryStore) accountsKey() string {
	return fmt.Sprintf("%s:accounts", s.prefix())
}

func (s *RedisRecoveryStore) prefix() string {
	if s != nil && s.Prefix != "" {
		return s.Prefix
	}
	return "gateway:downstream:recovery"
}

func recoveryBackoffDelay(attemptCount int) time.Duration {
	if attemptCount <= 0 {
		return defaultRecoveryBackoffBase
	}
	delay := defaultRecoveryBackoffBase
	for step := 1; step < attemptCount; step++ {
		delay *= 2
		if delay >= defaultRecoveryBackoffMax {
			return defaultRecoveryBackoffMax
		}
	}
	if delay > defaultRecoveryBackoffMax {
		return defaultRecoveryBackoffMax
	}
	return delay
}

func newRecoveryLockToken(instanceID string) string {
	return fmt.Sprintf("%s:%d", instanceID, time.Now().UTC().UnixNano())
}

func recoveryLockOwner(token string) string {
	for index, char := range token {
		if char == ':' {
			return token[:index]
		}
	}
	return token
}

func lockExpiresAt(ctx context.Context, client *redis.Client, key string, now time.Time) time.Time {
	ttl, err := client.TTL(ctx, key).Result()
	if err != nil || ttl <= 0 {
		return time.Time{}
	}
	return now.Add(ttl)
}

const completeRecoveryScript = `
local current = redis.call("GET", KEYS[1])
if current ~= ARGV[1] then
  return 0
end
redis.call("SET", KEYS[2], ARGV[2], "PX", ARGV[3])
redis.call("SADD", KEYS[3], ARGV[4])
redis.call("PEXPIRE", KEYS[3], ARGV[3])
redis.call("DEL", KEYS[1])
return 1
`

const deleteLockIfOwnedScript = `
local current = redis.call("GET", KEYS[1])
if current == ARGV[1] then
  return redis.call("DEL", KEYS[1])
end
return 0
`