From d02a2a2cfa6422f8e9d56ac1d0a268279d0b5a8b Mon Sep 17 00:00:00 2001
From: frrist <forrest@protocol.ai>
Date: Wed, 30 Jun 2021 16:40:43 -0700
Subject: [PATCH] polish(statetree): accept a context in statetree diff for
 timeouts

- this operation can take a while, this changes gives users the option
to abort if it takes "too long".
---
 chain/state/statetree.go | 57 ++++++++++++++++++++++------------------
 node/impl/full/state.go  |  2 +-
 2 files changed, 32 insertions(+), 27 deletions(-)

diff --git a/chain/state/statetree.go b/chain/state/statetree.go
index dbf150ecd..8705aeff8 100644
--- a/chain/state/statetree.go
+++ b/chain/state/statetree.go
@@ -547,7 +547,7 @@ 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 (
@@ -555,33 +555,38 @@ func Diff(oldTree, newTree *StateTree) (map[string]types.Actor, error) {
 		buf          = bytes.NewReader(nil)
 	)
 	if err := newTree.root.ForEach(&ncval, func(k string) error {
-		var act types.Actor
+		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)
+			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
+			}
+
+			if found && bytes.Equal(ocval.Raw, ncval.Raw) {
+				return nil // not changed
+			}
+
+			buf.Reset(ncval.Raw)
+			err = act.UnmarshalCBOR(buf)
+			buf.Reset(nil)
+
+			if err != nil {
+				return err
+			}
+
+			out[addr.String()] = act
+
+			return nil
 		}
-
-		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
-		}
-
-		buf.Reset(ncval.Raw)
-		err = act.UnmarshalCBOR(buf)
-		buf.Reset(nil)
-
-		if err != nil {
-			return err
-		}
-
-		out[addr.String()] = act
-
-		return nil
 	}); err != nil {
 		return nil, err
 	}
diff --git a/node/impl/full/state.go b/node/impl/full/state.go
index b3639c5e0..d8545ae13 100644
--- a/node/impl/full/state.go
+++ b/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) {