From 986613b00bc2e4edd3eb00723d458bb616a5f457 Mon Sep 17 00:00:00 2001
From: narumi <toucans-cutouts0f@icloud.com>
Date: Tue, 18 Jun 2024 18:40:50 +0800
Subject: [PATCH] add DailyDataTracker util

---
 pkg/strategy/common/fee_budget.go      | 53 +++++++++-----------------
 pkg/strategy/common/fee_budget_test.go |  2 +-
 pkg/util/daily_data_tracker.go         | 35 +++++++++++++++++
 3 files changed, 53 insertions(+), 37 deletions(-)
 create mode 100644 pkg/util/daily_data_tracker.go

diff --git a/pkg/strategy/common/fee_budget.go b/pkg/strategy/common/fee_budget.go
index a2d141b43..d7eb0a5d0 100644
--- a/pkg/strategy/common/fee_budget.go
+++ b/pkg/strategy/common/fee_budget.go
@@ -2,29 +2,29 @@ package common
 
 import (
 	"sync"
-	"time"
 
 	"github.com/c9s/bbgo/pkg/fixedpoint"
 	"github.com/c9s/bbgo/pkg/types"
+	"github.com/c9s/bbgo/pkg/util"
 	log "github.com/sirupsen/logrus"
 )
 
 type FeeBudget struct {
 	DailyFeeBudgets map[string]fixedpoint.Value `json:"dailyFeeBudgets,omitempty"`
-	State           *State                      `persistence:"state"`
+	DailyFeeTracker *util.DailyDataTracker      `persistence:"dailyFeeTracker"`
 
 	mu sync.Mutex
 }
 
 func (f *FeeBudget) Initialize() {
-	if f.State == nil {
-		f.State = &State{}
-		f.State.Reset()
+	if f.DailyFeeTracker == nil {
+		f.DailyFeeTracker = &util.DailyDataTracker{}
+		f.DailyFeeTracker.Reset()
 	}
 
-	if f.State.IsOver24Hours() {
+	if f.DailyFeeTracker.IsOver24Hours() {
 		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
 	}
 
-	if f.State.AccumulatedFees == nil {
+	if f.DailyFeeTracker.Data == nil {
 		return true
 	}
 
-	if f.State.IsOver24Hours() {
-		f.State.Reset()
+	if f.DailyFeeTracker.IsOver24Hours() {
+		f.DailyFeeTracker.Reset()
 		return true
 	}
 
 	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 {
 				log.Warnf("[FeeBudget] accumulative fee %s exceeded the fee budget %s, skipping...", fee.String(), budget.String())
 				return false
@@ -57,36 +57,17 @@ func (f *FeeBudget) IsBudgetAllowed() bool {
 func (f *FeeBudget) HandleTradeUpdate(trade types.Trade) {
 	log.Infof("[FeeBudget] received trade %s", trade.String())
 
-	if f.State.IsOver24Hours() {
-		f.State.Reset()
+	if f.DailyFeeTracker.IsOver24Hours() {
+		f.DailyFeeTracker.Reset()
 	}
 
 	// safe check
-	if f.State.AccumulatedFees == nil {
+	if f.DailyFeeTracker.Data == nil {
 		f.mu.Lock()
-		f.State.AccumulatedFees = make(map[string]fixedpoint.Value)
+		f.DailyFeeTracker.Data = make(map[string]fixedpoint.Value)
 		f.mu.Unlock()
 	}
 
-	f.State.AccumulatedFees[trade.FeeCurrency] = f.State.AccumulatedFees[trade.FeeCurrency].Add(trade.Fee)
-	log.Infof("[FeeBudget] accumulated fee: %s %s", f.State.AccumulatedFees[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)
+	f.DailyFeeTracker.Data[trade.FeeCurrency] = f.DailyFeeTracker.Data[trade.FeeCurrency].Add(trade.Fee)
+	log.Infof("[FeeBudget] accumulated fee: %s %s", f.DailyFeeTracker.Data[trade.FeeCurrency].String(), trade.FeeCurrency)
 }
diff --git a/pkg/strategy/common/fee_budget_test.go b/pkg/strategy/common/fee_budget_test.go
index d57658b15..b0265b272 100644
--- a/pkg/strategy/common/fee_budget_test.go
+++ b/pkg/strategy/common/fee_budget_test.go
@@ -50,7 +50,7 @@ func TestFeeBudget(t *testing.T) {
 		assert.Equal(t, c.expected, feeBudget.IsBudgetAllowed())
 
 		// 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())
 	}
 }
diff --git a/pkg/util/daily_data_tracker.go b/pkg/util/daily_data_tracker.go
new file mode 100644
index 000000000..317f2c880
--- /dev/null
+++ b/pkg/util/daily_data_tracker.go
@@ -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()
+}