github.com/onflow/flow-go@v0.33.17/cmd/util/ledger/migrations/storage_used_migration.go (about)

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