Add a migration which detects and filters out unreferenced slabs

cmd/util/cmd/execution-state-extract/cmd.go

@@ -46,6 +46,7 @@ var (
 	flagOutputPayloadFileName              string
 	flagOutputPayloadByAddresses           string
 	flagMaxAccountSize                     uint64
+	flagFilterUnreferencedSlabs            bool
 )
 
 var Cmd = &cobra.Command{
@@ -151,6 +152,9 @@ func init() {
 
 	Cmd.Flags().Uint64Var(&flagMaxAccountSize, "max-account-size", 0,
 		"max account size")
+
+	Cmd.Flags().BoolVar(&flagFilterUnreferencedSlabs, "filter-unreferenced-slabs", false,
+		"filter unreferenced slabs")
 }
 
 func run(*cobra.Command, []string) {
@@ -371,6 +375,7 @@ func run(*cobra.Command, []string) {
 		Prune:                             flagPrune,
 		MaxAccountSize:                    flagMaxAccountSize,
 		VerboseErrorOutput:                flagVerboseErrorOutput,
+		FilterUnreferencedSlabs:           flagFilterUnreferencedSlabs,
 	}
 
 	if len(flagInputPayloadFileName) > 0 {

cmd/util/ledger/migrations/cadence.go

@@ -384,6 +384,7 @@ type Options struct {
 	StagedContracts                   []StagedContract
 	Prune                             bool
 	MaxAccountSize                    uint64
+	FilterUnreferencedSlabs           bool
 }
 
 func NewCadence1Migrations(
@@ -412,6 +413,19 @@ func NewCadence1Migrations(
 		)
 	}
 
+	if opts.FilterUnreferencedSlabs {
+		migrations = append(migrations, NamedMigration{
+			Name: "filter-unreferenced-slabs-migration",
+			Migrate: NewAccountBasedMigration(
+				log,
+				opts.NWorker,
+				[]AccountBasedMigration{
+					NewFilterUnreferencedSlabsMigration(rwf),
+				},
+			),
+		})
+	}
+
 	if opts.Prune {
 		migration := NewCadence1PruneMigration(opts.ChainID, log)
 		if migration != nil {

cmd/util/ledger/migrations/filter_unreferenced_slabs_migration.go

@@ -0,0 +1,160 @@
+package migrations
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/onflow/atree"
+	"github.com/onflow/cadence/runtime"
+	"github.com/onflow/cadence/runtime/common"
+	"github.com/rs/zerolog"
+
+	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
+	"github.com/onflow/flow-go/ledger"
+	"github.com/onflow/flow-go/ledger/common/convert"
+	"github.com/onflow/flow-go/model/flow"
+)
+
+func StorageIDFromRegisterID(registerID flow.RegisterID) atree.StorageID {
+	storageID := atree.StorageID{
+		Address: atree.Address([]byte(registerID.Owner)),
+	}
+	copy(storageID.Index[:], registerID.Key[1:])
+	return storageID
+}
+
+type FilterUnreferencedSlabsMigration struct {
+	log zerolog.Logger
+	rw  reporters.ReportWriter
+}
+
+var _ AccountBasedMigration = &FilterUnreferencedSlabsMigration{}
+
+const filterUnreferencedSlabsName = "filter-unreferenced-slabs"
+
+func NewFilterUnreferencedSlabsMigration(
+	rwf reporters.ReportWriterFactory,
+) *FilterUnreferencedSlabsMigration {
+	return &FilterUnreferencedSlabsMigration{
+		rw: rwf.ReportWriter(filterUnreferencedSlabsName),
+	}
+}
+
+func (m *FilterUnreferencedSlabsMigration) InitMigration(
+	log zerolog.Logger,
+	_ []*ledger.Payload,
+	_ int,
+) error {
+	m.log = log.
+		With().
+		Str("migration", filterUnreferencedSlabsName).
+		Logger()
+
+	return nil
+}
+
+func (m *FilterUnreferencedSlabsMigration) MigrateAccount(
+	_ context.Context,
+	address common.Address,
+	oldPayloads []*ledger.Payload,
+) (
+	newPayloads []*ledger.Payload,
+	err error,
+) {
+	migrationRuntime, err := NewAtreeRegisterMigratorRuntime(address, oldPayloads)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create migrator runtime: %w", err)
+	}
+
+	storage := migrationRuntime.Storage
+
+	newPayloads = oldPayloads
+
+	err = checkStorageHealth(address, storage, oldPayloads)
+	if err == nil {
+		return
+	}
+
+	// The storage health check failed.
+	// This can happen if there are unreferenced root slabs.
+	// In this case, we filter out the unreferenced root slabs and all slabs they reference from the payloads.
+
+	var unreferencedRootSlabsErr runtime.UnreferencedRootSlabsError
+	if !errors.As(err, &unreferencedRootSlabsErr) {
+		return nil, fmt.Errorf("storage health check failed: %w", err)
+	}
+
+	m.log.Warn().
+		Err(err).
+		Str("account", address.Hex()).
+		Msg("filtering unreferenced root slabs")
+
+	// Create a set of unreferenced slabs: root slabs, and all slabs they reference.
+
+	unreferencedSlabIDs := map[atree.StorageID]struct{}{}
+	for _, rootSlabID := range unreferencedRootSlabsErr.UnreferencedRootSlabIDs {
+		unreferencedSlabIDs[rootSlabID] = struct{}{}
+
+		childReferences, _, err := storage.GetAllChildReferences(rootSlabID)
+		if err != nil {
+			return nil, fmt.Errorf(
+				"failed to get all child references for root slab %s: %w",
+				rootSlabID,
+				err,
+			)
+		}
+
+		for _, childSlabID := range childReferences {
+			unreferencedSlabIDs[childSlabID] = struct{}{}
+		}
+	}
+
+	// Filter out unreferenced slabs.
+
+	newCount := len(oldPayloads) - len(unreferencedSlabIDs)
+	newPayloads = make([]*ledger.Payload, 0, newCount)
+
+	filteredPayloads := make([]*ledger.Payload, 0, len(unreferencedSlabIDs))
+
+	for _, payload := range oldPayloads {
+		registerID, _, err := convert.PayloadToRegister(payload)
+		if err != nil {
+			return nil, fmt.Errorf("failed to convert payload to register: %w", err)
+		}
+
+		// Filter unreferenced slabs.
+		if registerID.IsSlabIndex() {
+			storageID := StorageIDFromRegisterID(registerID)
+			if _, ok := unreferencedSlabIDs[storageID]; ok {
+				filteredPayloads = append(filteredPayloads, payload)
+				continue
+			}
+		}
+
+		newPayloads = append(newPayloads, payload)
+	}
+
+	m.rw.Write(unreferencedSlabs{
+		Account:  address,
+		Payloads: filteredPayloads,
+	})
+
+	// Do NOT report the health check error here.
+	// The health check error is only reported if it is not due to unreferenced slabs.
+	// If it is due to unreferenced slabs, we filter them out and continue.
+
+	return newPayloads, nil
+}
+
+func (m *FilterUnreferencedSlabsMigration) Close() error {
+	// close the report writer so it flushes to file
+	m.rw.Close()
+
+	return nil
+}
+
+type unreferencedSlabs struct {
+	Account  common.Address    `json:"account"`
+	Payloads []*ledger.Payload `json:"payloads"`
+}