bbgo: add TradeStore prune func and its tests

This commit is contained in:
c9s 2022-12-06 00:15:09 +08:00
parent 6408224663
commit beb862be44
2 changed files with 57 additions and 1 deletions

View File

@ -8,12 +8,14 @@ import (
)
const TradeExpiryTime = 24 * time.Hour
const PruneTriggerNumOfTrades = 10_000
type TradeStore struct {
// any created trades for tracking trades
sync.Mutex
trades map[uint64]types.Trade
trades map[uint64]types.Trade
lastTradeTime time.Time
}
func NewTradeStore() *TradeStore {
@ -96,6 +98,13 @@ func (s *TradeStore) Add(trades ...types.Trade) {
for _, trade := range trades {
s.trades[trade.ID] = trade
s.touchLastTradeTime(trade)
}
}
func (s *TradeStore) touchLastTradeTime(trade types.Trade) {
if trade.Time.Time().After(s.lastTradeTime) {
s.lastTradeTime = trade.Time.Time()
}
}
@ -122,8 +131,16 @@ func (s *TradeStore) Prune(curTime time.Time) {
s.pruneExpiredTrades(curTime)
}
func (s *TradeStore) isCoolTrade(trade types.Trade) bool {
// if the time of last trade is over 1 hour, we call it's cool trade
return s.lastTradeTime != (time.Time{}) && time.Time(trade.Time).Sub(s.lastTradeTime) > time.Hour
}
func (s *TradeStore) BindStream(stream types.Stream) {
stream.OnTradeUpdate(func(trade types.Trade) {
s.Add(trade)
if s.isCoolTrade(trade) {
s.Prune(time.Time(trade.Time))
}
})
}

View File

@ -0,0 +1,39 @@
package bbgo
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/c9s/bbgo/pkg/types"
)
func TestTradeStore_isCoolTrade(t *testing.T) {
now := time.Now()
store := NewTradeStore()
store.lastTradeTime = now.Add(-2 * time.Hour)
ok := store.isCoolTrade(types.Trade{
Time: types.Time(now),
})
assert.True(t, ok)
store.lastTradeTime = now.Add(-2 * time.Minute)
ok = store.isCoolTrade(types.Trade{
Time: types.Time(now),
})
assert.False(t, ok)
}
func TestTradeStore_Prune(t *testing.T) {
now := time.Now()
store := NewTradeStore()
store.Add(
types.Trade{ID: 1, Time: types.Time(now.Add(-25 * time.Hour))},
types.Trade{ID: 2, Time: types.Time(now.Add(-23 * time.Hour))},
types.Trade{ID: 3, Time: types.Time(now.Add(-2 * time.Minute))},
types.Trade{ID: 4, Time: types.Time(now.Add(-1 * time.Minute))},
)
store.Prune(now)
assert.Equal(t, 3, len(store.trades))
}