From a6047c48405a33a17f74318c775f2c780096800d Mon Sep 17 00:00:00 2001
From: c9s <yoanlin93@gmail.com>
Date: Mon, 20 Feb 2023 16:52:39 +0800
Subject: [PATCH] grid2: implement CleanUp interface

---
 pkg/strategy/grid2/strategy.go | 38 ++++++++++++++++++++++++++--------
 1 file changed, 29 insertions(+), 9 deletions(-)

diff --git a/pkg/strategy/grid2/strategy.go b/pkg/strategy/grid2/strategy.go
index 69db5fea8..bd055e6c4 100644
--- a/pkg/strategy/grid2/strategy.go
+++ b/pkg/strategy/grid2/strategy.go
@@ -145,6 +145,9 @@ type Strategy struct {
 	GridProfitStats *GridProfitStats `persistence:"grid_profit_stats"`
 	Position        *types.Position  `persistence:"position"`
 
+	// ExchangeSession is an injection field
+	ExchangeSession *bbgo.ExchangeSession
+
 	grid              *Grid
 	session           *bbgo.ExchangeSession
 	orderQueryService types.ExchangeOrderQueryService
@@ -200,6 +203,17 @@ func (s *Strategy) Validate() error {
 }
 
 func (s *Strategy) Defaults() error {
+	if s.LogFields == nil {
+		s.LogFields = logrus.Fields{}
+	}
+
+	s.LogFields["symbol"] = s.Symbol
+	s.LogFields["strategy"] = ID
+	return nil
+}
+
+func (s *Strategy) Initialize() error {
+	s.logger = log.WithFields(s.LogFields)
 	return nil
 }
 
@@ -1392,6 +1406,20 @@ func (s *Strategy) newPrometheusLabels() prometheus.Labels {
 	return mergeLabels(s.PrometheusLabels, labels)
 }
 
+func (s *Strategy) CleanUp(ctx context.Context) error {
+	if s.ExchangeSession == nil {
+		return errors.New("ExchangeSession is nil, can not clean up")
+	}
+
+	openOrders, err := s.ExchangeSession.Exchange.QueryOpenOrders(ctx, s.Symbol)
+	if err != nil {
+		return err
+	}
+
+	err = s.ExchangeSession.Exchange.CancelOrders(ctx, openOrders...)
+	return errors.Wrapf(err, "can not cancel %s orders", s.Symbol)
+}
+
 func (s *Strategy) Run(ctx context.Context, _ bbgo.OrderExecutor, session *bbgo.ExchangeSession) error {
 	instanceID := s.InstanceID()
 
@@ -1401,14 +1429,6 @@ func (s *Strategy) Run(ctx context.Context, _ bbgo.OrderExecutor, session *bbgo.
 		s.orderQueryService = service
 	}
 
-	if s.LogFields == nil {
-		s.LogFields = logrus.Fields{}
-	}
-
-	s.LogFields["symbol"] = s.Symbol
-	s.LogFields["strategy"] = ID
-	s.logger = log.WithFields(s.LogFields)
-
 	if s.OrderGroupID == 0 {
 		s.OrderGroupID = util.FNV32(instanceID)
 	}
@@ -1592,4 +1612,4 @@ func (s *Strategy) openOrdersMismatches(ctx context.Context, session *bbgo.Excha
 	}
 
 	return false, nil
-}
\ No newline at end of file
+}