Add a migration which detects and filters out unreferenced slabs

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

@@ -372,7 +372,7 @@ func createTrieFromPayloads(logger zerolog.Logger, payloads []*ledger.Payload) (
 
 func newMigrations(
 	log zerolog.Logger,
-	dir string,
+	outputDir string,
 	runMigrations bool,
 	opts migrators.Options,
 ) []ledger.Migration {
@@ -382,10 +382,11 @@ func newMigrations(
 
 	log.Info().Msgf("initializing migrations")
 
-	rwf := reporters.NewReportFileWriterFactory(dir, log)
+	rwf := reporters.NewReportFileWriterFactory(outputDir, log)
 
 	namedMigrations := migrators.NewCadence1Migrations(
 		log,
+		outputDir,
 		rwf,
 		opts,
 	)

cmd/util/ledger/migrations/cadence.go

@@ -389,6 +389,7 @@ type Options struct {
 
 func NewCadence1Migrations(
 	log zerolog.Logger,
+	outputDir string,
 	rwf reporters.ReportWriterFactory,
 	opts Options,
 ) []NamedMigration {
@@ -420,7 +421,7 @@ func NewCadence1Migrations(
 				log,
 				opts.NWorker,
 				[]AccountBasedMigration{
-					NewFilterUnreferencedSlabsMigration(rwf),
+					NewFilterUnreferencedSlabsMigration(outputDir, rwf),
 				},
 			),
 		})

cmd/util/ledger/migrations/cadence_values_migration_test.go

@@ -79,6 +79,7 @@ func TestCadenceValuesMigration(t *testing.T) {
 
 	migrations := NewCadence1Migrations(
 		logger,
+		t.TempDir(),
 		rwf,
 		Options{
 			NWorker:              nWorker,
@@ -698,6 +699,7 @@ func TestBootstrappedStateMigration(t *testing.T) {
 
 	migrations := NewCadence1Migrations(
 		logger,
+		t.TempDir(),
 		rwf,
 		Options{
 			NWorker:              nWorker,
@@ -823,6 +825,7 @@ func TestProgramParsingError(t *testing.T) {
 
 	migrations := NewCadence1Migrations(
 		logger,
+		t.TempDir(),
 		rwf,
 		Options{
 			NWorker:              nWorker,
@@ -948,6 +951,7 @@ func TestCoreContractUsage(t *testing.T) {
 
 		migrations := NewCadence1Migrations(
 			logger,
+			t.TempDir(),
 			rwf,
 			Options{
 				NWorker:              nWorker,

cmd/util/ledger/migrations/filter_unreferenced_slabs_migration.go

@@ -4,13 +4,17 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"path"
+	"sync"
+	"time"
 
 	"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/cmd/util/ledger/util"
 	"github.com/onflow/flow-go/ledger"
 	"github.com/onflow/flow-go/ledger/common/convert"
 	"github.com/onflow/flow-go/model/flow"
@@ -25,19 +29,26 @@ func StorageIDFromRegisterID(registerID flow.RegisterID) atree.StorageID {
 }
 
 type FilterUnreferencedSlabsMigration struct {
-	log zerolog.Logger
-	rw  reporters.ReportWriter
+	log              zerolog.Logger
+	rw               reporters.ReportWriter
+	outputDir        string
+	mutex            sync.Mutex
+	filteredAccounts []common.Address
+	filteredPayloads []*ledger.Payload
+	payloadsFile     string
 }
 
 var _ AccountBasedMigration = &FilterUnreferencedSlabsMigration{}
 
 const filterUnreferencedSlabsName = "filter-unreferenced-slabs"
 
 func NewFilterUnreferencedSlabsMigration(
+	outputDir string,
 	rwf reporters.ReportWriterFactory,
 ) *FilterUnreferencedSlabsMigration {
 	return &FilterUnreferencedSlabsMigration{
-		rw: rwf.ReportWriter(filterUnreferencedSlabsName),
+		outputDir: outputDir,
+		rw:        rwf.ReportWriter(filterUnreferencedSlabsName),
 	}
 }
 
@@ -136,25 +147,76 @@ func (m *FilterUnreferencedSlabsMigration) MigrateAccount(
 	}
 
 	m.rw.Write(unreferencedSlabs{
-		Account:  address,
-		Payloads: filteredPayloads,
+		Account:      address,
+		PayloadCount: len(filteredPayloads),
 	})
 
+	m.mergeFilteredPayloads(
+		address,
+		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) mergeFilteredPayloads(
+	address common.Address,
+	payloads []*ledger.Payload,
+) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	m.filteredAccounts = append(m.filteredAccounts, address)
+	m.filteredPayloads = append(m.filteredPayloads, payloads...)
+}
+
 func (m *FilterUnreferencedSlabsMigration) Close() error {
 	// close the report writer so it flushes to file
 	m.rw.Close()
 
+	err := m.writeFilteredPayloads()
+	if err != nil {
+		return fmt.Errorf("failed to write filtered payloads to file: %w", err)
+	}
+
+	return nil
+}
+
+func (m *FilterUnreferencedSlabsMigration) writeFilteredPayloads() error {
+
+	m.payloadsFile = path.Join(
+		m.outputDir,
+		fmt.Sprintf("filtered_%d.payloads", int32(time.Now().Unix())),
+	)
+
+	writtenPayloadCount, err := util.CreatePayloadFile(
+		m.log,
+		m.payloadsFile,
+		m.filteredPayloads,
+		m.filteredAccounts,
+		true,
+	)
+
+	if err != nil {
+		return fmt.Errorf("failed to write all filtered payloads to file: %w", err)
+	}
+
+	if writtenPayloadCount != len(m.filteredPayloads) {
+		return fmt.Errorf(
+			"failed to write all filtered payloads to file: expected %d, got %d",
+			len(m.filteredPayloads),
+			writtenPayloadCount,
+		)
+	}
+
 	return nil
 }
 
 type unreferencedSlabs struct {
-	Account  common.Address    `json:"account"`
-	Payloads []*ledger.Payload `json:"payloads"`
+	Account      common.Address `json:"account"`
+	PayloadCount int            `json:"payload_count"`
 }

cmd/util/ledger/migrations/filter_unreferenced_slabs_migration_test.go

@@ -136,7 +136,7 @@ func TestFilterUnreferencedSlabs(t *testing.T) {
 	// Act
 
 	rwf := &testReportWriterFactory{}
-	migration := NewFilterUnreferencedSlabsMigration(rwf)
+	migration := NewFilterUnreferencedSlabsMigration(t.TempDir(), rwf)
 
 	log := zerolog.New(zerolog.NewTestWriter(t))
 
@@ -148,6 +148,9 @@ func TestFilterUnreferencedSlabs(t *testing.T) {
 	newPayloads, err := migration.MigrateAccount(ctx, testAddress, oldPayloads)
 	require.NoError(t, err)
 
+	err = migration.Close()
+	require.NoError(t, err)
+
 	// Assert
 
 	writer := rwf.reportWriters[filterUnreferencedSlabsName]
@@ -180,10 +183,23 @@ func TestFilterUnreferencedSlabs(t *testing.T) {
 	assert.Equal(t,
 		[]any{
 			unreferencedSlabs{
-				Account:  testAddress,
-				Payloads: expectedFilteredPayloads,
+				Account:      testAddress,
+				PayloadCount: len(expectedFilteredPayloads),
 			},
 		},
 		writer.entries,
 	)
+	assert.Equal(t,
+		expectedFilteredPayloads,
+		migration.filteredPayloads,
+	)
+	assert.Equal(t,
+		[]common.Address{testAddress},
+		migration.filteredAccounts,
+	)
+
+	readIsPartial, readFilteredPayloads, err := util.ReadPayloadFile(log, migration.payloadsFile)
+	require.NoError(t, err)
+	assert.True(t, readIsPartial)
+	assert.Equal(t, expectedFilteredPayloads, readFilteredPayloads)
 }

cmd/util/ledger/migrations/staged_contracts_migration_test.go

@@ -1797,6 +1797,7 @@ func TestConcurrentContractUpdate(t *testing.T) {
 
 	migrations := NewCadence1Migrations(
 		logger,
+		t.TempDir(),
 		rwf,
 		Options{
 			NWorker:              nWorker,