Updating of balaces added after cancelling leases to disabled aliases.

pkg/state/balances.go

@@ -339,6 +339,31 @@ func (s *balances) cancelInvalidLeaseIns(correctLeaseIns map[proto.Address]int64
 	return nil
 }
 
+func (s *balances) cancelLeases(changes map[proto.Address]balanceDiff, blockID proto.BlockID) error {
+	zap.S().Infof("Updating balances for cancelled leases")
+	for a, bd := range changes {
+		k := wavesBalanceKey{address: a}
+		r, err := s.newestWavesRecord(k.bytes(), true)
+		if err != nil {
+			return err
+		}
+		profile := r.balanceProfile
+		newProfile, err := bd.applyTo(&profile)
+		if err != nil {
+			return err
+		}
+		val := &wavesValue{leaseChange: true, profile: *newProfile}
+		if err := s.setWavesBalance(a, val, blockID); err != nil {
+			return err
+		}
+		zap.S().Infof("Balance of %s changed from (B: %d, LIn: %d, LOut: %d) to (B: %d, lIn: %d, lOut: %d)",
+			a.String(), profile.balance, profile.leaseIn, profile.leaseOut,
+			newProfile.balance, newProfile.leaseIn, newProfile.leaseOut)
+	}
+	zap.S().Infof("Finished to update balances")
+	return nil
+}
+
 type assetInfoFn func(crypto.Digest, bool) (*assetInfo, error)
 
 func (s *balances) nftList(addr proto.Address, limit uint64, after []byte, assetInfoById assetInfoFn) ([]crypto.Digest, error) {

pkg/state/leases.go

@@ -116,10 +116,10 @@ func (l *leases) cancelLeases(bySenders map[proto.Address]struct{}, blockID prot
 	return nil
 }
 
-func (l *leases) cancelLeasesToAliases(aliases map[string]struct{}, blockID proto.BlockID) error {
+func (l *leases) cancelLeasesToAliases(aliases map[string]struct{}, blockID proto.BlockID) (map[proto.Address]balanceDiff, error) {
 	leaseIter, err := l.hs.newNewestTopEntryIterator(lease, true)
 	if err != nil {
-		return errors.Wrap(err, "failed to create key iterator to cancel leases to stolen aliases")
+		return nil, errors.Wrap(err, "failed to create key iterator to cancel leases to stolen aliases")
 	}
 	defer func() {
 		leaseIter.Release()
@@ -129,30 +129,43 @@ func (l *leases) cancelLeasesToAliases(aliases map[string]struct{}, blockID prot
 	}()
 
 	// Iterate all the leases.
-	zap.S().Info("Started collecting leases")
+	zap.S().Info("Started cancelling leases to disabled aliases")
+	changes := make(map[proto.Address]balanceDiff)
 	for leaseIter.Next() {
 		keyBytes := keyvalue.SafeKey(leaseIter)
 		var key leaseKey
 		if err := key.unmarshal(keyBytes); err != nil {
-			return errors.Wrap(err, "failed ot unmarshal leasing key")
+			return nil, errors.Wrap(err, "failed ot unmarshal leasing key")
 		}
 		leaseBytes := keyvalue.SafeValue(leaseIter)
 		record := new(leasing)
 		if err := cbor.Unmarshal(leaseBytes, record); err != nil {
-			return errors.Wrap(err, "failed to unmarshal lease")
+			return nil, errors.Wrap(err, "failed to unmarshal lease")
 		}
 		if record.isActive() && record.RecipientAlias != nil {
 			if _, ok := aliases[record.RecipientAlias.Alias]; ok {
 				zap.S().Infof("State: canceling lease %s", key.leaseID.String())
 				record.Status = LeaseCanceled
 				if err := l.addLeasing(key.leaseID, record, blockID); err != nil {
-					return errors.Wrap(err, "failed to save lease to storage")
+					return nil, errors.Wrap(err, "failed to save lease to storage")
+				}
+				if diff, ok := changes[record.Sender]; ok {
+					diff.leaseOut += -int64(record.Amount)
+					changes[record.Sender] = diff
+				} else {
+					changes[record.Sender] = newBalanceDiff(0, 0, -int64(record.Amount), false)
+				}
+				if diff, ok := changes[record.Recipient]; ok {
+					diff.leaseIn += -int64(record.Amount)
+					changes[record.Recipient] = diff
+				} else {
+					changes[record.Recipient] = newBalanceDiff(0, -int64(record.Amount), 0, false)
 				}
 			}
 		}
 	}
-	zap.S().Info("Finished collecting leases")
-	return nil
+	zap.S().Info("Finished cancelling leases to disabled aliases")
+	return changes, nil
 }
 
 func (l *leases) validLeaseIns() (map[proto.Address]int64, error) {

pkg/state/state.go

@@ -1297,7 +1297,11 @@ func (s *stateManager) cancelLeases(height uint64, blockID proto.BlockID, initia
 		if err != nil {
 			return err
 		}
-		if err = s.stor.leases.cancelLeasesToAliases(disabledAliases, blockID); err != nil {
+		changes, err := s.stor.leases.cancelLeasesToAliases(disabledAliases, blockID)
+		if err != nil {
+			return err
+		}
+		if err := s.stor.balances.cancelLeases(changes, blockID); err != nil {
 			return err
 		}
 	}