Merge pull request #4481 from filecoin-project/feat/optim-state-cache

state: optimize state snapshot address cache

chain/state/statetree.go

@@ -30,12 +30,14 @@ type StateTree struct {
 	version     types.StateTreeVersion
 	info        cid.Cid
 	Store       cbor.IpldStore
+	lookupIDFun func(address.Address) (address.Address, error)
 
 	snaps *stateSnaps
 }
 
 type stateSnaps struct {
 	layers                        []*stateSnapLayer
+	lastMaybeNonEmptyResolveCache int
 }
 
 type stateSnapLayer struct {
@@ -67,7 +69,12 @@ func (ss *stateSnaps) addLayer() {
 
 func (ss *stateSnaps) dropLayer() {
 	ss.layers[len(ss.layers)-1] = nil // allow it to be GCed
+
 	ss.layers = ss.layers[:len(ss.layers)-1]
+
+	if ss.lastMaybeNonEmptyResolveCache == len(ss.layers) {
+		ss.lastMaybeNonEmptyResolveCache = len(ss.layers) - 1
+	}
 }
 
 func (ss *stateSnaps) mergeLastLayer() {
@@ -86,7 +93,13 @@ func (ss *stateSnaps) mergeLastLayer() {
 }
 
 func (ss *stateSnaps) resolveAddress(addr address.Address) (address.Address, bool) {
-	for i := len(ss.layers) - 1; i >= 0; i-- {
+	for i := ss.lastMaybeNonEmptyResolveCache; i >= 0; i-- {
+		if len(ss.layers[i].resolveCache) == 0 {
+			if ss.lastMaybeNonEmptyResolveCache == i {
+				ss.lastMaybeNonEmptyResolveCache = i - 1
+			}
+			continue
+		}
 		resa, ok := ss.layers[i].resolveCache[addr]
 		if ok {
 			return resa, true
@@ -97,6 +110,7 @@ func (ss *stateSnaps) resolveAddress(addr address.Address) (address.Address, boo
 
 func (ss *stateSnaps) cacheResolveAddress(addr, resa address.Address) {
 	ss.layers[len(ss.layers)-1].resolveCache[addr] = resa
+	ss.lastMaybeNonEmptyResolveCache = len(ss.layers) - 1
 }
 
 func (ss *stateSnaps) getActor(addr address.Address) (*types.Actor, error) {
@@ -160,13 +174,15 @@ func NewStateTree(cst cbor.IpldStore, ver types.StateTreeVersion) (*StateTree, e
 		return nil, err
 	}
 
-	return &StateTree{
+	s := &StateTree{
 		root:    root,
 		info:    info,
 		version: ver,
 		Store:   cst,
 		snaps:   newStateSnaps(),
-	}, nil
+	}
+	s.lookupIDFun = s.lookupIDinternal
+	return s, nil
 }
 
 func LoadStateTree(cst cbor.IpldStore, c cid.Cid) (*StateTree, error) {
@@ -190,13 +206,15 @@ func LoadStateTree(cst cbor.IpldStore, c cid.Cid) (*StateTree, error) {
 			return nil, err
 		}
 
-		return &StateTree{
+		s := &StateTree{
 			root:    nd,
 			info:    root.Info,
 			version: root.Version,
 			Store:   cst,
 			snaps:   newStateSnaps(),
-		}, nil
+		}
+		s.lookupIDFun = s.lookupIDinternal
+		return s, nil
 	default:
 		return nil, xerrors.Errorf("unsupported state tree version: %d", root.Version)
 	}
@@ -213,17 +231,7 @@ func (st *StateTree) SetActor(addr address.Address, act *types.Actor) error {
 	return nil
 }
 
-// LookupID gets the ID address of this actor's `addr` stored in the `InitActor`.
+func (st *StateTree) lookupIDinternal(addr address.Address) (address.Address, error) {
-func (st *StateTree) LookupID(addr address.Address) (address.Address, error) {
-	if addr.Protocol() == address.ID {
-		return addr, nil
-	}
-
-	resa, ok := st.snaps.resolveAddress(addr)
-	if ok {
-		return resa, nil
-	}
-
 	act, err := st.GetActor(init_.Address)
 	if err != nil {
 		return address.Undef, xerrors.Errorf("getting init actor: %w", err)
@@ -241,6 +249,23 @@ func (st *StateTree) LookupID(addr address.Address) (address.Address, error) {
 	if err != nil {
 		return address.Undef, xerrors.Errorf("resolve address %s: %w", addr, err)
 	}
+	return a, err
+}
+
+// LookupID gets the ID address of this actor's `addr` stored in the `InitActor`.
+func (st *StateTree) LookupID(addr address.Address) (address.Address, error) {
+	if addr.Protocol() == address.ID {
+		return addr, nil
+	}
+
+	resa, ok := st.snaps.resolveAddress(addr)
+	if ok {
+		return resa, nil
+	}
+	a, err := st.lookupIDFun(addr)
+	if err != nil {
+		return a, err
+	}
 
 	st.snaps.cacheResolveAddress(addr, a)
 

chain/state/statetree_test.go

@@ -73,6 +73,103 @@ func BenchmarkStateTreeSetFlush(b *testing.B) {
 	}
 }
 
+func TestResolveCache(t *testing.T) {
+	cst := cbor.NewMemCborStore()
+	st, err := NewStateTree(cst, VersionForNetwork(build.NewestNetworkVersion))
+	if err != nil {
+		t.Fatal(err)
+	}
+	nonId := address.NewForTestGetter()()
+	id, _ := address.NewIDAddress(1000)
+
+	st.lookupIDFun = func(a address.Address) (address.Address, error) {
+		if a == nonId {
+			return id, nil
+		}
+		return address.Undef, types.ErrActorNotFound
+	}
+
+	err = st.SetActor(nonId, &types.Actor{Nonce: 1})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	{
+		err = st.Snapshot(context.TODO())
+		if err != nil {
+			t.Fatal(err)
+		}
+		act, err := st.GetActor(nonId)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if act.Nonce != 1 {
+			t.Fatalf("expected nonce 1, got %d", act.Nonce)
+		}
+		err = st.SetActor(nonId, &types.Actor{Nonce: 2})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		act, err = st.GetActor(nonId)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if act.Nonce != 2 {
+			t.Fatalf("expected nonce 2, got %d", act.Nonce)
+		}
+
+		if err := st.Revert(); err != nil {
+			t.Fatal(err)
+		}
+		st.ClearSnapshot()
+	}
+
+	act, err := st.GetActor(nonId)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if act.Nonce != 1 {
+		t.Fatalf("expected nonce 1, got %d", act.Nonce)
+	}
+
+	{
+		err = st.Snapshot(context.TODO())
+		if err != nil {
+			t.Fatal(err)
+		}
+		act, err := st.GetActor(nonId)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if act.Nonce != 1 {
+			t.Fatalf("expected nonce 1, got %d", act.Nonce)
+		}
+		err = st.SetActor(nonId, &types.Actor{Nonce: 2})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		act, err = st.GetActor(nonId)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if act.Nonce != 2 {
+			t.Fatalf("expected nonce 2, got %d", act.Nonce)
+		}
+		st.ClearSnapshot()
+	}
+
+	act, err = st.GetActor(nonId)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if act.Nonce != 2 {
+		t.Fatalf("expected nonce 2, got %d", act.Nonce)
+	}
+
+}
+
 func BenchmarkStateTree10kGetActor(b *testing.B) {
 	cst := cbor.NewMemCborStore()
 	st, err := NewStateTree(cst, VersionForNetwork(build.NewestNetworkVersion))