Add a time bound to the UpgradeRoots pass

CHANGELOG.md

@@ -13,6 +13,10 @@ served it.
 
 ## [Unreleased]
 
+### Added
+
+- `Options.UpgradeTimeFactor` and `Options.UpgradeMaxTime`: opt-in time bounds for the `UpgradeRoots` pass. The pass stops starting re-solves once it has run for factor times the first solve's duration, or for the absolute time, whichever comes first. It returns the best answer so far and never an error. Zero (the default) means no bound. With a bound set, the answer can vary with machine load.
+
 ## [0.22.0] - 2026-11-14
 
 ### Fixed

resolver/export_test.go

@@ -3,6 +3,8 @@
 package resolver
 
 import (
+	"time"
+
 	"github.com/posit-dev/go-pubgrub/solver"
 	"github.com/posit-dev/go-pyresolver/pep440set"
 	"github.com/posit-dev/go-pyresolver/provider"
@@ -37,3 +39,10 @@ func OnSolver(f func(*solver.Solver[provider.Package, pep440set.Set])) (restore
 	onSolver = f
 	return func() { onSolver = nil }
 }
+
+// SetClock replaces the clock the upgrade time bound reads until restore is
+// called. Not safe with parallel tests.
+func SetClock(f func() time.Time) (restore func()) {
+	now = f
+	return func() { now = time.Now }
+}

resolver/resolver.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"slices"
 	"strings"
+	"time"
 
 	"github.com/posit-dev/go-pubgrub/solver"
 	"github.com/posit-dev/go-pyresolver/candidate"
@@ -188,6 +189,23 @@ type Options struct {
 	// It never fails a resolution: if the budget or ctx ends, or no re-solve
 	// helps, the best complete answer so far is returned.
 	UpgradeRoots int
+
+	// UpgradeTimeFactor, when above zero, bounds the UpgradeRoots pass in time:
+	// it stops starting new re-solves once the pass has run for this many times
+	// as long as the first solve took, in this Resolve call. A re-solve already
+	// running is not cut short. Easy solves stay fully upgraded; hard ones are
+	// bounded. Zero (the default) means no time bound.
+	//
+	// The bound reads the clock, so with it set the answer can differ between
+	// runs on a loaded machine. It never fails a resolution: the best answer so
+	// far is returned, as when the budget or ctx ends.
+	UpgradeTimeFactor float64
+
+	// UpgradeMaxTime, when above zero, is an absolute cap on the time the
+	// UpgradeRoots pass may run before it stops starting new re-solves. It
+	// applies together with UpgradeTimeFactor; whichever limit is reached first
+	// stops the pass. Zero (the default) means no cap.
+	UpgradeMaxTime time.Duration
 }
 
 // Resolution is a successful resolution: one version chosen for every package
@@ -351,11 +369,12 @@ func Resolve(
 	idx index.MetadataIndex,
 	opts Options,
 ) (*Resolution, error) {
+	start := now()
 	res, err := resolveOnce(ctx, reqs, idx, opts)
 	if err != nil || opts.UpgradeRoots <= 0 {
 		return res, err
 	}
-	return upgradeRoots(ctx, reqs, idx, opts, res), nil
+	return upgradeRoots(ctx, reqs, idx, opts, res, now().Sub(start)), nil
 }
 
 // providerOptions maps Options to the provider's options for one resolution.
@@ -481,6 +500,9 @@ func resolveOnce(
 // The fallback solve permits no transitive pin, so this is not expected.
 var errUnjustified = errors.New("resolver: no resolution justifies its transitive yank pins")
 
+// now is the clock the upgrade time bound reads. A test replaces it.
+var now = time.Now
+
 // onSolve, when set by a test, is called before each solve.
 var onSolve func()
 

resolver/upgrade.go

@@ -8,6 +8,7 @@ import (
 	"slices"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/posit-dev/go-pyresolver/index"
 	"github.com/posit-dev/go-pyresolver/provider"
@@ -203,15 +204,39 @@ func (r *upRoot) bounds(pin version.Version) []version.Version {
 	return out
 }
 
+// upgradeTimeLimit is how long the pass may run before it stops starting
+// re-solves, from the options and the first solve's duration. ok is false when
+// neither bound is set.
+func upgradeTimeLimit(opts Options, baseDur time.Duration) (limit time.Duration, ok bool) {
+	if opts.UpgradeMaxTime > 0 {
+		limit, ok = opts.UpgradeMaxTime, true
+	}
+	if opts.UpgradeTimeFactor > 0 {
+		// A huge factor is no bound; it must not overflow the Duration.
+		if f := float64(baseDur) * opts.UpgradeTimeFactor; f < math.MaxInt64 {
+			if d := time.Duration(f); !ok || d < limit {
+				limit, ok = d, true
+			}
+		}
+	}
+	return limit, ok
+}
+
 // upgradeRoots is the pass behind Options.UpgradeRoots. base is a valid answer
-// to reqs; the result is base or a valid answer that scores strictly better.
+// to reqs, and baseDur is how long it took; the result is base or a valid
+// answer that scores strictly better.
 func upgradeRoots(
 	ctx context.Context,
 	reqs []requirement.Requirement,
 	idx index.MetadataIndex,
 	opts Options,
 	base *Resolution,
+	baseDur time.Duration,
 ) *Resolution {
+	passStart := now()
+	limit, timed := upgradeTimeLimit(opts, baseDur)
+	// timeUp is checked before each re-solve; one already running is not cut.
+	timeUp := func() bool { return timed && now().Sub(passStart) >= limit }
 	u := &usability{
 		p:    provider.New(ctx, idx, providerOptions(reqs, opts)),
 		memo: map[usableKey]bool{},
@@ -256,7 +281,7 @@ func upgradeRoots(
 		return got
 	}
 
-	for budget > 0 && ctx.Err() == nil {
+	for budget > 0 && ctx.Err() == nil && !timeUp() {
 		// The untried root furthest behind (at least one level); the first by
 		// name on a tie. A root behind only in versions is not tried.
 		var pick *upRoot
@@ -304,7 +329,7 @@ func upgradeRoots(
 		lo := accept(got) // highest level index known to work
 		hi := len(bounds) // lowest level index known to fail
 		triedTop := false
-		for hi-lo > 1 && budget > 0 && ctx.Err() == nil {
+		for hi-lo > 1 && budget > 0 && ctx.Err() == nil && !timeUp() {
 			mid := hi - 1
 			if triedTop {
 				mid = (lo + hi) / 2

resolver/upgrade_time_test.go

@@ -0,0 +1,153 @@
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+package resolver_test
+
+import (
+	"context"
+	"math"
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/posit-dev/go-pyresolver/index"
+	"github.com/posit-dev/go-pyresolver/resolver"
+)
+
+const solveCost = 10 * time.Second
+
+// resolveTimed resolves reqs on a fake clock where each solve takes solveCost.
+// It returns the pins and the number of re-solves started.
+func resolveTimed(t *testing.T, idx *index.MockIndex, reqs []string, mutate func(*resolver.Options)) (map[string]string, int) {
+	t.Helper()
+	clock := time.Unix(0, 0)
+	defer resolver.SetClock(func() time.Time { return clock })()
+	solves := 0
+	defer resolver.OnSolve(func() { solves++; clock = clock.Add(solveCost) })()
+	opts := testOptions(t)
+	opts.UpgradeRoots = 40
+	mutate(&opts)
+	res, err := resolver.Resolve(context.Background(), mustRequirements(t, reqs...), idx, opts)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return pins(t, res), solves - 1
+}
+
+// The first re-solve raises a. Later ones raise nothing.
+func earlyGain() (*index.MockIndex, []string) {
+	return upgradeFixture().
+		AddVersion("c", "1.0").AddVersion("c", "2.0").AddVersion("c", "3.0").
+		AddVersion("d", "1.0").AddVersion("d", "2.0", "c<2"), []string{"a", "b", "c", "d"}
+}
+
+// The pass tries a first (furthest behind) and fails: b==5.0 needs a<2, so the
+// re-solve has no answer. Only the second re-solve raises z.
+func lateGain() (*index.MockIndex, []string) {
+	idx := index.NewMockIndex("test").
+		AddVersion("a", "1.0").AddVersion("a", "2.0").AddVersion("a", "3.0").AddVersion("a", "4.0").
+		AddVersion("b", "5.0", "a<2").
+		AddVersion("z", "1.0").AddVersion("z", "2.0").AddVersion("z", "3.0").
+		AddVersion("y", "1.0").AddVersion("y", "2.0", "z<2")
+	return idx, []string{"a", "b==5.0", "y", "z"}
+}
+
+var (
+	earlyA   = map[string]string{"a": "3.0", "b": "1.0", "c": "1.0", "d": "2.0"}
+	lateBase = map[string]string{"a": "1.0", "b": "5.0", "y": "2.0", "z": "1.0"}
+	lateZ    = map[string]string{"a": "1.0", "b": "5.0", "y": "1.0", "z": "3.0"}
+)
+
+func TestUpgradeTimeFactorZeroIsUnbounded(t *testing.T) {
+	idx, reqs := lateGain()
+	got, attempts := resolveTimed(t, idx, reqs, func(*resolver.Options) {})
+	if !reflect.DeepEqual(got, lateZ) || attempts != 3 {
+		t.Errorf("got %v after %d attempts, want %v after 3", got, attempts, lateZ)
+	}
+}
+
+// Each re-solve takes as long as the first solve, so factor F allows ceil(F)
+// attempts: one starts while the pass has run for less than F solves.
+func TestUpgradeTimeFactorStopsAfterExpectedAttempts(t *testing.T) {
+	type want struct {
+		attempts int
+		pins     map[string]string
+	}
+	for _, tc := range []struct {
+		name    string
+		fixture func() (*index.MockIndex, []string)
+		factor  float64
+		want    want
+	}{
+		{"late 0.5", lateGain, 0.5, want{1, lateBase}}, // 0 < 0.5 starts one; 1 >= 0.5 stops
+		{"late 1", lateGain, 1, want{1, lateBase}},
+		{"late 1.5", lateGain, 1.5, want{2, lateZ}},
+		{"late 2", lateGain, 2, want{2, lateZ}},
+		{"late 50", lateGain, 50, want{3, lateZ}}, // the pass ends by itself
+		{"late +Inf", lateGain, math.Inf(1), want{3, lateZ}},
+		{"early 1", earlyGain, 1, want{1, earlyA}}, // best so far is kept
+		{"early 50", earlyGain, 50, want{3, earlyA}},
+	} {
+		idx, reqs := tc.fixture()
+		got, attempts := resolveTimed(t, idx, reqs, func(o *resolver.Options) { o.UpgradeTimeFactor = tc.factor })
+		if !reflect.DeepEqual(got, tc.want.pins) || attempts != tc.want.attempts {
+			t.Errorf("%s: got %v after %d attempts, want %v after %d",
+				tc.name, got, attempts, tc.want.pins, tc.want.attempts)
+		}
+	}
+}
+
+func TestUpgradeMaxTimeStopsPass(t *testing.T) {
+	idx, reqs := lateGain()
+	got, attempts := resolveTimed(t, idx, reqs, func(o *resolver.Options) { o.UpgradeMaxTime = solveCost })
+	if !reflect.DeepEqual(got, lateBase) || attempts != 1 {
+		t.Errorf("got %v after %d attempts, want %v after 1", got, attempts, lateBase)
+	}
+}
+
+// The tighter of the two bounds wins.
+func TestUpgradeTimeBoundsCombine(t *testing.T) {
+	idx, reqs := lateGain()
+	got, attempts := resolveTimed(t, idx, reqs, func(o *resolver.Options) {
+		o.UpgradeTimeFactor = 10
+		o.UpgradeMaxTime = solveCost
+	})
+	if !reflect.DeepEqual(got, lateBase) || attempts != 1 {
+		t.Errorf("max tighter: got %v after %d attempts", got, attempts)
+	}
+	got, attempts = resolveTimed(t, idx, reqs, func(o *resolver.Options) {
+		o.UpgradeTimeFactor = 1
+		o.UpgradeMaxTime = time.Hour
+	})
+	if !reflect.DeepEqual(got, lateBase) || attempts != 1 {
+		t.Errorf("factor tighter: got %v after %d attempts", got, attempts)
+	}
+}
+
+// A bound spent before the first re-solve returns the first answer, no error.
+func TestUpgradeTimeBoundBeforeFirstAttemptReturnsBase(t *testing.T) {
+	idx := upgradeFixture()
+	clock := time.Unix(0, 0)
+	defer resolver.SetClock(func() time.Time { return clock })()
+	opts := testOptions(t)
+	opts.UpgradeRoots = 12
+	opts.UpgradeTimeFactor = 1 // the fake clock never moves: base took 0
+	res, err := resolver.Resolve(context.Background(), mustRequirements(t, "a", "b"), idx, opts)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if want := map[string]string{"a": "1.0", "b": "2.0"}; !reflect.DeepEqual(pins(t, res), want) {
+		t.Errorf("got %v, want the base answer %v", pins(t, res), want)
+	}
+}
+
+func TestUpgradeTimeBoundIsDeterministicOnAFakeClock(t *testing.T) {
+	idx, reqs := earlyGain()
+	set := func(o *resolver.Options) { o.UpgradeTimeFactor = 1 }
+	first, n1 := resolveTimed(t, idx, reqs, set)
+	for range 5 {
+		got, n := resolveTimed(t, idx, reqs, set)
+		if !reflect.DeepEqual(got, first) || n != n1 {
+			t.Fatalf("run differs: %v/%d vs %v/%d", got, n, first, n1)
+		}
+	}
+}