feat: add url shortner

pkg/shorturl/queue.go

@@ -0,0 +1,181 @@
+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
+}