Merge pull request #6639 from filecoin-project/frrist/context-in-stat…

…etree-diff

polish(statetree): accept a context in statetree diff for timeouts

chain/state/statetree.go

@@ -547,41 +547,46 @@ func (st *StateTree) Version() types.StateTreeVersion {
 	return st.version
 }
 
-func Diff(oldTree, newTree *StateTree) (map[string]types.Actor, error) {
+func Diff(ctx context.Context, oldTree, newTree *StateTree) (map[string]types.Actor, error) {
 	out := map[string]types.Actor{}
 
 	var (
 		ncval, ocval cbg.Deferred
 		buf          = bytes.NewReader(nil)
 	)
 	if err := newTree.root.ForEach(&ncval, func(k string) error {
-		var act types.Actor
-
-		addr, err := address.NewFromBytes([]byte(k))
-		if err != nil {
-			return xerrors.Errorf("address in state tree was not valid: %w", err)
-		}
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			var act types.Actor
+
+			addr, err := address.NewFromBytes([]byte(k))
+			if err != nil {
+				return xerrors.Errorf("address in state tree was not valid: %w", err)
+			}
 
-		found, err := oldTree.root.Get(abi.AddrKey(addr), &ocval)
-		if err != nil {
-			return err
-		}
+			found, err := oldTree.root.Get(abi.AddrKey(addr), &ocval)
+			if err != nil {
+				return err
+			}
 
-		if found && bytes.Equal(ocval.Raw, ncval.Raw) {
-			return nil // not changed
-		}
+			if found && bytes.Equal(ocval.Raw, ncval.Raw) {
+				return nil // not changed
+			}
 
-		buf.Reset(ncval.Raw)
-		err = act.UnmarshalCBOR(buf)
-		buf.Reset(nil)
+			buf.Reset(ncval.Raw)
+			err = act.UnmarshalCBOR(buf)
+			buf.Reset(nil)
 
-		if err != nil {
-			return err
-		}
+			if err != nil {
+				return err
+			}
 
-		out[addr.String()] = act
+			out[addr.String()] = act
 
-		return nil
+			return nil
+		}
 	}); err != nil {
 		return nil, err
 	}

node/impl/full/state.go

@@ -705,7 +705,7 @@ func (a *StateAPI) StateChangedActors(ctx context.Context, old cid.Cid, new cid.
 		return nil, xerrors.Errorf("failed to load new state tree: %w", err)
 	}
 
-	return state.Diff(oldTree, newTree)
+	return state.Diff(ctx, oldTree, newTree)
 }
 
 func (a *StateAPI) StateMinerSectorCount(ctx context.Context, addr address.Address, tsk types.TipSetKey) (api.MinerSectors, error) {