From 480346a65624533dd44c17d6d61f4766986248b4 Mon Sep 17 00:00:00 2001
From: forcodedancing <just.haha.it@gmail.com>
Date: Mon, 19 Jun 2023 14:00:56 +0800
Subject: [PATCH] add payment for bucket migration

---
 .../msg_server_complete_migrate_bucket.go     | 13 ++++++
 x/storage/keeper/payment.go                   | 40 ++++++++++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/x/storage/keeper/msg_server_complete_migrate_bucket.go b/x/storage/keeper/msg_server_complete_migrate_bucket.go
index 9cd357f34..e9712518c 100644
--- a/x/storage/keeper/msg_server_complete_migrate_bucket.go
+++ b/x/storage/keeper/msg_server_complete_migrate_bucket.go
@@ -42,10 +42,23 @@ func (k msgServer) CompleteMigrateBucket(goCtx context.Context, msg *types.MsgCo
 		return nil, virtualgroupmoduletypes.ErrGVGFamilyNotExist
 	}
 
+	oldBucketInfo := &types.BucketInfo{
+		PaymentAddress:             bucketInfo.PaymentAddress,
+		PrimarySpId:                bucketInfo.PrimarySpId,
+		GlobalVirtualGroupFamilyId: bucketInfo.GlobalVirtualGroupFamilyId,
+		ChargedReadQuota:           bucketInfo.ChargedReadQuota,
+		BillingInfo:                bucketInfo.BillingInfo,
+	}
+
 	bucketInfo.PrimarySpId = migrationBucketInfo.DstSpId
 	bucketInfo.GlobalVirtualGroupFamilyId = msg.GlobalVirtualGroupFamilyId
 	k.SetBucketInfo(ctx, bucketInfo)
 	k.DeleteMigrationBucketInfo(ctx, bucketInfo.Id)
 
+	err := k.ChargeBucketMigration(ctx, oldBucketInfo, bucketInfo)
+	if err != nil {
+		return nil, types.ErrMigtationBucketFailed.Wrapf("update payment info failed.")
+	}
+
 	return &types.MsgCompleteMigrateBucketResponse{}, nil
 }
diff --git a/x/storage/keeper/payment.go b/x/storage/keeper/payment.go
index 8ddde7810..96ba4e810 100644
--- a/x/storage/keeper/payment.go
+++ b/x/storage/keeper/payment.go
@@ -175,7 +175,8 @@ func (k Keeper) ChargeViaBucketChange(ctx sdk.Context, bucketInfo *storagetypes.
 	return nil
 }
 
-func (k Keeper) GetBucketBill(ctx sdk.Context, bucketInfo *storagetypes.BucketInfo) (userFlows types.UserFlows, err error) {
+func (k Keeper) GetBucketBill(ctx sdk.Context, bucketInfo *storagetypes.BucketInfo, settle ...bool) (userFlows types.UserFlows, err error) {
+	doSettle := settle != nil && settle[0]
 	userFlows.From = sdk.MustAccAddressFromHex(bucketInfo.PaymentAddress)
 	if bucketInfo.BillingInfo.TotalChargeSize == 0 && bucketInfo.ChargedReadQuota == 0 {
 		return userFlows, nil
@@ -197,6 +198,13 @@ func (k Keeper) GetBucketBill(ctx sdk.Context, bucketInfo *storagetypes.BucketIn
 		return userFlows, fmt.Errorf("get GVG family failed: %d", bucketInfo.GlobalVirtualGroupFamilyId)
 	}
 
+	if doSettle {
+		err := k.virtualGroupKeeper.SettleAndDistributeGVGFamily(ctx, primarySp.Id, gvgFamily)
+		if err != nil {
+			return userFlows, fmt.Errorf("settle GVG family failed: %d, err: %s", gvgFamily.Id, err.Error())
+		}
+	}
+
 	// primary sp total rate
 	primaryTotalFlowRate := price.ReadPrice.MulInt(sdkmath.NewIntFromUint64(bucketInfo.ChargedReadQuota)).TruncateInt()
 
@@ -220,6 +228,13 @@ func (k Keeper) GetBucketBill(ctx sdk.Context, bucketInfo *storagetypes.BucketIn
 			return userFlows, fmt.Errorf("get GVG failed: %d", lvg.GlobalVirtualGroupId)
 		}
 
+		if doSettle {
+			err := k.virtualGroupKeeper.SettleAndDistributeGVG(ctx, gvg)
+			if err != nil {
+				return userFlows, fmt.Errorf("settle GVG failed: %d, err: %s", gvg.Id, err.Error())
+			}
+		}
+
 		secondaryRate := price.SecondaryStorePrice.MulInt(sdkmath.NewIntFromUint64(lvgStoreSize.TotalChargeSize)).TruncateInt()
 		secondaryRate = secondaryRate.MulRaw(int64(len(gvg.SecondarySpIds)))
 		if secondaryRate.IsPositive() {
@@ -357,3 +372,26 @@ func (k Keeper) GetChargeSize(ctx sdk.Context, payloadSize uint64, ts int64) (si
 		return payloadSize, nil
 	}
 }
+
+func (k Keeper) ChargeBucketMigration(ctx sdk.Context, oldBucketInfo, newBucketInfo *storagetypes.BucketInfo) error {
+	// settle and get previous bill
+	prevBill, err := k.GetBucketBill(ctx, oldBucketInfo, true)
+	if err != nil {
+		return fmt.Errorf("settle and get bucket bill failed, bucket: %s, err: %s", oldBucketInfo.BucketName, err.Error())
+	}
+
+	newBucketInfo.BillingInfo.PriceTime = ctx.BlockTime().Unix()
+	// calculate new bill
+	newBill, err := k.GetBucketBill(ctx, newBucketInfo)
+	if err != nil {
+		return fmt.Errorf("get new bucket bill failed: %w", err)
+	}
+
+	// charge according to bill change
+	err = k.ChargeAccordingToBillChange(ctx, prevBill, newBill)
+	if err != nil {
+		ctx.Logger().Error("charge via bucket change failed", "err", err.Error())
+		return err
+	}
+	return nil
+}