counters/internal/store/store.go

// Package store holds the Redis counters. Writes (INCR) go to the master,
// reads (MGET for /counters) go to the replicas; in Sentinel mode go-redis
// discovers both, in static mode the addresses come from the config.
//
// Key format, shared with the rspamd Lua plugin and read by migadu_web
// (apps/core/lib/core/redis_sentinel.ex):
//
//	cnt/<A|D|O>/<id>/d/<YYYY-MM-DD>/<i|o>
//
// id is the address, the domain or the organisation id; the day is the UTC
// calendar day; every key lives 31 days so the web app can chart a month.
package store

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

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

	"migadu_counters/internal/config"
)

const (
	// Prefix of every counter key.
	Prefix = "cnt"
	// TTL is refreshed on every increment, like the Lua plugin did.
	TTL = 31 * 24 * time.Hour

	// MarkerLimit and MarkerWarn name the two notification markers of a counter.
	MarkerLimit = "limit"
	MarkerWarn  = "warn"
	// MarkerTTL outlives the UTC day the marker belongs to.
	MarkerTTL = 48 * time.Hour
)

// ErrNotApplied wraps an Incr failure where Redis answered and rejected the
// whole transaction (EXECABORT, READONLY, LOADING, ...): nothing was applied,
// so running it again cannot double count.
var ErrNotApplied = errors.New("transaction not applied")

// Key builds a counter key: Key("A", "x@example.com", "2026-10-13", "i").
func Key(level, id, date, flow string) string {
	return Prefix + "/" + level + "/" + id + "/d/" + date + "/" + flow
}

// MarkerKey builds the key recording that a counter's notification went out.
// It uses its own prefix because the web app scans cnt/*<address>*.
func MarkerKey(level, id, date, flow, kind string) string {
	return "ntf/" + level + "/" + id + "/d/" + date + "/" + flow + "/" + kind
}

// DayKey returns the UTC calendar day of t, as used in the keys. The web app
// builds its keys with Date.utc_today(), so this must never use local time.
func DayKey(t time.Time) string { return t.UTC().Format("2006-01-02") }

// Store is a write client plus one or more read clients.
type Store struct {
	write   *redis.Client
	readers []*redis.Client
	next    atomic.Uint64
}

// New builds the clients from the config. Nothing is dialled until the first
// command, so an unreachable Redis does not prevent the daemon from starting.
func New(cfg config.RedisConfig) (*Store, error) {
	if cfg.DB != 0 {
		return nil, fmt.Errorf("redis.db must be 0 (got %d): migadu_web reads the message counters from DB 0", cfg.DB)
	}
	timeout := cfg.Timeout()

	if cfg.SentinelMode() {
		base := redis.FailoverOptions{
			MasterName:       cfg.MasterName,
			SentinelAddrs:    cfg.Sentinels,
			SentinelPassword: cfg.SentinelPassword,
			Password:         cfg.Password,
			DB:               0,
			PoolSize:         cfg.PoolSize,
			DialTimeout:      timeout,
			ReadTimeout:      timeout,
			WriteTimeout:     timeout,

			ContextTimeoutEnabled: true,
		}
		write := redis.NewFailoverClient(&base)
		// ReplicaOnly picks a random replica and falls back to the master when
		// no replica is known.
		replica := base
		replica.ReplicaOnly = true
		read := redis.NewFailoverClient(&replica)
		return &Store{write: write, readers: []*redis.Client{read}}, nil
	}

	if cfg.WriteAddr == "" {
		return nil, errors.New("redis.write_addr is required when no sentinels are configured")
	}
	opts := func(addr string) *redis.Options {
		return &redis.Options{
			Addr:         addr,
			Password:     cfg.Password,
			DB:           0,
			PoolSize:     cfg.PoolSize,
			DialTimeout:  timeout,
			ReadTimeout:  timeout,
			WriteTimeout: timeout,

			ContextTimeoutEnabled: true,
		}
	}
	addrs := cfg.ReadAddrs
	if len(addrs) == 0 {
		addrs = []string{cfg.WriteAddr}
	}
	readers := make([]*redis.Client, 0, len(addrs))
	for _, a := range addrs {
		readers = append(readers, redis.NewClient(opts(a)))
	}
	return &Store{write: redis.NewClient(opts(cfg.WriteAddr)), readers: readers}, nil
}

// NewWithClients wires pre-built clients (tests). With no readers the write
// client is read from.
func NewWithClients(write *redis.Client, readers []*redis.Client) *Store {
	if len(readers) == 0 {
		readers = []*redis.Client{write}
	}
	return &Store{write: write, readers: readers}
}

// reader returns the next read client, round-robin.
func (s *Store) reader() *redis.Client {
	n := s.next.Add(1) - 1
	return s.readers[n%uint64(len(s.readers))]
}

// IncrResult is the outcome of one key's INCR inside the transaction.
type IncrResult struct {
	Count int64 // value after the increment; meaningless when Err != nil
	Err   error // this key's own failure (e.g. it holds a non-number)
}

// Incr increments every key by one and refreshes its TTL in one MULTI/EXEC,
// returning one result per key in order. EXEC's replies are the
// authoritative counts, so no read-after-write is needed and replica lag
// does not matter. A per-key failure (a key holding a non-number) is reported
// in that key's result; the other keys are still incremented. The returned
// error is set only when the transaction as a whole failed, in which case
// OutcomeUnknown tells whether it may nevertheless have been applied.
func (s *Store) Incr(ctx context.Context, keys []string) ([]IncrResult, error) {
	if len(keys) == 0 {
		return nil, nil
	}
	cmds := make([]*redis.IntCmd, len(keys))
	all, err := s.write.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
		for i, k := range keys {
			cmds[i] = pipe.Incr(ctx, k)
			pipe.Expire(ctx, k, TTL)
		}
		return nil
	})
	out := make([]IncrResult, len(keys))
	if err != nil {
		// TxPipelined returns the first failed command's error. A Redis reply
		// error (redis.Error) means EXEC answered and every command has a
		// definite result: value or its own error. Anything else (dial, write,
		// read timeout) is a transport failure of the whole transaction.
		var rerr redis.Error
		if !errors.As(err, &rerr) {
			return nil, fmt.Errorf("redis incr: %w", err)
		}
		if allFailed(all) {
			return nil, fmt.Errorf("redis incr: %w: %w", ErrNotApplied, err)
		}
	}
	for i, c := range cmds {
		if cerr := c.Err(); cerr != nil {
			out[i].Err = fmt.Errorf("redis incr %s: %w", keys[i], cerr)
			continue
		}
		out[i].Count = c.Val()
	}
	return out, nil
}

// allFailed tells a rejected transaction from per-key failures: an INCR on a
// non-number fails alone and its EXPIRE still succeeds, so only a transaction
// Redis refused as a whole leaves every command with an error.
func allFailed(cmds []redis.Cmder) bool {
	for _, c := range cmds {
		if c.Err() == nil {
			return false
		}
	}
	return len(cmds) > 0
}

// OutcomeUnknown reports whether a failed Incr may still have been applied
// by Redis: the request may have been written but the reply lost (timeout,
// EOF, reset). A refused connection, a dial failure or a cancelled context
// never reached Redis and is a definite failure.
func OutcomeUnknown(err error) bool {
	if err == nil || errors.Is(err, context.Canceled) || errors.Is(err, ErrNotApplied) {
		return false
	}
	// No connection could be had from the pool, so no command was sent. go-redis
	// keeps this error in an internal package; only its text is available.
	if strings.Contains(err.Error(), "connection pool timeout") {
		return false
	}
	var opErr *net.OpError
	if errors.As(err, &opErr) && opErr.Op == "dial" {
		return false
	}
	if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
		return true
	}
	var nerr net.Error
	if errors.As(err, &nerr) && nerr.Timeout() {
		return true
	}
	msg := err.Error()
	return strings.Contains(msg, "timeout") || strings.Contains(msg, "reset") || strings.Contains(msg, "broken pipe") || strings.Contains(msg, "EOF")
}

// MarkNotified records that a counter's notification was sent. It overwrites
// any older marker: a counter the web app reset may legitimately cross again.
func (s *Store) MarkNotified(ctx context.Context, key string) error {
	return s.write.Set(ctx, key, "1", MarkerTTL).Err()
}

// ClaimNotified atomically takes the right to send a notification. It returns
// true for exactly one caller while the marker exists.
func (s *Store) ClaimNotified(ctx context.Context, key string) (bool, error) {
	return s.write.SetNX(ctx, key, "1", MarkerTTL).Result()
}

// GetMaster reads the keys from the write side (the master) with one MGET;
// used to establish the outcome of an Incr whose reply was lost, when a
// replica might not have the increment yet. A missing key reads as 0.
func (s *Store) GetMaster(ctx context.Context, keys []string) ([]int64, error) {
	return mget(ctx, s.write, keys)
}

// Get reads the keys from a replica with one MGET; a missing key reads as 0.
func (s *Store) Get(ctx context.Context, keys []string) ([]int64, error) {
	return mget(ctx, s.reader(), keys)
}

func mget(ctx context.Context, c *redis.Client, keys []string) ([]int64, error) {
	if len(keys) == 0 {
		return nil, nil
	}
	vals, err := c.MGet(ctx, keys...).Result()
	if err != nil {
		return nil, fmt.Errorf("redis mget: %w", err)
	}
	out := make([]int64, len(keys))
	for i, v := range vals {
		str, ok := v.(string)
		if !ok {
			continue // nil = missing key
		}
		n, err := strconv.ParseInt(str, 10, 64)
		if err != nil {
			return nil, fmt.Errorf("redis key %s holds non-integer value %q", keys[i], str)
		}
		out[i] = n
	}
	return out, nil
}

// Ping checks the write side (the master), which is what /health cares about.
func (s *Store) Ping(ctx context.Context) error { return s.write.Ping(ctx).Err() }

// Close closes every client.
func (s *Store) Close() error {
	errs := []error{s.write.Close()}
	for _, r := range s.readers {
		if r != s.write {
			errs = append(errs, r.Close())
		}
	}
	return errors.Join(errs...)
}