traversal: budget tests, well-typed errors, more error info, and fix …

…off-by-one.
ipld · Sep 28, 2021 · e917bba · e917bba
1 parent 3ac6a8f
commit e917bba
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 17 deletions.
diff --git a/traversal/fns.go b/traversal/fns.go
@@ -2,6 +2,7 @@ package traversal
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/ipld/go-ipld-prime/datamodel"
 	"github.com/ipld/go-ipld-prime/linking"
@@ -80,3 +81,17 @@ type SkipMe struct{}
 func (SkipMe) Error() string {
 	return "skip"
 }
+
+type ErrBudgetExceeded struct {
+	BudgetKind string // "node"|"link"
+	Path       datamodel.Path
+	Link       datamodel.Link // only present if BudgetKind=="link"
+}
+
+func (e *ErrBudgetExceeded) Error() string {
+	msg := fmt.Sprintf("traversal budget exceeded: budget for %ss reached zero as we reached path %q", e.BudgetKind, e.Path)
+	if e.Link != nil {
+		msg += fmt.Sprintf(" (link: %q)", e.Link)
+	}
+	return msg
+}
diff --git a/traversal/focus.go b/traversal/focus.go
@@ -91,7 +91,7 @@ func (prog *Progress) get(n datamodel.Node, p datamodel.Path, trackProgress bool
 		if prog.Budget != nil {
 			prog.Budget.NodeBudget--
 			if prog.Budget.NodeBudget <= 0 {
-				return nil, fmt.Errorf("traversal budget for nodes visited exceeded")
+				return nil, &ErrBudgetExceeded{BudgetKind: "node", Path: prog.Path}
 			}
 		}
 		// Traverse the segment.
@@ -119,15 +119,15 @@ func (prog *Progress) get(n datamodel.Node, p datamodel.Path, trackProgress bool
 		}
 		// Dereference any links.
 		for n.Kind() == datamodel.Kind_Link {
+			lnk, _ := n.AsLink()
 			// Check the budget!
 			if prog.Budget != nil {
-				prog.Budget.LinkBudget--
 				if prog.Budget.LinkBudget <= 0 {
-					return nil, fmt.Errorf("traversal budget for links exceeded")
+					return nil, &ErrBudgetExceeded{BudgetKind: "link", Path: prog.Path, Link: lnk}
 				}
+				prog.Budget.LinkBudget--
 			}
 			// Put together the context info we'll offer to the loader and prototypeChooser.
-			lnk, _ := n.AsLink()
 			lnkCtx := linking.LinkContext{
 				Ctx:        prog.Cfg.Ctx,
 				LinkPath:   p.Truncate(i),
@@ -218,10 +218,10 @@ func (prog Progress) focusedTransform(n datamodel.Node, na datamodel.NodeAssembl
 	seg, p2 := p.Shift()
 	// Check the budget!
 	if prog.Budget != nil {
-		prog.Budget.NodeBudget--
 		if prog.Budget.NodeBudget <= 0 {
-			return fmt.Errorf("traversal budget for nodes visited exceeded")
+			return &ErrBudgetExceeded{BudgetKind: "node", Path: prog.Path}
 		}
+		prog.Budget.NodeBudget--
 	}
 	// Special branch for if we've entered createParent mode in an earlier step.
 	//  This needs slightly different logic because there's no prior node to reference
@@ -341,12 +341,13 @@ func (prog Progress) focusedTransform(n datamodel.Node, na datamodel.NodeAssembl
 		}
 		return la.Finish()
 	case datamodel.Kind_Link:
+		lnk, _ := n.AsLink()
 		// Check the budget!
 		if prog.Budget != nil {
-			prog.Budget.LinkBudget--
 			if prog.Budget.LinkBudget <= 0 {
-				return fmt.Errorf("traversal budget for links exceeded")
+				return &ErrBudgetExceeded{BudgetKind: "link", Path: prog.Path, Link: lnk}
 			}
+			prog.Budget.LinkBudget--
 		}
 		// Put together the context info we'll offer to the loader and prototypeChooser.
 		lnkCtx := linking.LinkContext{
@@ -355,7 +356,6 @@ func (prog Progress) focusedTransform(n datamodel.Node, na datamodel.NodeAssembl
 			LinkNode:   n,
 			ParentNode: nil, // TODO inconvenient that we don't have this.  maybe this whole case should be a helper function.
 		}
-		lnk, _ := n.AsLink()
 		// Pick what in-memory format we will build.
 		np, err := prog.Cfg.LinkTargetNodePrototypeChooser(lnk, lnkCtx)
 		if err != nil {

diff --git a/traversal/walk.go b/traversal/walk.go
@@ -88,10 +88,10 @@ func (prog Progress) WalkAdv(n datamodel.Node, s selector.Selector, fn AdvVisitF
 func (prog Progress) walkAdv(n datamodel.Node, s selector.Selector, fn AdvVisitFn) error {
 	// Check the budget!
 	if prog.Budget != nil {
-		prog.Budget.NodeBudget--
 		if prog.Budget.NodeBudget <= 0 {
-			return fmt.Errorf("traversal budget for nodes visited exceeded")
+			return &ErrBudgetExceeded{BudgetKind: "node", Path: prog.Path}
 		}
+		prog.Budget.NodeBudget--
 	}
 	// Decide if this node is matched -- do callbacks as appropriate.
 	if s.Decide(n) {
@@ -190,18 +190,18 @@ func (prog Progress) walkAdv_iterateSelective(n datamodel.Node, attn []datamodel
 }
 
 func (prog Progress) loadLink(v datamodel.Node, parent datamodel.Node) (datamodel.Node, error) {
+	lnk, err := v.AsLink()
+	if err != nil {
+		return nil, err
+	}
 	// Check the budget!
 	if prog.Budget != nil {
-		prog.Budget.LinkBudget--
 		if prog.Budget.LinkBudget <= 0 {
-			return nil, fmt.Errorf("traversal budget for links exceeded")
+			return nil, &ErrBudgetExceeded{BudgetKind: "link", Path: prog.Path, Link: lnk}
 		}
+		prog.Budget.LinkBudget--
 	}
 	// Put together the context info we'll offer to the loader and prototypeChooser.
-	lnk, err := v.AsLink()
-	if err != nil {
-		return nil, err
-	}
 	lnkCtx := linking.LinkContext{
 		Ctx:        prog.Cfg.Ctx,
 		LinkPath:   prog.Path,

diff --git a/traversal/walk_test.go b/traversal/walk_test.go
@@ -3,6 +3,7 @@ package traversal_test
 import (
 	"testing"
 
+	qt "github.com/frankban/quicktest"
 	. "github.com/warpfork/go-wish"
 
 	_ "github.com/ipld/go-ipld-prime/codec/dagjson"
@@ -257,3 +258,67 @@ func TestWalkMatching(t *testing.T) {
 		Wish(t, order, ShouldEqual, 7)
 	})
 }
+
+func TestWalkBudgets(t *testing.T) {
+	ssb := builder.NewSelectorSpecBuilder(basicnode.Prototype.Any)
+	t.Run("node-budget-halts", func(t *testing.T) {
+		ss := ssb.ExploreFields(func(efsb builder.ExploreFieldsSpecBuilder) {
+			efsb.Insert("foo", ssb.Matcher())
+			efsb.Insert("bar", ssb.Matcher())
+		})
+		s, err := ss.Selector()
+		qt.Assert(t, err, qt.Equals, nil)
+		var order int
+		prog := traversal.Progress{}
+		prog.Budget = &traversal.Budget{
+			NodeBudget: 2, // should reach root, then "foo", then stop.
+		}
+		err = prog.WalkMatching(middleMapNode, s, func(prog traversal.Progress, n datamodel.Node) error {
+			switch order {
+			case 0:
+				qt.Assert(t, n, qt.CmpEquals(), basicnode.NewBool(true))
+				qt.Assert(t, prog.Path.String(), qt.Equals, "foo")
+			}
+			order++
+			return nil
+		})
+		qt.Check(t, order, qt.Equals, 1) // because it should've stopped early
+		qt.Assert(t, err, qt.Not(qt.Equals), nil)
+		qt.Check(t, err.Error(), qt.Equals, `traversal budget exceeded: budget for nodes reached zero as we reached path "bar"`)
+	})
+	t.Run("link-budget-halts", func(t *testing.T) {
+		ss := ssb.ExploreAll(ssb.Matcher())
+		s, err := ss.Selector()
+		qt.Assert(t, err, qt.Equals, nil)
+		var order int
+		lsys := cidlink.DefaultLinkSystem()
+		lsys.StorageReadOpener = (&store).OpenRead
+		err = traversal.Progress{
+			Cfg: &traversal.Config{
+				LinkSystem:                     lsys,
+				LinkTargetNodePrototypeChooser: basicnode.Chooser,
+			},
+			Budget: &traversal.Budget{
+				NodeBudget: 9000,
+				LinkBudget: 3,
+			},
+		}.WalkMatching(middleListNode, s, func(prog traversal.Progress, n datamodel.Node) error {
+			switch order {
+			case 0:
+				qt.Assert(t, n, qt.CmpEquals(), basicnode.NewString("alpha"))
+				qt.Assert(t, prog.Path.String(), qt.Equals, "0")
+			case 1:
+				qt.Assert(t, n, qt.CmpEquals(), basicnode.NewString("alpha"))
+				qt.Assert(t, prog.Path.String(), qt.Equals, "1")
+			case 2:
+				qt.Assert(t, n, qt.CmpEquals(), basicnode.NewString("beta"))
+				qt.Assert(t, prog.Path.String(), qt.Equals, "2")
+			}
+			order++
+			return nil
+		})
+		qt.Check(t, order, qt.Equals, 3)
+		qt.Assert(t, err, qt.Not(qt.Equals), nil)
+		qt.Check(t, err.Error(), qt.Equals, `traversal budget exceeded: budget for links reached zero as we reached path "3" (link: "baguqeeyexkjwnfy")`)
+	})
+}