fm_be/internal/repository/mysql/idempotency_store.go

package mysql

import (
	"context"
	"errors"
	"strings"
	"time"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"wucher/internal/domain/transient"
	"wucher/internal/queue"
)

const (
	idempotencyStateProcessing = "processing"
	idempotencyStateDone       = "done"
)

type QueueIdempotencyStore struct {
	db     *gorm.DB
	prefix string
	now    func() time.Time
}

func NewQueueIdempotencyStore(db *gorm.DB, prefix string) *QueueIdempotencyStore {
	return &QueueIdempotencyStore{
		db:     db,
		prefix: strings.TrimSpace(prefix),
		now:    func() time.Time { return time.Now().UTC() },
	}
}

func (s *QueueIdempotencyStore) Acquire(ctx context.Context, key string, processingTTL time.Duration) (queue.IdempotencyState, error) {
	if s == nil || s.db == nil {
		return queue.IdempotencyStateAcquired, errors.New("queue idempotency db is required")
	}
	storeKey := s.key(key)
	if storeKey == "" {
		return "", errors.New("queue idempotency key is required")
	}
	if processingTTL <= 0 {
		processingTTL = 15 * time.Minute
	}

	now := s.now().UTC()
	expiresAt := now.Add(processingTTL)
	record := transient.QueueIdempotencyRecord{
		IdempotencyKey: storeKey,
		State:          idempotencyStateProcessing,
		ExpiresAt:      expiresAt,
	}
	create := s.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).Create(&record)
	if create.Error != nil {
		return "", create.Error
	}
	if create.RowsAffected == 1 {
		return queue.IdempotencyStateAcquired, nil
	}

	existing, err := s.load(ctx, storeKey)
	if errors.Is(err, gorm.ErrRecordNotFound) {
		retry := s.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).Create(&record)
		if retry.Error != nil {
			return "", retry.Error
		}
		if retry.RowsAffected == 1 {
			return queue.IdempotencyStateAcquired, nil
		}
		return queue.IdempotencyStateInProgress, nil
	}
	if err != nil {
		return "", err
	}

	if !existing.ExpiresAt.After(now) {
		update := s.db.WithContext(ctx).
			Model(&transient.QueueIdempotencyRecord{}).
			Where("idempotency_key = ? AND expires_at <= ?", storeKey, now).
			Updates(map[string]any{
				"state":      idempotencyStateProcessing,
				"expires_at": expiresAt,
				"updated_at": now,
			})
		if update.Error != nil {
			return "", update.Error
		}
		if update.RowsAffected == 1 {
			return queue.IdempotencyStateAcquired, nil
		}
		existing, err = s.load(ctx, storeKey)
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return queue.IdempotencyStateAcquired, nil
		}
		if err != nil {
			return "", err
		}
	}

	if existing.State == idempotencyStateDone {
		return queue.IdempotencyStateDuplicate, nil
	}
	return queue.IdempotencyStateInProgress, nil
}

func (s *QueueIdempotencyStore) MarkCompleted(ctx context.Context, key string, ttl time.Duration) error {
	if s == nil || s.db == nil {
		return errors.New("queue idempotency db is required")
	}
	storeKey := s.key(key)
	if storeKey == "" {
		return errors.New("queue idempotency key is required")
	}
	if ttl <= 0 {
		ttl = 24 * time.Hour
	}
	now := s.now().UTC()
	record := transient.QueueIdempotencyRecord{
		IdempotencyKey: storeKey,
		State:          idempotencyStateDone,
		ExpiresAt:      now.Add(ttl),
	}
	return s.db.WithContext(ctx).Clauses(clause.OnConflict{
		Columns: []clause.Column{{Name: "idempotency_key"}},
		DoUpdates: clause.Assignments(map[string]any{
			"state":      record.State,
			"expires_at": record.ExpiresAt,
			"updated_at": now,
		}),
	}).Create(&record).Error
}

func (s *QueueIdempotencyStore) Release(ctx context.Context, key string) error {
	if s == nil || s.db == nil {
		return errors.New("queue idempotency db is required")
	}
	storeKey := s.key(key)
	if storeKey == "" {
		return nil
	}
	return s.db.WithContext(ctx).Delete(&transient.QueueIdempotencyRecord{}, "idempotency_key = ?", storeKey).Error
}

func (s *QueueIdempotencyStore) load(ctx context.Context, key string) (*transient.QueueIdempotencyRecord, error) {
	var record transient.QueueIdempotencyRecord
	if err := s.db.WithContext(ctx).Where("idempotency_key = ?", key).First(&record).Error; err != nil {
		return nil, err
	}
	return &record, nil
}

func (s *QueueIdempotencyStore) key(key string) string {
	key = strings.TrimSpace(key)
	if key == "" {
		return ""
	}
	return s.prefix + key
}