test pre-migrations

chain/stmgr/forks_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"sync"
 	"testing"
 
 	"github.com/ipfs/go-cid"
@@ -317,3 +318,149 @@ func TestForkRefuseCall(t *testing.T) {
 		}
 	}
 }
+
+func TestForkPreMigration(t *testing.T) {
+	logging.SetAllLoggers(logging.LevelInfo)
+
+	cg, err := gen.NewGenerator()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	fooCid, err := abi.CidBuilder.Sum([]byte("foo"))
+	require.NoError(t, err)
+
+	barCid, err := abi.CidBuilder.Sum([]byte("bar"))
+	require.NoError(t, err)
+
+	failCid, err := abi.CidBuilder.Sum([]byte("fail"))
+	require.NoError(t, err)
+
+	var wait20 sync.WaitGroup
+	wait20.Add(3)
+
+	wasCanceled := make(chan struct{})
+
+	checkCache := func(t *testing.T, cache MigrationCache) {
+		found, value, err := cache.Read("foo")
+		require.NoError(t, err)
+		require.True(t, found)
+		require.Equal(t, fooCid, value)
+
+		found, value, err = cache.Read("bar")
+		require.NoError(t, err)
+		require.True(t, found)
+		require.Equal(t, barCid, value)
+
+		found, _, err = cache.Read("fail")
+		require.NoError(t, err)
+		require.False(t, found)
+	}
+
+	sm, err := NewStateManagerWithUpgradeSchedule(
+		cg.ChainStore(), UpgradeSchedule{{
+			Network: 1,
+			Height:  testForkHeight,
+			Migration: func(ctx context.Context, sm *StateManager, cache MigrationCache, cb ExecCallback,
+				root cid.Cid, height abi.ChainEpoch, ts *types.TipSet) (cid.Cid, error) {
+
+				// Make sure the test that should be canceled, is canceled.
+				select {
+				case <-wasCanceled:
+				case <-ctx.Done():
+					return cid.Undef, ctx.Err()
+				}
+
+				// the cache should be setup correctly.
+				checkCache(t, cache)
+
+				return root, nil
+			},
+			PreMigrations: []PreMigration{{
+				When: 20,
+				PreMigration: func(ctx context.Context, _ *StateManager, cache MigrationCache,
+					_ cid.Cid, _ abi.ChainEpoch, _ *types.TipSet) error {
+					wait20.Done()
+					wait20.Wait()
+
+					err := cache.Write("foo", fooCid)
+					require.NoError(t, err)
+
+					return nil
+				},
+			}, {
+				When: 20,
+				PreMigration: func(ctx context.Context, _ *StateManager, cache MigrationCache,
+					_ cid.Cid, _ abi.ChainEpoch, _ *types.TipSet) error {
+					wait20.Done()
+					wait20.Wait()
+
+					err := cache.Write("bar", barCid)
+					require.NoError(t, err)
+
+					return nil
+				},
+			}, {
+				When: 20,
+				PreMigration: func(ctx context.Context, _ *StateManager, cache MigrationCache,
+					_ cid.Cid, _ abi.ChainEpoch, _ *types.TipSet) error {
+					wait20.Done()
+					wait20.Wait()
+
+					err := cache.Write("fail", failCid)
+					require.NoError(t, err)
+
+					// Fail this migration. The cached entry should not be persisted.
+					return fmt.Errorf("failed")
+				},
+			}, {
+				When: 10,
+				PreMigration: func(ctx context.Context, _ *StateManager, cache MigrationCache,
+					_ cid.Cid, _ abi.ChainEpoch, _ *types.TipSet) error {
+
+					checkCache(t, cache)
+
+					return nil
+				},
+			}, {
+				When:     15,
+				NotAfter: 5,
+				PreMigration: func(ctx context.Context, _ *StateManager, cache MigrationCache,
+					_ cid.Cid, _ abi.ChainEpoch, _ *types.TipSet) error {
+
+					<-ctx.Done()
+					close(wasCanceled)
+
+					return nil
+				},
+			}}},
+		})
+	if err != nil {
+		t.Fatal(err)
+	}
+	require.NoError(t, sm.Start(context.Background()))
+	defer func() {
+		require.NoError(t, sm.Stop(context.Background()))
+	}()
+
+	inv := vm.NewActorRegistry()
+	inv.Register(nil, testActor{})
+
+	sm.SetVMConstructor(func(ctx context.Context, vmopt *vm.VMOpts) (*vm.VM, error) {
+		nvm, err := vm.NewVM(ctx, vmopt)
+		if err != nil {
+			return nil, err
+		}
+		nvm.SetInvoker(inv)
+		return nvm, nil
+	})
+
+	cg.SetStateManager(sm)
+
+	for i := 0; i < 50; i++ {
+		_, err := cg.NextTipSet()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}