spsays-bff/pkg/shorturl/queue.go

package shorturl

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

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

// Redis key prefixes.
const (
	pendingKeyPrefix = "share:pending:"
	cacheKeyPrefix   = "share:cache:"
	createQueueKey   = "share:queue"
	clickQueueKey    = "share:click-queue"
	deadLetterKey    = "share:dead-letter"
)

// Lifetimes.
const (
	PendingTTL = 1 * time.Hour
	CacheTTL   = 7 * 24 * time.Hour
)

// PendingPayload is what the POST handler stores in Redis before the worker
// drains it into Postgres.
type PendingPayload struct {
	URL            string    `json:"url"`
	CreatorIP      string    `json:"creator_ip,omitempty"`
	CreatorUA      string    `json:"creator_ua,omitempty"`
	CreatorReferer string    `json:"creator_referer,omitempty"`
	CreatorCountry string    `json:"creator_country,omitempty"`
	CreatedAt      time.Time `json:"created_at"`
}

// ClickEvent is what the resolve handler enqueues per click.
type ClickEvent struct {
	Code      string    `json:"code"`
	IP        string    `json:"ip,omitempty"`
	UA        string    `json:"ua,omitempty"`
	Referer   string    `json:"referer,omitempty"`
	Country   string    `json:"country,omitempty"`
	IsBot     bool      `json:"is_bot"`
	ClickedAt time.Time `json:"clicked_at"`
}

// Queue provides write-behind primitives backed by Redis. The HTTP handlers
// only ever talk to Queue; the workers drain its lists into Postgres.
type Queue struct {
	rdb *redis.Client
}

func NewQueue(rdb *redis.Client) *Queue {
	return &Queue{rdb: rdb}
}

// AllocateCode generates a random unused code via SETNX on the pending key.
// The pending key is created as a placeholder ('_') so concurrent allocators
// won't collide; MarkPending then overwrites it with the real payload.
func (q *Queue) AllocateCode(ctx context.Context) (string, error) {
	for attempt := 0; attempt < 5; attempt++ {
		code, err := RandomCode()
		if err != nil {
			return "", err
		}
		ok, err := q.rdb.SetNX(ctx, pendingKeyPrefix+code, "_", PendingTTL).Result()
		if err != nil {
			return "", fmt.Errorf("setnx pending: %w", err)
		}
		if ok {
			return code, nil
		}
	}
	return "", errors.New("could not allocate unique code after 5 attempts")
}

// MarkPending overwrites the placeholder with the real payload.
func (q *Queue) MarkPending(ctx context.Context, code string, p PendingPayload) error {
	b, err := json.Marshal(p)
	if err != nil {
		return fmt.Errorf("marshal pending: %w", err)
	}
	return q.rdb.Set(ctx, pendingKeyPrefix+code, b, PendingTTL).Err()
}

// LookupPending returns the in-flight payload, or (zero, false) if not present.
// Used by the resolve handler when a code was created very recently and the
// create-worker hasn't yet flushed it to Postgres.
func (q *Queue) LookupPending(ctx context.Context, code string) (PendingPayload, bool) {
	val, err := q.rdb.Get(ctx, pendingKeyPrefix+code).Result()
	if err != nil {
		return PendingPayload{}, false
	}
	if val == "_" || val == "" {
		return PendingPayload{}, false // placeholder, not yet marked
	}
	var p PendingPayload
	if err := json.Unmarshal([]byte(val), &p); err != nil {
		return PendingPayload{}, false
	}
	return p, true
}

// UnmarkPending deletes the pending entry once the worker has persisted to PG.
func (q *Queue) UnmarkPending(ctx context.Context, code string) error {
	return q.rdb.Del(ctx, pendingKeyPrefix+code).Err()
}

// EnqueueCreate pushes a code onto the create-worker queue.
func (q *Queue) EnqueueCreate(ctx context.Context, code string) error {
	return q.rdb.LPush(ctx, createQueueKey, code).Err()
}

// DequeueCreate blocks until a code is available (or timeout). Returns
// ("", nil) on timeout, ("", err) on real error, (code, nil) on hit.
func (q *Queue) DequeueCreate(ctx context.Context, timeout time.Duration) (string, error) {
	res, err := q.rdb.BRPop(ctx, timeout, createQueueKey).Result()
	if err == redis.Nil {
		return "", nil
	}
	if err != nil {
		return "", err
	}
	if len(res) < 2 {
		return "", nil
	}
	return res[1], nil
}

// PushDeadLetter records an unprocessable create-job for manual inspection.
func (q *Queue) PushDeadLetter(ctx context.Context, code, reason string) error {
	payload := map[string]any{"code": code, "reason": reason, "at": time.Now().UTC()}
	b, _ := json.Marshal(payload)
	return q.rdb.LPush(ctx, deadLetterKey, b).Err()
}

// MarkCached stores a resolved long_url for fast future reads. Refreshed
// (TTL reset) on every successful resolve so popular links stay hot.
func (q *Queue) MarkCached(ctx context.Context, code, longURL string) error {
	return q.rdb.Set(ctx, cacheKeyPrefix+code, longURL, CacheTTL).Err()
}

// LookupCached returns the cached long_url, or (empty, false) on miss.
func (q *Queue) LookupCached(ctx context.Context, code string) (string, bool) {
	val, err := q.rdb.Get(ctx, cacheKeyPrefix+code).Result()
	if err != nil || val == "" {
		return "", false
	}
	return val, true
}

// EnqueueClick pushes a click event onto the click-worker queue.
func (q *Queue) EnqueueClick(ctx context.Context, ev ClickEvent) error {
	b, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	return q.rdb.LPush(ctx, clickQueueKey, b).Err()
}

// DequeueClick blocks until a click event is available.
func (q *Queue) DequeueClick(ctx context.Context, timeout time.Duration) (ClickEvent, bool, error) {
	res, err := q.rdb.BRPop(ctx, timeout, clickQueueKey).Result()
	if err == redis.Nil {
		return ClickEvent{}, false, nil
	}
	if err != nil {
		return ClickEvent{}, false, err
	}
	if len(res) < 2 {
		return ClickEvent{}, false, nil
	}
	var ev ClickEvent
	if err := json.Unmarshal([]byte(res[1]), &ev); err != nil {
		return ClickEvent{}, false, err
	}
	return ev, true, nil
}