This commit is contained in:
なるみ 2024-09-19 09:41:05 +08:00 committed by GitHub
commit 1cb0dabb56
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 53 additions and 37 deletions

View File

@ -2,29 +2,29 @@ package common
import ( import (
"sync" "sync"
"time"
"github.com/c9s/bbgo/pkg/fixedpoint" "github.com/c9s/bbgo/pkg/fixedpoint"
"github.com/c9s/bbgo/pkg/types" "github.com/c9s/bbgo/pkg/types"
"github.com/c9s/bbgo/pkg/util"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
type FeeBudget struct { type FeeBudget struct {
DailyFeeBudgets map[string]fixedpoint.Value `json:"dailyFeeBudgets,omitempty"` DailyFeeBudgets map[string]fixedpoint.Value `json:"dailyFeeBudgets,omitempty"`
State *State `persistence:"state"` DailyFeeTracker *util.DailyDataTracker `persistence:"dailyFeeTracker"`
mu sync.Mutex mu sync.Mutex
} }
func (f *FeeBudget) Initialize() { func (f *FeeBudget) Initialize() {
if f.State == nil { if f.DailyFeeTracker == nil {
f.State = &State{} f.DailyFeeTracker = &util.DailyDataTracker{}
f.State.Reset() f.DailyFeeTracker.Reset()
} }
if f.State.IsOver24Hours() { if f.DailyFeeTracker.IsOver24Hours() {
log.Warn("[FeeBudget] state is over 24 hours, resetting to zero") log.Warn("[FeeBudget] state is over 24 hours, resetting to zero")
f.State.Reset() f.DailyFeeTracker.Reset()
} }
} }
@ -33,17 +33,17 @@ func (f *FeeBudget) IsBudgetAllowed() bool {
return true return true
} }
if f.State.AccumulatedFees == nil { if f.DailyFeeTracker.Data == nil {
return true return true
} }
if f.State.IsOver24Hours() { if f.DailyFeeTracker.IsOver24Hours() {
f.State.Reset() f.DailyFeeTracker.Reset()
return true return true
} }
for asset, budget := range f.DailyFeeBudgets { for asset, budget := range f.DailyFeeBudgets {
if fee, ok := f.State.AccumulatedFees[asset]; ok { if fee, ok := f.DailyFeeTracker.Data[asset]; ok {
if fee.Compare(budget) >= 0 { if fee.Compare(budget) >= 0 {
log.Warnf("[FeeBudget] accumulative fee %s exceeded the fee budget %s, skipping...", fee.String(), budget.String()) log.Warnf("[FeeBudget] accumulative fee %s exceeded the fee budget %s, skipping...", fee.String(), budget.String())
return false return false
@ -57,36 +57,17 @@ func (f *FeeBudget) IsBudgetAllowed() bool {
func (f *FeeBudget) HandleTradeUpdate(trade types.Trade) { func (f *FeeBudget) HandleTradeUpdate(trade types.Trade) {
log.Infof("[FeeBudget] received trade %s", trade.String()) log.Infof("[FeeBudget] received trade %s", trade.String())
if f.State.IsOver24Hours() { if f.DailyFeeTracker.IsOver24Hours() {
f.State.Reset() f.DailyFeeTracker.Reset()
} }
// safe check // safe check
if f.State.AccumulatedFees == nil { if f.DailyFeeTracker.Data == nil {
f.mu.Lock() f.mu.Lock()
f.State.AccumulatedFees = make(map[string]fixedpoint.Value) f.DailyFeeTracker.Data = make(map[string]fixedpoint.Value)
f.mu.Unlock() f.mu.Unlock()
} }
f.State.AccumulatedFees[trade.FeeCurrency] = f.State.AccumulatedFees[trade.FeeCurrency].Add(trade.Fee) f.DailyFeeTracker.Data[trade.FeeCurrency] = f.DailyFeeTracker.Data[trade.FeeCurrency].Add(trade.Fee)
log.Infof("[FeeBudget] accumulated fee: %s %s", f.State.AccumulatedFees[trade.FeeCurrency].String(), trade.FeeCurrency) log.Infof("[FeeBudget] accumulated fee: %s %s", f.DailyFeeTracker.Data[trade.FeeCurrency].String(), trade.FeeCurrency)
}
type State struct {
AccumulatedFeeStartedAt time.Time `json:"accumulatedFeeStartedAt,omitempty"`
AccumulatedFees map[string]fixedpoint.Value `json:"accumulatedFees,omitempty"`
}
func (s *State) IsOver24Hours() bool {
return time.Since(s.AccumulatedFeeStartedAt) >= 24*time.Hour
}
func (s *State) Reset() {
t := time.Now()
dateTime := time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, t.Location())
log.Infof("[State] resetting accumulated started time to: %s", dateTime)
s.AccumulatedFeeStartedAt = dateTime
s.AccumulatedFees = make(map[string]fixedpoint.Value)
} }

View File

@ -50,7 +50,7 @@ func TestFeeBudget(t *testing.T) {
assert.Equal(t, c.expected, feeBudget.IsBudgetAllowed()) assert.Equal(t, c.expected, feeBudget.IsBudgetAllowed())
// test reset // test reset
feeBudget.State.AccumulatedFeeStartedAt = feeBudget.State.AccumulatedFeeStartedAt.Add(-24 * time.Hour) feeBudget.DailyFeeTracker.StartedAt = feeBudget.DailyFeeTracker.StartedAt.Add(-24 * time.Hour)
assert.True(t, feeBudget.IsBudgetAllowed()) assert.True(t, feeBudget.IsBudgetAllowed())
} }
} }

View File

@ -0,0 +1,35 @@
package util
import (
"time"
"github.com/c9s/bbgo/pkg/types"
log "github.com/sirupsen/logrus"
)
type DailyDataTracker struct {
StartedAt time.Time `json:"startedAt,omitempty"`
Data types.ValueMap `json:"data,omitempty"`
}
func (d *DailyDataTracker) IsOver24Hours() bool {
return time.Since(d.StartedAt) >= 24*time.Hour
}
func (d *DailyDataTracker) ResetTime() {
t := time.Now()
dateTime := time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, t.Location())
log.Infof("[Countdown] resetting accumulated started time to: %s", dateTime)
d.StartedAt = dateTime
}
func (d *DailyDataTracker) ResetData() {
d.Data = make(types.ValueMap)
}
func (d *DailyDataTracker) Reset() {
d.ResetTime()
d.ResetData()
}