diff --git a/cmd/server/main.go b/cmd/server/main.go index 05dfdb5..972895e 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -231,7 +231,7 @@ func main() { deduper := subscriber.NewCacheDeduper(redisCache, cfg.SubscriptionDedupTTL) for _, sub := range cfg.ActiveSubscriptions() { sub := sub - cons, err := subscriber.New(sub, lookup, notifySvc.Process, deduper) + cons, err := subscriber.New(sub, lookup, notifySvc.Process, deduper, redisCache) if err != nil { slog.Error("subscriber init", "name", sub.Name, "error", err) os.Exit(1) diff --git a/internal/cache/kv.go b/internal/cache/kv.go new file mode 100644 index 0000000..3f1c770 --- /dev/null +++ b/internal/cache/kv.go @@ -0,0 +1,46 @@ +package cache + +import ( + "context" + "fmt" + "time" + + "github.com/redis/go-redis/v9" +) + +type KVWrite struct { + Key string + Val []byte + TTL time.Duration + Delete bool +} + +func (c *Cache) GetRaw(ctx context.Context, key string) ([]byte, error) { + if c == nil || c.rdb == nil { + return nil, fmt.Errorf("redis unavailable") + } + b, err := c.rdb.Get(ctx, key).Bytes() + if err == redis.Nil { + return nil, nil + } + return b, err +} + +func (c *Cache) TxWrite(ctx context.Context, writes []KVWrite) error { + if c == nil || c.rdb == nil { + return fmt.Errorf("redis unavailable") + } + if len(writes) == 0 { + return nil + } + pipe := c.rdb.TxPipeline() + for _, w := range writes { + if w.Delete { + pipe.Del(ctx, w.Key) + continue + } + pipe.Set(ctx, w.Key, w.Val, w.TTL) + } + _, err := pipe.Exec(ctx) + return err +} diff --git a/internal/subscriber/handle.go b/internal/subscriber/handle.go index 3e72a21..6fe9416 100644 --- a/internal/subscriber/handle.go +++ b/internal/subscriber/handle.go @@ -8,6 +8,7 @@ import ( "aiaa-notification-service/internal/model" "aiaa-notification-service/internal/notify" + "aiaa-notification-service/internal/subscriber/tradesignal" ) const retryHeader = "x-retry-count" @@ -82,6 +83,10 @@ func HandleMessage(ctx context.Context, in HandleInput, conv MessageConverter, l event, data, err := conv.Convert(in.Body) if err != nil { hash := MessageHash(in.Body) + if errors.Is(err, tradesignal.ErrPositionStore) { + slog.Warn("position store failed, retry", "hash", hash, "raw", raw, "error", errText(err)) + return DecideRetry(RetryCount(in.Headers), in.MaxRetry) + } slog.Warn("invalid signal, ack", "hash", hash, "raw", raw, "error", errText(err)) return DispositionAck } diff --git a/internal/subscriber/handle_test.go b/internal/subscriber/handle_test.go index 68b02b2..79cf199 100644 --- a/internal/subscriber/handle_test.go +++ b/internal/subscriber/handle_test.go @@ -108,6 +108,28 @@ func TestHandleProcessErrorRetryThenDLQ(t *testing.T) { } } +func TestHandlePositionStoreErrorRetries(t *testing.T) { + d := HandleMessage(context.Background(), HandleInput{ + Body: []byte(`{"action":"OPEN"}`), SourceName: "s", MaxRetry: 3, + }, stubConverter{err: tradesignal.ErrPositionStore}, + func(context.Context, string) (*model.Source, error) { return enabledSrc(), nil }, + func(context.Context, notify.Request) (notify.Result, error) { + t.Fatal("process should not run") + return notify.Result{}, nil + }) + if d != DispositionRetry { + t.Fatalf("%v", d) + } +} + +type stubConverter struct { + err error +} + +func (s stubConverter) Convert([]byte) (string, map[string]interface{}, error) { + return "", nil, s.err +} + func TestHandleSuccessAckPassesEventAndFormatted(t *testing.T) { var got notify.Request d := HandleMessage(context.Background(), HandleInput{ diff --git a/internal/subscriber/subscriber.go b/internal/subscriber/subscriber.go index 7995aef..f032f0d 100644 --- a/internal/subscriber/subscriber.go +++ b/internal/subscriber/subscriber.go @@ -7,6 +7,7 @@ import ( "net/url" "time" + "aiaa-notification-service/internal/cache" "aiaa-notification-service/internal/config" "aiaa-notification-service/internal/subscriber/cryptostrategy" "aiaa-notification-service/internal/subscriber/tradesignal" @@ -22,11 +23,11 @@ type Subscriber struct { deduper Deduper } -func New(cfg config.SubscriptionConfig, lookup SourceLookup, process ProcessFunc, deduper Deduper) (*Subscriber, error) { +func New(cfg config.SubscriptionConfig, lookup SourceLookup, process ProcessFunc, deduper Deduper, redisCache *cache.Cache) (*Subscriber, error) { var conv MessageConverter switch cfg.Formatter { case "trade_signal": - conv = tradesignal.NewConverter(cfg.StrategyOverrides) + conv = tradesignal.NewConverterWithCache(cfg.StrategyOverrides, redisCache) case "crypto_strategy": conv = cryptostrategy.NewConverter() default: diff --git a/internal/subscriber/tradesignal/convert.go b/internal/subscriber/tradesignal/convert.go index 6e169c0..1bed335 100644 --- a/internal/subscriber/tradesignal/convert.go +++ b/internal/subscriber/tradesignal/convert.go @@ -6,10 +6,12 @@ import ( "fmt" "strings" + "aiaa-notification-service/internal/cache" "aiaa-notification-service/internal/config" ) var ErrInvalidSignal = errors.New("invalid signal") +var ErrPositionStore = errors.New("position store") type Converter struct { overrides map[string]config.StrategyOverride @@ -17,9 +19,13 @@ type Converter struct { } func NewConverter(overrides map[string]config.StrategyOverride) *Converter { + return NewConverterWithCache(overrides, nil) +} + +func NewConverterWithCache(overrides map[string]config.StrategyOverride, c *cache.Cache) *Converter { return &Converter{ overrides: overrides, - positions: NewTracker(), + positions: NewTrackerWithStore(newRedisStore(c)), } } @@ -39,7 +45,10 @@ func (c *Converter) Convert(body []byte) (string, map[string]interface{}, error) return "trade.message", data, nil } out := Apply(&sig, c.overrideFor(sig.StrategyCode)) - snap := c.positions.Apply(out) + snap, err := c.positions.Apply(out) + if err != nil { + return "", nil, fmt.Errorf("%w: %v", ErrPositionStore, err) + } var opts FormatOptions if snap.HasAvg { avg := snap.AvgPrice diff --git a/internal/subscriber/tradesignal/position.go b/internal/subscriber/tradesignal/position.go index 5f5563f..1ec696f 100644 --- a/internal/subscriber/tradesignal/position.go +++ b/internal/subscriber/tradesignal/position.go @@ -14,9 +14,9 @@ const ( ) type Snapshot struct { - AvgPrice float64 - Size float64 - HasAvg bool + AvgPrice float64 `json:"avgPrice"` + Size float64 `json:"size"` + HasAvg bool `json:"hasAvg"` } type state struct { @@ -26,71 +26,75 @@ type state struct { } type Tracker struct { - mu sync.Mutex - positions map[string]*state - applied map[string]Snapshot + mu sync.Mutex + store positionStore } func NewTracker() *Tracker { - return &Tracker{ - positions: make(map[string]*state), - applied: make(map[string]Snapshot), - } + return NewTrackerWithStore(newMemoryStore()) } -func (t *Tracker) Apply(signal *Signal) Snapshot { +func NewTrackerWithStore(store positionStore) *Tracker { + if store == nil { + store = newMemoryStore() + } + return &Tracker{store: store} +} + +func (t *Tracker) Apply(signal *Signal) (Snapshot, error) { if signal == nil { - return Snapshot{} + return Snapshot{}, nil } t.mu.Lock() defer t.mu.Unlock() if signal.SignalID != "" { - if snap, ok := t.applied[signal.SignalID]; ok { - return snap + if snap, ok, err := t.store.loadApplied(signal.SignalID); err != nil { + return Snapshot{}, err + } else if ok { + return snap, nil } } key := positionKey(signal.StrategyCode, signal.Symbol, signal.Side) action := strings.ToUpper(signal.Action) - st := t.positions[key] + st, err := t.store.load(key) + if err != nil { + return Snapshot{}, err + } var snap Snapshot + del := false switch action { case "OPEN": st = openPosition(signal) snap = snapshotFrom(st) - if st != nil { - t.positions[key] = st - } else { - delete(t.positions, key) + if st == nil { + del = true } case "ADD": st = addPosition(st, signal) snap = snapshotFrom(st) - if st != nil { - t.positions[key] = st - } case "REDUCE": snap = snapshotFrom(st) st = reducePosition(st, signal) if st == nil || st.size <= 0 { - delete(t.positions, key) - } else { - t.positions[key] = st + del = true + st = nil } case "CLOSE": snap = snapshotFrom(st) - delete(t.positions, key) + del = true + st = nil default: snap = snapshotFrom(st) } - if signal.SignalID != "" { - t.applied[signal.SignalID] = snap + if err := t.store.commit(key, st, del, signal.SignalID, snap); err != nil { + return Snapshot{}, err } - return snap + return snap, nil } func openPosition(signal *Signal) *state { diff --git a/internal/subscriber/tradesignal/store.go b/internal/subscriber/tradesignal/store.go new file mode 100644 index 0000000..f8eb9ca --- /dev/null +++ b/internal/subscriber/tradesignal/store.go @@ -0,0 +1,145 @@ +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) +} diff --git a/internal/subscriber/tradesignal/store_test.go b/internal/subscriber/tradesignal/store_test.go new file mode 100644 index 0000000..ae264fa --- /dev/null +++ b/internal/subscriber/tradesignal/store_test.go @@ -0,0 +1,117 @@ +package tradesignal + +import ( + "context" + "errors" + "math" + "sync" + "testing" + + "aiaa-notification-service/internal/cache" +) + +type fakeKV struct { + mu sync.Mutex + data map[string][]byte + fail bool +} + +func newFakeKV() *fakeKV { + return &fakeKV{data: make(map[string][]byte)} +} + +func (f *fakeKV) GetRaw(_ context.Context, key string) ([]byte, error) { + if f.fail { + return nil, errors.New("redis down") + } + f.mu.Lock() + defer f.mu.Unlock() + if b, ok := f.data[key]; ok { + return append([]byte(nil), b...), nil + } + return nil, nil +} + +func (f *fakeKV) TxWrite(_ context.Context, writes []cache.KVWrite) error { + if f.fail { + return errors.New("redis down") + } + f.mu.Lock() + defer f.mu.Unlock() + for _, w := range writes { + if w.Delete { + delete(f.data, w.Key) + continue + } + f.data[w.Key] = append([]byte(nil), w.Val...) + } + return nil +} + +func TestTrackerPersistsAcrossMemoryInstances(t *testing.T) { + store := newMemoryStore() + mustApply(t, NewTrackerWithStore(store), &Signal{ + SignalID: "p1", StrategyCode: "BLONG", Symbol: "PEPEUSDT", Side: "SHORT", + Action: "OPEN", Quantity: ptr(10), Price: 0.00000059, + }) + snap := mustApply(t, NewTrackerWithStore(store), &Signal{ + SignalID: "p2", StrategyCode: "BLONG", Symbol: "PEPEUSDT", Side: "SHORT", + Action: "ADD", Quantity: ptr(10), Price: 0.00000061, + }) + if math.Abs(snap.AvgPrice-0.0000006) > 1e-12 { + t.Fatalf("shared memory store avg=%v", snap.AvgPrice) + } +} + +func TestTrackerPersistsAcrossRedisInstances(t *testing.T) { + store := &redisStore{c: newFakeKV()} + mustApply(t, NewTrackerWithStore(store), &Signal{ + SignalID: "r1", StrategyCode: "BLONG", Symbol: "PEPEUSDT", Side: "SHORT", + Action: "OPEN", Quantity: ptr(10), Price: 100, + }) + snap := mustApply(t, NewTrackerWithStore(store), &Signal{ + SignalID: "r2", StrategyCode: "BLONG", Symbol: "PEPEUSDT", Side: "SHORT", + Action: "ADD", Quantity: ptr(10), Price: 200, + }) + if math.Abs(snap.AvgPrice-150) > 1e-9 { + t.Fatalf("shared redis store avg=%v", snap.AvgPrice) + } +} + +func TestTrackerRedisIdempotentAcrossInstances(t *testing.T) { + store := &redisStore{c: newFakeKV()} + open := &Signal{ + SignalID: "same", StrategyCode: "BLONG", Symbol: "BTCUSDT", Side: "LONG", + Action: "OPEN", Quantity: ptr(1), Price: 100, + } + mustApply(t, NewTrackerWithStore(store), open) + mustApply(t, NewTrackerWithStore(store), open) + snap := mustApply(t, NewTrackerWithStore(store), &Signal{ + SignalID: "add", StrategyCode: "BLONG", Symbol: "BTCUSDT", Side: "LONG", + Action: "ADD", Quantity: ptr(1), Price: 200, + }) + if math.Abs(snap.AvgPrice-150) > 1e-9 { + t.Fatalf("replayed open should not double size, avg=%v", snap.AvgPrice) + } +} + +func TestTrackerRedisStoreError(t *testing.T) { + tr := NewTrackerWithStore(&redisStore{c: &fakeKV{fail: true, data: map[string][]byte{}}}) + _, err := tr.Apply(&Signal{ + SignalID: "e1", StrategyCode: "BLONG", Symbol: "BTCUSDT", Side: "LONG", + Action: "OPEN", Quantity: ptr(1), Price: 100, + }) + if err == nil { + t.Fatal("expected store error") + } +} + +func TestConvertPositionStoreError(t *testing.T) { + c := &Converter{ + positions: NewTrackerWithStore(&redisStore{c: &fakeKV{fail: true, data: map[string][]byte{}}}), + } + _, _, err := c.Convert([]byte(`{"action":"OPEN","symbol":"BTCUSDT","price":1,"quantity":1}`)) + if !errors.Is(err, ErrPositionStore) { + t.Fatalf("err=%v", err) + } +} diff --git a/internal/subscriber/tradesignal/transform_test.go b/internal/subscriber/tradesignal/transform_test.go index debac58..c383605 100644 --- a/internal/subscriber/tradesignal/transform_test.go +++ b/internal/subscriber/tradesignal/transform_test.go @@ -77,6 +77,15 @@ func TestApplyDoesNotChangeMarginRatioOnlySignals(t *testing.T) { } } +func mustApply(t *testing.T, tr *Tracker, signal *Signal) Snapshot { + t.Helper() + snap, err := tr.Apply(signal) + if err != nil { + t.Fatal(err) + } + return snap +} + func TestAvgPriceOpenAndAdd(t *testing.T) { tr := NewTracker() @@ -89,7 +98,7 @@ func TestAvgPriceOpenAndAdd(t *testing.T) { Quantity: ptr(2), Price: 100, } - snap := tr.Apply(open) + snap := mustApply(t, tr, open) if !snap.HasAvg || snap.AvgPrice != 100 { t.Fatalf("open avg=%v has=%v", snap.AvgPrice, snap.HasAvg) } @@ -103,7 +112,7 @@ func TestAvgPriceOpenAndAdd(t *testing.T) { Quantity: ptr(2), Price: 200, } - snap = tr.Apply(add) + snap = mustApply(t, tr, add) if !snap.HasAvg || math.Abs(snap.AvgPrice-150) > 1e-9 { t.Fatalf("expected avg 150, got %v", snap.AvgPrice) } @@ -121,7 +130,7 @@ func TestAvgPriceWithMarginRatio(t *testing.T) { AmountMarginRatio: ptr(0.1), Price: 100, } - tr.Apply(open) + mustApply(t, tr, open) add := &Signal{ SignalID: "r2", @@ -132,7 +141,7 @@ func TestAvgPriceWithMarginRatio(t *testing.T) { AmountMarginRatio: ptr(0.1), Price: 200, } - snap := tr.Apply(add) + snap := mustApply(t, tr, add) if !snap.HasAvg || math.Abs(snap.AvgPrice-150) > 1e-9 { t.Fatalf("expected weighted avg 150, got %v", snap.AvgPrice) } @@ -140,7 +149,7 @@ func TestAvgPriceWithMarginRatio(t *testing.T) { func TestCloseKeepsEntryAvgInSnapshot(t *testing.T) { tr := NewTracker() - tr.Apply(&Signal{ + mustApply(t, tr, &Signal{ SignalID: "c1", StrategyCode: "BLONG", Symbol: "BTCUSDT", @@ -150,7 +159,7 @@ func TestCloseKeepsEntryAvgInSnapshot(t *testing.T) { Price: 64000, }) - snap := tr.Apply(&Signal{ + snap := mustApply(t, tr, &Signal{ SignalID: "c2", StrategyCode: "BLONG", Symbol: "BTCUSDT", @@ -162,7 +171,7 @@ func TestCloseKeepsEntryAvgInSnapshot(t *testing.T) { t.Fatalf("close should report entry avg 64000, got %v", snap.AvgPrice) } - snap = tr.Apply(&Signal{ + snap = mustApply(t, tr, &Signal{ SignalID: "c3", StrategyCode: "BLONG", Symbol: "BTCUSDT", @@ -187,10 +196,10 @@ func TestSignalIDIdempotent(t *testing.T) { Quantity: ptr(1), Price: 100, } - tr.Apply(sig) - tr.Apply(sig) + mustApply(t, tr, sig) + mustApply(t, tr, sig) - snap := tr.Apply(&Signal{ + snap := mustApply(t, tr, &Signal{ SignalID: "dup2", StrategyCode: "BLONG", Symbol: "BTCUSDT", @@ -206,11 +215,11 @@ func TestSignalIDIdempotent(t *testing.T) { func TestDifferentSideIsolated(t *testing.T) { tr := NewTracker() - tr.Apply(&Signal{ + mustApply(t, tr, &Signal{ SignalID: "l1", StrategyCode: "BLONG", Symbol: "BTCUSDT", Side: "LONG", Action: "OPEN", Quantity: ptr(1), Price: 100, }) - snap := tr.Apply(&Signal{ + snap := mustApply(t, tr, &Signal{ SignalID: "s1", StrategyCode: "BLONG", Symbol: "BTCUSDT", Side: "SHORT", Action: "OPEN", Quantity: ptr(1), Price: 200, })