summercms/modules/pier/store_redis.go

package pier

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"sync/atomic"
	"time"

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

// Redis key and channel layout under cooler's application prefix.
const (
	redisSessionKey = "mcp:s:"     // <prefix>mcp:s:<id> holds the JSON record
	redisBindingKey = "mcp:b:"     // <prefix>mcp:b:<binding> is the set of its session ids
	redisEventsKey  = "mcp:events" // <prefix>mcp:events carries JSON events
)

// redisStore keeps session records in Redis with PX TTLs and fans events
// out through pub/sub, so several instances share sessions and react to
// each other's events.
type redisStore struct {
	client *redis.Client
	prefix string
}

func newRedisStore(client *redis.Client, prefix string) *redisStore {
	return &redisStore{client: client, prefix: prefix}
}

func (s *redisStore) sessionKey(id string) string { return s.prefix + redisSessionKey + id }

func (s *redisStore) bindingKey(binding string) string { return s.prefix + redisBindingKey + binding }

func (s *redisStore) channel() string { return s.prefix + redisEventsKey }

func (s *redisStore) Put(ctx context.Context, id string, rec Record, ttl time.Duration) error {
	raw, err := json.Marshal(rec)
	if err != nil {
		return err
	}
	pipe := s.client.TxPipeline()
	pipe.Set(ctx, s.sessionKey(id), raw, ttl)
	pipe.SAdd(ctx, s.bindingKey(rec.Binding), id)
	pipe.PExpire(ctx, s.bindingKey(rec.Binding), ttl)
	_, err = pipe.Exec(ctx)
	return err
}

// Update rewrites an existing session record (SET XX) and refreshes the
// binding set; a record that is gone stays gone.
func (s *redisStore) Update(ctx context.Context, id string, rec Record, ttl time.Duration) (bool, error) {
	raw, err := json.Marshal(rec)
	if err != nil {
		return false, err
	}
	err = s.client.SetArgs(ctx, s.sessionKey(id), raw, redis.SetArgs{Mode: "XX", TTL: ttl}).Err()
	if errors.Is(err, redis.Nil) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	pipe := s.client.TxPipeline()
	pipe.SAdd(ctx, s.bindingKey(rec.Binding), id)
	pipe.PExpire(ctx, s.bindingKey(rec.Binding), ttl)
	_, err = pipe.Exec(ctx)
	return true, err
}

func (s *redisStore) Get(ctx context.Context, id string) (Record, bool, error) {
	raw, err := s.client.Get(ctx, s.sessionKey(id)).Bytes()
	if errors.Is(err, redis.Nil) {
		return Record{}, false, nil
	}
	if err != nil {
		return Record{}, false, err
	}
	var rec Record
	if err := json.Unmarshal(raw, &rec); err != nil {
		return Record{}, false, err
	}
	return rec, true, nil
}

// Touch slides the record's TTL and its binding set's with it.
func (s *redisStore) Touch(ctx context.Context, id string, ttl time.Duration) error {
	rec, ok, err := s.Get(ctx, id)
	if err != nil || !ok {
		return err
	}
	pipe := s.client.TxPipeline()
	pipe.PExpire(ctx, s.sessionKey(id), ttl)
	pipe.PExpire(ctx, s.bindingKey(rec.Binding), ttl)
	_, err = pipe.Exec(ctx)
	return err
}

func (s *redisStore) Delete(ctx context.Context, id string) error {
	rec, ok, err := s.Get(ctx, id)
	if err != nil {
		return err
	}
	pipe := s.client.TxPipeline()
	pipe.Del(ctx, s.sessionKey(id))
	if ok {
		pipe.SRem(ctx, s.bindingKey(rec.Binding), id)
	}
	_, err = pipe.Exec(ctx)
	return err
}

// Count answers the live sessions of a binding, pruning members whose
// record has expired from the set.
func (s *redisStore) Count(ctx context.Context, binding string) (int, error) {
	members, err := s.client.SMembers(ctx, s.bindingKey(binding)).Result()
	if err != nil {
		return 0, err
	}
	n := 0
	for _, id := range members {
		exists, err := s.client.Exists(ctx, s.sessionKey(id)).Result()
		if err != nil {
			return 0, err
		}
		if exists > 0 {
			n++
			continue
		}
		_ = s.client.SRem(ctx, s.bindingKey(binding), id).Err()
	}
	return n, nil
}

// reserveScript is the atomic per-binding reservation: it drops the binding
// set's members whose record has expired, refuses at max (0) and otherwise
// writes the record with its PX TTL, adds the id to the set and slides the
// set's TTL (1). KEYS: the binding set, the record. ARGV: the record key
// prefix, max, the record JSON, the TTL in milliseconds, the session id.
var reserveScript = redis.NewScript(`
local members = redis.call('SMEMBERS', KEYS[1])
local n = 0
for _, m in ipairs(members) do
	if redis.call('EXISTS', ARGV[1] .. m) == 1 then
		n = n + 1
	else
		redis.call('SREM', KEYS[1], m)
	end
end
if n >= tonumber(ARGV[2]) then
	return 0
end
redis.call('SET', KEYS[2], ARGV[3], 'PX', ARGV[4])
redis.call('SADD', KEYS[1], ARGV[5])
redis.call('PEXPIRE', KEYS[1], ARGV[4])
return 1
`)

// Reserve runs reserveScript, so the count and the write are one step on
// the server.
func (s *redisStore) Reserve(ctx context.Context, binding, id string, rec Record, max int, ttl time.Duration) (bool, error) {
	raw, err := json.Marshal(rec)
	if err != nil {
		return false, err
	}
	ms := ttl.Milliseconds()
	if ms < 1 {
		ms = 1
	}
	n, err := reserveScript.Run(ctx, s.client,
		[]string{s.bindingKey(binding), s.sessionKey(id)},
		s.prefix+redisSessionKey, max, string(raw), ms, id,
	).Int()
	if err != nil {
		return false, err
	}
	return n == 1, nil
}

// Release drops a reservation: the record and its binding set member.
func (s *redisStore) Release(ctx context.Context, binding, id string) error {
	pipe := s.client.TxPipeline()
	pipe.Del(ctx, s.sessionKey(id))
	pipe.SRem(ctx, s.bindingKey(binding), id)
	_, err := pipe.Exec(ctx)
	return err
}

func (s *redisStore) Publish(ctx context.Context, ev Event) error {
	raw, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	return s.client.Publish(ctx, s.channel(), raw).Err()
}

// Subscribe delivers every event published on the channel to fn, on its
// own goroutine, until ctx ends.
func (s *redisStore) Subscribe(ctx context.Context, fn func(Event)) error {
	sub := s.client.Subscribe(ctx, s.channel())
	if _, err := sub.Receive(ctx); err != nil {
		_ = sub.Close()
		return err
	}
	go func() {
		defer sub.Close()
		ch := sub.Channel()
		for {
			select {
			case <-ctx.Done():
				return
			case msg, ok := <-ch:
				if !ok {
					return
				}
				var ev Event
				if err := json.Unmarshal([]byte(msg.Payload), &ev); err != nil {
					continue
				}
				fn(ev)
			}
		}
	}()
	return nil
}

// lazyStore opens the Redis store on first use, so nothing dials at mount
// (route:list and tests never dial) and a missing redis.url is a 503 on the
// first request, logged once, not a boot failure (cooler's rule).
type lazyStore struct {
	open    func(ctx context.Context) (Store, error)
	mu      sync.Mutex
	store   Store
	warned  atomic.Bool
	ready   atomic.Bool // the store is open
	opening atomic.Bool // a kick is opening it
}

// get returns the store, opening it on the first call and retrying on a
// later call after a failed open.
func (l *lazyStore) get(ctx context.Context) (Store, error) {
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.store != nil {
		return l.store, nil
	}
	s, err := l.open(ctx)
	if err != nil {
		if !l.warned.Swap(true) {
			logger().Error("pier: session store could not be opened", "error", err)
		}
		return nil, err
	}
	l.store = s
	l.ready.Store(true)
	return s, nil
}

// kick opens the store in the background when it is not open yet and no
// other kick is opening it; the caller never waits.
func (l *lazyStore) kick() {
	if l.ready.Load() || !l.opening.CompareAndSwap(false, true) {
		return
	}
	go func() {
		defer l.opening.Store(false)
		_, _ = l.get(context.Background())
	}()
}

// warm opens the store at mount in the background, retrying with a backoff
// from one second doubling to thirty, so the event subscription exists
// before traffic arrives (A-WR-05). It returns once the store is open or
// ctx ends; each failure is logged once per backoff step.
func (l *lazyStore) warm(ctx context.Context) {
	backoff := time.Second
	for {
		if _, err := l.get(ctx); err == nil {
			return
		} else if ctx.Err() == nil {
			logger().Warn("pier: session store not open yet, retrying", "retry_in", backoff.String(), "error", err)
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(backoff):
		}
		backoff = min(backoff*2, 30*time.Second)
	}
}

func (l *lazyStore) Put(ctx context.Context, id string, rec Record, ttl time.Duration) error {
	s, err := l.get(ctx)
	if err != nil {
		return err
	}
	return s.Put(ctx, id, rec, ttl)
}

func (l *lazyStore) Get(ctx context.Context, id string) (Record, bool, error) {
	s, err := l.get(ctx)
	if err != nil {
		return Record{}, false, err
	}
	return s.Get(ctx, id)
}

// Update forwards to the opened store's Update (Put for a store without one).
func (l *lazyStore) Update(ctx context.Context, id string, rec Record, ttl time.Duration) (bool, error) {
	s, err := l.get(ctx)
	if err != nil {
		return false, err
	}
	if u, ok := s.(updater); ok {
		return u.Update(ctx, id, rec, ttl)
	}
	return true, s.Put(ctx, id, rec, ttl)
}

// Reserve forwards to the opened store's Reserve (Count then Put for a
// store without one).
func (l *lazyStore) Reserve(ctx context.Context, binding, id string, rec Record, max int, ttl time.Duration) (bool, error) {
	s, err := l.get(ctx)
	if err != nil {
		return false, err
	}
	if r, ok := s.(reserver); ok {
		return r.Reserve(ctx, binding, id, rec, max, ttl)
	}
	n, err := s.Count(ctx, binding)
	if err != nil || n >= max {
		return false, err
	}
	return true, s.Put(ctx, id, rec, ttl)
}

// Release forwards to the opened store's Release (Delete for a store
// without one).
func (l *lazyStore) Release(ctx context.Context, binding, id string) error {
	s, err := l.get(ctx)
	if err != nil {
		return err
	}
	if r, ok := s.(reserver); ok {
		return r.Release(ctx, binding, id)
	}
	return s.Delete(ctx, id)
}

func (l *lazyStore) Touch(ctx context.Context, id string, ttl time.Duration) error {
	s, err := l.get(ctx)
	if err != nil {
		return err
	}
	return s.Touch(ctx, id, ttl)
}

func (l *lazyStore) Delete(ctx context.Context, id string) error {
	s, err := l.get(ctx)
	if err != nil {
		return err
	}
	return s.Delete(ctx, id)
}

func (l *lazyStore) Count(ctx context.Context, binding string) (int, error) {
	s, err := l.get(ctx)
	if err != nil {
		return 0, err
	}
	return s.Count(ctx, binding)
}

func (l *lazyStore) Publish(ctx context.Context, ev Event) error {
	s, err := l.get(ctx)
	if err != nil {
		return err
	}
	return s.Publish(ctx, ev)
}

func (l *lazyStore) Subscribe(ctx context.Context, fn func(Event)) error {
	s, err := l.get(ctx)
	if err != nil {
		return err
	}
	return s.Subscribe(ctx, fn)
}

// redisURLError names the key a redis store needs.
func redisURLError(url string) error {
	if strings.TrimSpace(url) == "" {
		return fmt.Errorf("pier: redis.url is required when %s is %s (set it in config/redis.yaml or SUMMER_REDIS__URL)", SessionStoreKey, StoreRedis)
	}
	return nil
}