aiaa-notification-server/internal/subscriber/tradesignal/store.go

package tradesignal

import (
	"context"
	"encoding/json"
	"sync"
	"time"

	"aiaa-notification-service/internal/cache"
)

const (
	positionKeyPrefix = "notify:position:"
	appliedKeyPrefix  = "notify:position:applied:"
	positionTTL       = 30 * 24 * time.Hour
	appliedTTL        = 7 * 24 * time.Hour
	storeTimeout      = 2 * time.Second
)

type persistedState struct {
	Avg  float64 `json:"avg"`
	Size float64 `json:"size"`
	Mode int     `json:"mode"`
}

type positionStore interface {
	load(key string) (*state, error)
	commit(key string, st *state, del bool, signalID string, snap Snapshot) error
	loadApplied(signalID string) (Snapshot, bool, error)
}

type memoryStore struct {
	mu        sync.Mutex
	positions map[string]*state
	applied   map[string]Snapshot
}

func newMemoryStore() *memoryStore {
	return &memoryStore{
		positions: make(map[string]*state),
		applied:   make(map[string]Snapshot),
	}
}

func (m *memoryStore) load(key string) (*state, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.positions[key], nil
}

func (m *memoryStore) loadApplied(signalID string) (Snapshot, bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	snap, ok := m.applied[signalID]
	return snap, ok, nil
}

func (m *memoryStore) commit(key string, st *state, del bool, signalID string, snap Snapshot) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if del {
		delete(m.positions, key)
	} else if st != nil {
		m.positions[key] = st
	}
	if signalID != "" {
		m.applied[signalID] = snap
	}
	return nil
}

type kvClient interface {
	GetRaw(ctx context.Context, key string) ([]byte, error)
	TxWrite(ctx context.Context, writes []cache.KVWrite) error
}

type redisStore struct {
	c kvClient
}

func newRedisStore(c *cache.Cache) positionStore {
	if c == nil {
		return newMemoryStore()
	}
	return &redisStore{c: c}
}

func (s *redisStore) ctx() (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), storeTimeout)
}

func (s *redisStore) load(key string) (*state, error) {
	ctx, cancel := s.ctx()
	defer cancel()
	b, err := s.c.GetRaw(ctx, positionKeyPrefix+key)
	if err != nil || len(b) == 0 {
		return nil, err
	}
	var p persistedState
	if err := json.Unmarshal(b, &p); err != nil {
		return nil, err
	}
	return &state{avg: p.Avg, size: p.Size, mode: mode(p.Mode)}, nil
}

func (s *redisStore) loadApplied(signalID string) (Snapshot, bool, error) {
	if signalID == "" {
		return Snapshot{}, false, nil
	}
	ctx, cancel := s.ctx()
	defer cancel()
	b, err := s.c.GetRaw(ctx, appliedKeyPrefix+signalID)
	if err != nil || len(b) == 0 {
		return Snapshot{}, false, err
	}
	var snap Snapshot
	if err := json.Unmarshal(b, &snap); err != nil {
		return Snapshot{}, false, err
	}
	return snap, true, nil
}

func (s *redisStore) commit(key string, st *state, del bool, signalID string, snap Snapshot) error {
	ctx, cancel := s.ctx()
	defer cancel()
	writes := make([]cache.KVWrite, 0, 2)
	posKey := positionKeyPrefix + key
	if del {
		writes = append(writes, cache.KVWrite{Key: posKey, Delete: true})
	} else if st != nil {
		b, err := json.Marshal(persistedState{Avg: st.avg, Size: st.size, Mode: int(st.mode)})
		if err != nil {
			return err
		}
		writes = append(writes, cache.KVWrite{Key: posKey, Val: b, TTL: positionTTL})
	}
	if signalID != "" {
		b, err := json.Marshal(snap)
		if err != nil {
			return err
		}
		writes = append(writes, cache.KVWrite{Key: appliedKeyPrefix + signalID, Val: b, TTL: appliedTTL})
	}
	return s.c.TxWrite(ctx, writes)
}