traversal: implement monotonically decrementing budgets.

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"
@@ -36,6 +37,7 @@ type Progress struct {
 		Path datamodel.Path
 		Link datamodel.Link
 	}
+	Budget *Budget // If present, tracks "budgets" for how many more steps we're willing to take before we should halt.
 }
 
 type Config struct {
@@ -44,6 +46,18 @@ type Config struct {
 	LinkTargetNodePrototypeChooser LinkTargetNodePrototypeChooser // Chooser for Node implementations to produce during automatic link traversal.
 }
 
+type Budget struct {
+	// Fields below are described as "monotonically-decrementing", because that's what the traversal library will do with them,
+	// but they are user-accessable and can be reset to higher numbers again by code in the visitor callbacks.  This is not recommended (why?), but possible.
+
+	// If you set any budgets (by having a non-nil Progress.Budget field), you must set some value for all of them.
+	// Traversal halts when _any_ of the budgets reaches zero.
+	// The max value of an int (math.MaxInt64) is acceptable for any budget you don't care about.
+
+	NodeBudget int64 // A monotonically-decrementing "budget" for how many more nodes we're willing to visit before halting.
+	LinkBudget int64 // A monotonically-decrementing "budget" for how many more links we're willing to load before halting.  (This is not aware of any caching; it's purely in terms of links encountered and traversed.)
+}
+
 // LinkTargetNodePrototypeChooser is a function that returns a NodePrototype based on
 // the information in a Link and/or its LinkContext.
 //
@@ -67,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 while on path %q", e.BudgetKind, e.Path)
+	if e.Link != nil {
+		msg += fmt.Sprintf(" (link: %q)", e.Link)
+	}
+	return msg
+}

traversal/focus.go

@@ -87,6 +87,13 @@ func (prog *Progress) get(n datamodel.Node, p datamodel.Path, trackProgress bool
 	segments := p.Segments()
 	var prev datamodel.Node // for LinkContext
 	for i, seg := range segments {
+		// Check the budget!
+		if prog.Budget != nil {
+			prog.Budget.NodeBudget--
+			if prog.Budget.NodeBudget <= 0 {
+				return nil, &ErrBudgetExceeded{BudgetKind: "node", Path: prog.Path}
+			}
+		}
 		// Traverse the segment.
 		switch n.Kind() {
 		case datamodel.Kind_Invalid:
@@ -113,6 +120,14 @@ 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 {
+				if prog.Budget.LinkBudget <= 0 {
+					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.
 			lnkCtx := linking.LinkContext{
 				Ctx:        prog.Cfg.Ctx,
 				LinkPath:   p.Truncate(i),
@@ -201,6 +216,13 @@ func (prog Progress) focusedTransform(n datamodel.Node, na datamodel.NodeAssembl
 		return na.AssignNode(n2)
 	}
 	seg, p2 := p.Shift()
+	// Check the budget!
+	if prog.Budget != nil {
+		if prog.Budget.NodeBudget <= 0 {
+			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
 	//   (and we wouldn't want to waste time creating a dummy one).
@@ -319,13 +341,21 @@ 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 {
+			if prog.Budget.LinkBudget <= 0 {
+				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{
 			Ctx:        prog.Cfg.Ctx,
 			LinkPath:   prog.Path,
 			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 {

traversal/walk.go

@@ -86,6 +86,14 @@ 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 {
+		if prog.Budget.NodeBudget <= 0 {
+			return &ErrBudgetExceeded{BudgetKind: "node", Path: prog.Path}
+		}
+		prog.Budget.NodeBudget--
+	}
+	// Decide if this node is matched -- do callbacks as appropriate.
 	if s.Decide(n) {
 		if err := fn(prog, n, VisitReason_SelectionMatch); err != nil {
 			return err
@@ -95,12 +103,14 @@ func (prog Progress) walkAdv(n datamodel.Node, s selector.Selector, fn AdvVisitF
 			return err
 		}
 	}
+	// If we're handling scalars (e.g. not maps and lists) we can return now.
 	nk := n.Kind()
 	switch nk {
 	case datamodel.Kind_Map, datamodel.Kind_List: // continue
 	default:
 		return nil
 	}
+	// For maps and lists: recurse (in one of two ways, depending on if the selector also states specific interests).
 	attn := s.Interests()
 	if attn == nil {
 		return prog.walkAdv_iterateAll(n, s, fn)
@@ -184,6 +194,14 @@ func (prog Progress) loadLink(v datamodel.Node, parent datamodel.Node) (datamode
 	if err != nil {
 		return nil, err
 	}
+	// Check the budget!
+	if prog.Budget != nil {
+		if prog.Budget.LinkBudget <= 0 {
+			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.
 	lnkCtx := linking.LinkContext{
 		Ctx:        prog.Cfg.Ctx,
 		LinkPath:   prog.Path,

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 while on 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 while on path "3" (link: "baguqeeyexkjwnfy")`)
+	})
+}