From 7134f51d38d4967b8a050c8476dee5e82f8057a7 Mon Sep 17 00:00:00 2001
From: kbearXD <s199180626@gmail.com>
Date: Wed, 22 May 2024 11:37:13 +0800
Subject: [PATCH] FEATURE: [dca2] change state recovery logic

---
 pkg/exchange/retry/account.go    | 43 ++++++++++++++++++++++++++++++++
 pkg/exchange/retry/order.go      | 12 ---------
 pkg/strategy/dca2/recover.go     | 14 ++++-------
 pkg/strategy/dca2/take_profit.go | 22 ++++++++++++++++
 4 files changed, 70 insertions(+), 21 deletions(-)
 create mode 100644 pkg/exchange/retry/account.go

diff --git a/pkg/exchange/retry/account.go b/pkg/exchange/retry/account.go
new file mode 100644
index 000000000..1e142c034
--- /dev/null
+++ b/pkg/exchange/retry/account.go
@@ -0,0 +1,43 @@
+package retry
+
+import (
+	"context"
+
+	"github.com/c9s/bbgo/pkg/types"
+)
+
+func QueryAccountUntilSuccessful(
+	ctx context.Context, ex types.ExchangeAccountService,
+) (account *types.Account, err error) {
+	var op = func() (err2 error) {
+		account, err2 = ex.QueryAccount(ctx)
+		return err2
+	}
+
+	err = GeneralBackoff(ctx, op)
+	return account, err
+}
+
+func QueryAccountBalancesUntilSuccessful(
+	ctx context.Context, ex types.ExchangeAccountService,
+) (bals types.BalanceMap, err error) {
+	var op = func() (err2 error) {
+		bals, err2 = ex.QueryAccountBalances(ctx)
+		return err2
+	}
+
+	err = GeneralBackoff(ctx, op)
+	return bals, err
+}
+
+func QueryAccountBalancesUntilSuccessfulLite(
+	ctx context.Context, ex types.ExchangeAccountService,
+) (bals types.BalanceMap, err error) {
+	var op = func() (err2 error) {
+		bals, err2 = ex.QueryAccountBalances(ctx)
+		return err2
+	}
+
+	err = GeneralLiteBackoff(ctx, op)
+	return bals, err
+}
diff --git a/pkg/exchange/retry/order.go b/pkg/exchange/retry/order.go
index d27f297ef..00a0a7fac 100644
--- a/pkg/exchange/retry/order.go
+++ b/pkg/exchange/retry/order.go
@@ -180,18 +180,6 @@ func QueryOrderTradesUntilSuccessfulLite(
 	return trades, err
 }
 
-func QueryAccountUntilSuccessful(
-	ctx context.Context, ex types.ExchangeAccountService,
-) (account *types.Account, err error) {
-	var op = func() (err2 error) {
-		account, err2 = ex.QueryAccount(ctx)
-		return err2
-	}
-
-	err = GeneralBackoff(ctx, op)
-	return account, err
-}
-
 func QueryOrderUntilSuccessful(
 	ctx context.Context, query types.ExchangeOrderQueryService, opts types.OrderQuery,
 ) (order *types.Order, err error) {
diff --git a/pkg/strategy/dca2/recover.go b/pkg/strategy/dca2/recover.go
index 23f72c9b8..f28abb154 100644
--- a/pkg/strategy/dca2/recover.go
+++ b/pkg/strategy/dca2/recover.go
@@ -113,13 +113,13 @@ func recoverState(ctx context.Context, maxOrderCount int, currentRound Round, or
 		}
 	}
 
-	// all open-position orders are still not filled -> OpenPositionReady
-	if filledCnt == 0 && cancelledCnt == 0 {
+	// no order is filled -> OpenPositionReady
+	if filledCnt == 0 {
 		return OpenPositionReady, nil
 	}
 
 	// there are at least one open-position orders filled
-	if filledCnt > 0 && cancelledCnt == 0 {
+	if cancelledCnt == 0 {
 		if openedCnt > 0 {
 			return OpenPositionOrderFilled, nil
 		} else {
@@ -128,12 +128,8 @@ func recoverState(ctx context.Context, maxOrderCount int, currentRound Round, or
 		}
 	}
 
-	// there are at last one open-position orders cancelled ->
-	if cancelledCnt > 0 {
-		return OpenPositionOrdersCancelling, nil
-	}
-
-	return None, fmt.Errorf("unexpected order status combination (opened, filled, cancelled) = (%d, %d, %d)", openedCnt, filledCnt, cancelledCnt)
+	// there are at last one open-position orders cancelled and at least one filled order -> open position order cancelling
+	return OpenPositionOrdersCancelling, nil
 }
 
 func recoverPosition(ctx context.Context, position *types.Position, currentRound Round, queryService types.ExchangeOrderQueryService) error {
diff --git a/pkg/strategy/dca2/take_profit.go b/pkg/strategy/dca2/take_profit.go
index a21ebbafc..a5255bb45 100644
--- a/pkg/strategy/dca2/take_profit.go
+++ b/pkg/strategy/dca2/take_profit.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/c9s/bbgo/pkg/exchange/retry"
 	"github.com/c9s/bbgo/pkg/fixedpoint"
 	"github.com/c9s/bbgo/pkg/types"
 	"github.com/pkg/errors"
@@ -39,6 +40,27 @@ func (s *Strategy) placeTakeProfitOrders(ctx context.Context) error {
 	s.logger.Infof("position of this round before place the take-profit order: %s", roundPosition.String())
 
 	order := generateTakeProfitOrder(s.Market, s.TakeProfitRatio, roundPosition, s.OrderGroupID)
+
+	// verify the volume of order
+	bals, err := retry.QueryAccountBalancesUntilSuccessfulLite(ctx, s.ExchangeSession.Exchange)
+	if err != nil {
+		return errors.Wrapf(err, "failed to query balance to verify")
+	}
+
+	bal, exist := bals[s.Market.BaseCurrency]
+	if !exist {
+		return fmt.Errorf("there is no %s in the balances %+v", s.Market.BaseCurrency, bals)
+	}
+
+	quantityDiff := bal.Available.Sub(order.Quantity)
+	if quantityDiff.Sign() < 0 {
+		return fmt.Errorf("the balance (%s) is not enough for the order (%s)", bal.String(), order.Quantity.String())
+	}
+
+	if quantityDiff.Compare(s.Market.MinQuantity) > 0 {
+		s.logger.Warnf("the diff between balance (%s) and the take-profit order (%s) is larger than min quantity %s", bal.String(), order.Quantity.String(), s.Market.MinQuantity.String())
+	}
+
 	createdOrders, err := s.OrderExecutor.SubmitOrders(ctx, order)
 	if err != nil {
 		return err