modelRT/repository/redis/data_object_redis_change.go

// Package redis provides Redis persistence helpers.
package redis

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strconv"
	"time"

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

const redisChangeRestoreTimeout = 5 * time.Second

type redisHashChange struct {
	key      string
	field    string
	oldValue string
	newValue string
}

type redisZSetChange struct {
	key       string
	oldValues []redisclient.Z
	newValues []redisclient.Z
}

// RedisChangeSet keeps the Redis changes belonging to one PostgreSQL
// transaction. Changes are prepared first and applied together immediately
// before the PostgreSQL transaction is committed.
type RedisChangeSet struct {
	client      *redisclient.Client
	hashChanges []redisHashChange
	zsetChanges []redisZSetChange
	applied     bool
}

func NewRedisChangeSet(client *redisclient.Client) *RedisChangeSet {
	return &RedisChangeSet{client: client}
}

func (changes *RedisChangeSet) AddHashChange(
	ctx context.Context,
	canonicalKey, field string,
	value any,
) error {
	if changes == nil || changes.client == nil {
		return fmt.Errorf("redis client is not initialized")
	}
	if canonicalKey == "" {
		return fmt.Errorf("canonical redis key is empty")
	}

	newValue, err := redisChangeString(value)
	if err != nil {
		return fmt.Errorf("encode redis data-object value: %w", err)
	}
	oldValue, err := changes.client.HGet(ctx, canonicalKey, field).Result()
	if err != nil {
		return fmt.Errorf("query canonical redis hash %q field %q: %w", canonicalKey, field, err)
	}
	changes.hashChanges = append(changes.hashChanges, redisHashChange{
		key:      canonicalKey,
		field:    field,
		oldValue: oldValue,
		newValue: newValue,
	})
	return nil
}

func (changes *RedisChangeSet) AddMeasurementValueChange(
	ctx context.Context,
	key string,
	value float64,
	timestamp time.Time,
	replace bool,
) error {
	if changes == nil || changes.client == nil {
		return fmt.Errorf("redis client is not initialized")
	}
	if key == "" {
		return fmt.Errorf("measurement redis key is empty")
	}
	keyType, err := changes.client.Type(ctx, key).Result()
	if err != nil {
		return fmt.Errorf("query measurement redis key type for %q: %w", key, err)
	}
	if keyType != "none" && keyType != "zset" {
		return fmt.Errorf("measurement redis key %q has type %q, expected zset", key, keyType)
	}
	oldValues, err := changes.client.ZRangeWithScores(ctx, key, 0, -1).Result()
	if err != nil {
		return fmt.Errorf("query measurement redis values for %q: %w", key, err)
	}

	newMember := strconv.FormatInt(timestamp.UnixNano(), 10)
	newValues := []redisclient.Z{{Score: value, Member: newMember}}
	if !replace {
		newValues = mergeRedisZValues(oldValues, newValues...)
	}
	changes.zsetChanges = append(changes.zsetChanges, redisZSetChange{
		key:       key,
		oldValues: normalizeRedisZValues(oldValues),
		newValues: normalizeRedisZValues(newValues),
	})
	return nil
}

func (changes *RedisChangeSet) Apply(ctx context.Context) error {
	if changes == nil || changes.client == nil {
		return fmt.Errorf("redis client is not initialized")
	}
	if len(changes.hashChanges) == 0 && len(changes.zsetChanges) == 0 {
		return nil
	}

	keys := changes.keys()
	err := changes.client.Watch(ctx, func(tx *redisclient.Tx) error {
		if err := changes.verify(ctx, tx, false); err != nil {
			return err
		}
		_, err := tx.TxPipelined(ctx, func(pipe redisclient.Pipeliner) error {
			for _, change := range changes.hashChanges {
				pipe.HSet(ctx, change.key, change.field, change.newValue)
			}
			for _, change := range changes.zsetChanges {
				pipe.Del(ctx, change.key)
				if len(change.newValues) > 0 {
					pipe.ZAdd(ctx, change.key, change.newValues...)
				}
			}
			return nil
		})
		return err
	}, keys...)
	if err != nil {
		// A connection error can leave EXEC's outcome unknown. Restore only
		// when Redis still contains either the prepared or the applied state.
		restoreCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), redisChangeRestoreTimeout)
		defer cancel()
		if restoreErr := changes.restoreAfterApplyFailure(restoreCtx); restoreErr != nil {
			return fmt.Errorf("apply redis changes: %w; restore redis changes: %v", err, restoreErr)
		}
		return fmt.Errorf("apply redis changes: %w", err)
	}
	changes.applied = true
	return nil
}

// Revert restores Redis after a PostgreSQL commit failure. It uses WATCH and
// only restores values that still match this change set.
func (changes *RedisChangeSet) Revert(ctx context.Context) error {
	if changes == nil || !changes.applied {
		return nil
	}
	return changes.restoreOldValues(ctx, true)
}

func (changes *RedisChangeSet) restoreOldValues(ctx context.Context, compareNew bool) error {
	keys := changes.keys()
	return changes.client.Watch(ctx, func(tx *redisclient.Tx) error {
		if compareNew {
			if err := changes.verify(ctx, tx, true); err != nil {
				return err
			}
		}
		_, err := tx.TxPipelined(ctx, func(pipe redisclient.Pipeliner) error {
			for _, change := range changes.hashChanges {
				pipe.HSet(ctx, change.key, change.field, change.oldValue)
			}
			for _, change := range changes.zsetChanges {
				pipe.Del(ctx, change.key)
				if len(change.oldValues) > 0 {
					pipe.ZAdd(ctx, change.key, change.oldValues...)
				}
			}
			return nil
		})
		return err
	}, keys...)
}

func (changes *RedisChangeSet) restoreAfterApplyFailure(ctx context.Context) error {
	keys := changes.keys()
	return changes.client.Watch(ctx, func(tx *redisclient.Tx) error {
		for _, change := range changes.hashChanges {
			actual, err := tx.HGet(ctx, change.key, change.field).Result()
			if err != nil {
				return fmt.Errorf("verify redis hash %q field %q after apply failure: %w", change.key, change.field, err)
			}
			if actual != change.oldValue && actual != change.newValue {
				return fmt.Errorf("redis hash %q field %q changed concurrently", change.key, change.field)
			}
		}
		for _, change := range changes.zsetChanges {
			actual, err := tx.ZRangeWithScores(ctx, change.key, 0, -1).Result()
			if err != nil {
				return fmt.Errorf("verify redis zset %q after apply failure: %w", change.key, err)
			}
			if !equalRedisZValues(actual, change.oldValues) && !equalRedisZValues(actual, change.newValues) {
				return fmt.Errorf("redis zset %q changed concurrently", change.key)
			}
		}
		_, err := tx.TxPipelined(ctx, func(pipe redisclient.Pipeliner) error {
			for _, change := range changes.hashChanges {
				pipe.HSet(ctx, change.key, change.field, change.oldValue)
			}
			for _, change := range changes.zsetChanges {
				pipe.Del(ctx, change.key)
				if len(change.oldValues) > 0 {
					pipe.ZAdd(ctx, change.key, change.oldValues...)
				}
			}
			return nil
		})
		return err
	}, keys...)
}

func (changes *RedisChangeSet) verify(ctx context.Context, tx *redisclient.Tx, expectNew bool) error {
	for _, change := range changes.hashChanges {
		expected := change.oldValue
		if expectNew {
			expected = change.newValue
		}
		actual, err := tx.HGet(ctx, change.key, change.field).Result()
		if err != nil {
			return fmt.Errorf("verify redis hash %q field %q: %w", change.key, change.field, err)
		}
		if actual != expected {
			return fmt.Errorf("redis hash %q field %q changed concurrently", change.key, change.field)
		}
	}
	for _, change := range changes.zsetChanges {
		expected := change.oldValues
		if expectNew {
			expected = change.newValues
		}
		actual, err := tx.ZRangeWithScores(ctx, change.key, 0, -1).Result()
		if err != nil {
			return fmt.Errorf("verify redis zset %q: %w", change.key, err)
		}
		if !equalRedisZValues(actual, expected) {
			return fmt.Errorf("redis zset %q changed concurrently", change.key)
		}
	}
	return nil
}

func (changes *RedisChangeSet) keys() []string {
	seen := make(map[string]struct{}, len(changes.hashChanges)+len(changes.zsetChanges))
	keys := make([]string, 0, len(seen))
	for _, change := range changes.hashChanges {
		if _, ok := seen[change.key]; !ok {
			seen[change.key] = struct{}{}
			keys = append(keys, change.key)
		}
	}
	for _, change := range changes.zsetChanges {
		if _, ok := seen[change.key]; !ok {
			seen[change.key] = struct{}{}
			keys = append(keys, change.key)
		}
	}
	sort.Strings(keys)
	return keys
}

func redisChangeString(value any) (string, error) {
	switch typedValue := value.(type) {
	case string:
		return typedValue, nil
	case []byte:
		return string(typedValue), nil
	case nil:
		return "null", nil
	case bool:
		return strconv.FormatBool(typedValue), nil
	case int:
		return strconv.Itoa(typedValue), nil
	case int16:
		return strconv.FormatInt(int64(typedValue), 10), nil
	case int64:
		return strconv.FormatInt(typedValue, 10), nil
	case float64:
		return strconv.FormatFloat(typedValue, 'f', -1, 64), nil
	default:
		encoded, err := json.Marshal(typedValue)
		if err != nil {
			return "", err
		}
		return string(encoded), nil
	}
}

func cloneRedisZValues(values []redisclient.Z) []redisclient.Z {
	cloned := make([]redisclient.Z, len(values))
	copy(cloned, values)
	return cloned
}

func mergeRedisZValues(current []redisclient.Z, additions ...redisclient.Z) []redisclient.Z {
	valuesByMember := make(map[string]redisclient.Z, len(current)+len(additions))
	for _, value := range current {
		valuesByMember[fmt.Sprint(value.Member)] = value
	}
	for _, value := range additions {
		valuesByMember[fmt.Sprint(value.Member)] = value
	}
	values := make([]redisclient.Z, 0, len(valuesByMember))
	for _, value := range valuesByMember {
		values = append(values, value)
	}
	return normalizeRedisZValues(values)
}

func normalizeRedisZValues(values []redisclient.Z) []redisclient.Z {
	normalized := cloneRedisZValues(values)
	sort.Slice(normalized, func(i, j int) bool {
		if normalized[i].Score != normalized[j].Score {
			return normalized[i].Score < normalized[j].Score
		}
		return fmt.Sprint(normalized[i].Member) < fmt.Sprint(normalized[j].Member)
	})
	return normalized
}

func equalRedisZValues(left, right []redisclient.Z) bool {
	left = normalizeRedisZValues(left)
	right = normalizeRedisZValues(right)
	if len(left) != len(right) {
		return false
	}
	for index := range left {
		if left[index].Score != right[index].Score ||
			fmt.Sprint(left[index].Member) != fmt.Sprint(right[index].Member) {
			return false
		}
	}
	return true
}