github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/cmd/util/ledger/migrations/storage_used_migration.go (about)

     1  package migrations
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"github.com/onflow/cadence/runtime/common"
     8  	"github.com/rs/zerolog"
     9  	"github.com/rs/zerolog/log"
    10  
    11  	"github.com/onflow/flow-go/cmd/util/ledger/util/registers"
    12  	"github.com/onflow/flow-go/fvm/environment"
    13  	"github.com/onflow/flow-go/model/flow"
    14  )
    15  
    16  // AccountUsageMigration iterates through each payload, and calculate the storage usage
    17  // and update the accounts status with the updated storage usage
    18  type AccountUsageMigration struct {
    19  	log zerolog.Logger
    20  }
    21  
    22  var _ AccountBasedMigration = &AccountUsageMigration{}
    23  
    24  func (m *AccountUsageMigration) InitMigration(
    25  	log zerolog.Logger,
    26  	_ *registers.ByAccount,
    27  	_ int,
    28  ) error {
    29  	m.log = log.With().Str("component", "AccountUsageMigration").Logger()
    30  	return nil
    31  }
    32  
    33  const oldAccountStatusSize = 25
    34  
    35  func (m *AccountUsageMigration) Close() error {
    36  	return nil
    37  }
    38  
    39  func (m *AccountUsageMigration) MigrateAccount(
    40  	_ context.Context,
    41  	address common.Address,
    42  	accountRegisters *registers.AccountRegisters,
    43  ) error {
    44  
    45  	var status *environment.AccountStatus
    46  	var statusValue []byte
    47  	actualSize := uint64(0)
    48  
    49  	// Find the account status register,
    50  	// and calculate the storage usage
    51  	err := accountRegisters.ForEach(func(owner, key string, value []byte) error {
    52  
    53  		if key == flow.AccountStatusKey {
    54  			statusValue = value
    55  
    56  			var err error
    57  			status, err = environment.AccountStatusFromBytes(value)
    58  			if err != nil {
    59  				return fmt.Errorf("could not parse account status: %w", err)
    60  			}
    61  		}
    62  
    63  		actualSize += uint64(environment.RegisterSize(
    64  			flow.RegisterID{
    65  				Owner: owner,
    66  				Key:   key,
    67  			},
    68  			value,
    69  		))
    70  
    71  		return nil
    72  	})
    73  	if err != nil {
    74  		return fmt.Errorf(
    75  			"could not iterate through registers of account %s: %w",
    76  			address.HexWithPrefix(),
    77  			err,
    78  		)
    79  	}
    80  
    81  	if status == nil {
    82  		log.Error().
    83  			Str("account", address.HexWithPrefix()).
    84  			Msgf("could not find account status register")
    85  		return nil
    86  	}
    87  
    88  	isOldVersionOfStatusRegister := len(statusValue) == oldAccountStatusSize
    89  
    90  	same := m.compareUsage(isOldVersionOfStatusRegister, status, actualSize, address)
    91  	if same {
    92  		// there is no problem with the usage, return
    93  		return nil
    94  	}
    95  
    96  	if isOldVersionOfStatusRegister {
    97  		// size will grow by 8 bytes because of the on-the-fly
    98  		// migration of account status in AccountStatusFromBytes
    99  		actualSize += 8
   100  	}
   101  
   102  	// update storage used
   103  	status.SetStorageUsed(actualSize)
   104  
   105  	err = accountRegisters.Set(
   106  		string(address[:]),
   107  		flow.AccountStatusKey,
   108  		status.ToBytes(),
   109  	)
   110  	if err != nil {
   111  		return fmt.Errorf("could not update account status register: %w", err)
   112  	}
   113  
   114  	return nil
   115  }
   116  
   117  func (m *AccountUsageMigration) compareUsage(
   118  	isOldVersionOfStatusRegister bool,
   119  	status *environment.AccountStatus,
   120  	actualSize uint64,
   121  	address common.Address,
   122  ) bool {
   123  	oldSize := status.StorageUsed()
   124  	if isOldVersionOfStatusRegister {
   125  		// size will be reported as 8 bytes larger than the actual size due to on-the-fly
   126  		// migration of account status in AccountStatusFromBytes
   127  		oldSize -= 8
   128  	}
   129  
   130  	if oldSize != actualSize {
   131  		m.log.Warn().
   132  			Uint64("old_size", oldSize).
   133  			Uint64("new_size", actualSize).
   134  			Str("account", address.HexWithPrefix()).
   135  			Msg("account storage used usage mismatch")
   136  		return false
   137  	}
   138  	return true
   139  }