First commit Go counter service

2026-10-15 02:12:18 +02:00
commit c43ff17627
59 changed files with 11205 additions and 0 deletions
@@ -0,0 +1,238 @@
+// 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
+)
+
+// 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
+}
+
+// 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,
+		}
+		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,
+		}
+	}
+	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))
+	_, 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)
+		}
+	}
+	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
+}
+
+// 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) {
+		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")
+}
+
+// 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...)
+}