Re-derive the composite write mode on rollback

sei-db/state_db/sc/composite/store.go

@@ -404,10 +404,7 @@ func (cs *CompositeCommitStore) LoadLatest() error {
 		}
 	}
 
-	if err := cs.resolveCurrentWriteMode(true); err != nil {
-		return fmt.Errorf("failed to resolve write mode: %w", err)
-	}
-	if err := cs.buildRouter(); err != nil {
+	if err := cs.adoptPersistedWriteMode(); err != nil {
 		return err
 	}
 	// After the router, because the gating this reads gets its answer from migration metadata through
@@ -491,6 +488,15 @@ func (cs *CompositeCommitStore) LoadVersionReadOnly(targetVersion int64) (_ type
 	return ro, nil
 }
 
+// adoptPersistedWriteMode sets the effective write mode from the opened backends' persisted state and
+// installs a router for it.
+func (cs *CompositeCommitStore) adoptPersistedWriteMode() error {
+	if err := cs.resolveCurrentWriteMode(true); err != nil {
+		return fmt.Errorf("failed to resolve write mode: %w", err)
+	}
+	return cs.buildRouter()
+}
+
 // resolveCurrentWriteMode sets cs.currentWriteMode after the backends have been
 // opened. For a fixed configured mode this is a copy; for types.Auto the
 // mode is derived from the migration metadata persisted in flatkv. A nil
@@ -1440,21 +1446,22 @@ func (cs *CompositeCommitStore) Rollback(targetVersion int64) error {
 	// `seid rollback` prints and rootmulti caches in rs.lastCommitInfo)
 	// diverges from the canonical AppHash for the target height. The gates
 	// re-latch correctly on the next call against the rolled-back metadata.
-	//
-	// Note: currentWriteMode is not re-derived here. It only matters when a
-	// rollback crosses a seam whose latest-derived mode differs from the
-	// target's (e.g. a rollback all the way across a completed migration);
-	// that in-process view self-heals on the next `seid start`, which
-	// re-derives the mode from the rolled-back metadata.
 	cs.latticeAppendLatched.Store(false)
 	cs.memiavlHashExcluded.Store(false)
 
+	// The effective mode and its router were derived at the pre-rollback height. A rollback across a
+	// migration seam (back below bank completion, or to before activation) leaves them describing state
+	// the store no longer holds, so derive both again from the rolled-back metadata, as a restart would.
+	if err := cs.adoptPersistedWriteMode(); err != nil {
+		return err
+	}
+
 	// Rollback is offline (no commit cycle in flight); clear the per-block
 	// migration-advance gate defensively.
 	cs.migrationAdvancedThisCommit = false
 
-	// After the latch resets above, so the rebuilt info reflects the rolled-back metadata rather than
-	// the gating that was latched at the pre-rollback height.
+	// After the latch resets and the re-derived mode above, so the rebuilt info reflects the rolled-back
+	// metadata rather than the gating that was latched at the pre-rollback height.
 	return cs.refreshLastCommitInfo()
 }
 

sei-db/state_db/sc/composite/store_migration_test.go

@@ -1165,24 +1165,7 @@ func TestComposite_MigrateBank_RollbackAcrossCompletionBoundary(t *testing.T) {
 	const batch = 8
 
 	cs := openAutoStore(t, dir, batch)
-
-	// Bootstrap bank+evm state under MemiavlOnly (bank-heavy so migrate_bank
-	// later spans several blocks), then walk the ladder to migrate_bank.
-	for i := 0; i < 12; i++ {
-		require.NoError(t, cs.ApplyChangeSets(workload.generateBlock(6, 0, 0, 40, 0)))
-		_, err := cs.Commit(cs.Version() + 1)
-		require.NoError(t, err)
-	}
-	require.NoError(t, cs.SetWriteMode(types.MigrateEVM))
-	runUntilAtMigrationVersion(t, cs, workload, migration.Version1_MigrateEVM, 400)
-	require.NoError(t, cs.SetWriteMode(types.EVMMigrated))
-	runBlocks(t, cs, workload, 2)
-	require.NoError(t, cs.SetWriteMode(types.MigrateAllButBank))
-	runUntilAtMigrationVersion(t, cs, workload, migration.Version2_MigrateAllButBank, 400)
-	require.NoError(t, cs.SetWriteMode(types.AllMigratedButBank))
-	runBlocks(t, cs, workload, 2)
-
-	require.NoError(t, cs.SetWriteMode(types.MigrateBank))
+	walkToMigrateBank(t, cs, workload)
 	require.Equal(t, types.MigrateBank, cs.loadWriteMode())
 
 	// One migrate_bank block that does not complete the migration: memiavl
@@ -1220,6 +1203,170 @@ func TestComposite_MigrateBank_RollbackAcrossCompletionBoundary(t *testing.T) {
 		"post-restart commit info across the bank-completion boundary must re-include memiavl")
 }
 
+// TestComposite_Auto_RollbackAcrossBankCompletionAfterRestart rolls an Auto store that was reopened
+// after the bank migration completed, so its mode was derived as FlatKVOnly, back below completion. The
+// rolled-back store must report the canonical commit info and recommit the rolled-back blocks exactly,
+// both in the same process and after a restart.
+func TestComposite_Auto_RollbackAcrossBankCompletionAfterRestart(t *testing.T) {
+	dir := t.TempDir()
+	workload := newMigrationWorkload(0xB0A7)
+	const batch = 8
+
+	cs := openAutoStore(t, dir, batch)
+	walkToMigrateBank(t, cs, workload)
+
+	target := commitRecorded(t, cs, workload.generateBlock(0, 2, 1, 0, 2))
+	require.Greater(t, len(target.info.StoreInfos), 1,
+		"sanity: an in-progress bank migration still carries memiavl's store infos")
+	rolledBack := commitUntilAtMigrationVersion(t, cs, workload, migration.Version3_FlatKVOnly, 800)
+	require.NoError(t, cs.Close())
+
+	cs = openAutoStore(t, dir, batch)
+	require.Equal(t, types.FlatKVOnly, cs.loadWriteMode(), "sanity: a reopened store derives FlatKVOnly")
+	require.NoError(t, cs.Rollback(target.info.Version))
+	require.Equal(t, types.MigrateBank, cs.loadWriteMode())
+	requireCommitInfoEqual(t, target.info, cs.LastCommitInfo(), "post-rollback commit info (same process)")
+	replayRecorded(t, cs, rolledBack, "replay after rollback (same process)")
+	require.NoError(t, cs.Close())
+
+	cs = openAutoStore(t, dir, batch)
+	require.NoError(t, cs.Rollback(target.info.Version))
+	require.NoError(t, cs.Close())
+	cs = openAutoStore(t, dir, batch)
+	defer func() { _ = cs.Close() }()
+	require.Equal(t, types.MigrateBank, cs.loadWriteMode())
+	requireCommitInfoEqual(t, target.info, cs.LastCommitInfo(), "post-rollback commit info (after restart)")
+	replayRecorded(t, cs, rolledBack, "replay after rollback (after restart)")
+}
+
+// TestComposite_Auto_RollbackToActivationAfterEVMMigration rolls an Auto store that was reopened after
+// the EVM migration completed back to the last height before activation. The rolled-back store must be
+// MemiavlOnly with flatkv closed, report the canonical commit info, and reproduce activation and the
+// whole migration when the blocks are recommitted, both in the same process and after a restart.
+func TestComposite_Auto_RollbackToActivationAfterEVMMigration(t *testing.T) {
+	dir := t.TempDir()
+	workload := newMigrationWorkload(0xE7E7)
+	const batch = 8
+
+	cs := openAutoStore(t, dir, batch)
+	var preActivation recordedBlock
+	for i := 0; i < 12; i++ {
+		preActivation = commitRecorded(t, cs, workload.generateBlock(6, 0, 0, 4, 0))
+	}
+	require.False(t, hasLattice(preActivation.info), "sanity: pre-activation commit info has no evm_lattice")
+	require.NoError(t, cs.SetWriteMode(types.MigrateEVM))
+	rolledBack := commitUntilAtMigrationVersion(t, cs, workload, migration.Version1_MigrateEVM, 400)
+	require.NoError(t, cs.Close())
+
+	requireRolledBackToActivation := func(cs *CompositeCommitStore, msg string) {
+		t.Helper()
+		require.Equal(t, preActivation.info.Version, cs.Version(), msg)
+		require.Equal(t, types.MemiavlOnly, cs.loadWriteMode(), msg)
+		require.Nil(t, cs.loadFlatKV(), "%s: flatkv must be closed before activation", msg)
+		requireCommitInfoEqual(t, preActivation.info, cs.LastCommitInfo(), msg)
+	}
+
+	cs = openAutoStore(t, dir, batch)
+	require.Equal(t, types.EVMMigrated, cs.loadWriteMode(), "sanity: a reopened store derives EVMMigrated")
+	require.NoError(t, cs.Rollback(preActivation.info.Version))
+	requireRolledBackToActivation(cs, "post-rollback (same process)")
+	require.NoError(t, cs.SetWriteMode(types.MigrateEVM))
+	replayRecorded(t, cs, rolledBack, "replay after rollback (same process)")
+	require.NoError(t, cs.Close())
+
+	cs = openAutoStore(t, dir, batch)
+	require.NoError(t, cs.Rollback(preActivation.info.Version))
+	require.NoError(t, cs.Close())
+	cs = openAutoStore(t, dir, batch)
+	defer func() { _ = cs.Close() }()
+	requireRolledBackToActivation(cs, "post-rollback (after restart)")
+	require.NoError(t, cs.SetWriteMode(types.MigrateEVM))
+	replayRecorded(t, cs, rolledBack, "replay after rollback (after restart)")
+}
+
+// walkToMigrateBank bootstraps bank-heavy state under MemiavlOnly, so the bank migration spans several
+// blocks, then walks cs through the migration ladder and leaves it in MigrateBank.
+func walkToMigrateBank(t *testing.T, cs *CompositeCommitStore, workload *migrationWorkload) {
+	t.Helper()
+	for i := 0; i < 12; i++ {
+		require.NoError(t, cs.ApplyChangeSets(workload.generateBlock(6, 0, 0, 40, 0)))
+		_, err := cs.Commit(cs.Version() + 1)
+		require.NoError(t, err)
+	}
+	require.NoError(t, cs.SetWriteMode(types.MigrateEVM))
+	runUntilAtMigrationVersion(t, cs, workload, migration.Version1_MigrateEVM, 400)
+	require.NoError(t, cs.SetWriteMode(types.EVMMigrated))
+	runBlocks(t, cs, workload, 2)
+	require.NoError(t, cs.SetWriteMode(types.MigrateAllButBank))
+	runUntilAtMigrationVersion(t, cs, workload, migration.Version2_MigrateAllButBank, 400)
+	require.NoError(t, cs.SetWriteMode(types.AllMigratedButBank))
+	runBlocks(t, cs, workload, 2)
+	require.NoError(t, cs.SetWriteMode(types.MigrateBank))
+}
+
+// recordedBlock is a committed block's change sets and the commit info committing them produced.
+type recordedBlock struct {
+	changes []*proto.NamedChangeSet
+	info    *proto.CommitInfo
+}
+
+// commitRecorded commits changes as the next block and returns its record.
+func commitRecorded(t *testing.T, cs *CompositeCommitStore, changes []*proto.NamedChangeSet) recordedBlock {
+	t.Helper()
+	recorded := cloneChangeSets(t, changes)
+	require.NoError(t, cs.ApplyChangeSets(changes))
+	_, err := cs.Commit(cs.Version() + 1)
+	require.NoError(t, err)
+	return recordedBlock{changes: recorded, info: cloneCommitInfo(cs.LastCommitInfo())}
+}
+
+// commitUntilAtMigrationVersion is runUntilAtMigrationVersion that returns the record of every block it
+// committed.
+func commitUntilAtMigrationVersion(
+	t *testing.T,
+	cs *CompositeCommitStore,
+	workload *migrationWorkload,
+	version uint64,
+	maxBlocks int,
+) []recordedBlock {
+	t.Helper()
+	var blocks []recordedBlock
+	for i := 0; i < maxBlocks; i++ {
+		blocks = append(blocks, commitRecorded(t, cs, workload.generateBlock(0, 2, 1, 1, 1)))
+		done, err := migration.IsAtVersion(flatKVReaderFor(cs), version)
+		require.NoError(t, err)
+		if done {
+			return blocks
+		}
+	}
+	t.Fatalf("migration did not reach version %d within %d blocks", version, maxBlocks)
+	return nil
+}
+
+// replayRecorded recommits blocks and requires each to reproduce its recorded commit info.
+func replayRecorded(t *testing.T, cs *CompositeCommitStore, blocks []recordedBlock, msg string) {
+	t.Helper()
+	for _, block := range blocks {
+		require.NoError(t, cs.ApplyChangeSets(cloneChangeSets(t, block.changes)))
+		_, err := cs.Commit(cs.Version() + 1)
+		require.NoError(t, err)
+		requireCommitInfoEqual(t, block.info, cs.LastCommitInfo(), fmt.Sprintf("%s: block %d", msg, block.info.Version))
+	}
+}
+
+// cloneChangeSets returns a deep copy of changes.
+func cloneChangeSets(t *testing.T, changes []*proto.NamedChangeSet) []*proto.NamedChangeSet {
+	t.Helper()
+	clones := make([]*proto.NamedChangeSet, len(changes))
+	for i, changeSet := range changes {
+		bz, err := changeSet.Marshal()
+		require.NoError(t, err)
+		clones[i] = &proto.NamedChangeSet{}
+		require.NoError(t, clones[i].Unmarshal(bz))
+	}
+	return clones
+}
+
 // TestMigrateEVMBeforeTheBoundaryDrainsTheHashStream pins the drain a committing flatkv needs while it
 // is still outside the AppHash.
 //