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 }