Merge commit from fork

core/images/handlers.go

@@ -23,6 +23,7 @@ import (
 	"slices"
 	"sort"
 	"strings"
+	"sync"
 
 	"github.com/containerd/containerd/v2/core/content"
 	"github.com/containerd/errdefs"
@@ -32,6 +33,16 @@ import (
 	"golang.org/x/sync/semaphore"
 )
 
+const (
+	// defaultMaxConcurrency bounds the number of concurrent handler
+	// goroutines when Dispatch is called with a nil limiter.
+	defaultMaxConcurrency = 32
+
+	// maxReferences caps the references in a single Dispatch or Walk.
+	// Duplicate references count toward this limit.
+	maxReferences = 10_000
+)
+
 var (
 	// ErrSkipDesc is used to skip processing of a descriptor and
 	// its descendants.
@@ -88,9 +99,26 @@ func Handlers(handlers ...Handler) HandlerFunc {
 //
 // This differs from dispatch in that each sibling resource is considered
 // synchronously.
+//
+// Each call is limited to 10,000 references, including duplicates.
+// Exceeding this limit returns an error wrapping
+// [errdefs.ErrResourceExhausted].
 func Walk(ctx context.Context, handler Handler, descs ...ocispec.Descriptor) error {
-	for _, desc := range descs {
+	w := &walker{}
+	return w.walk(ctx, handler, descs...)
+}
 
+type walker struct {
+	referenceCount int
+}
+
+func (w *walker) walk(ctx context.Context, handler Handler, descs ...ocispec.Descriptor) error {
+	if w.referenceCount+len(descs) > maxReferences {
+		return fmt.Errorf("too many descriptors (limit %d): %w", maxReferences, errdefs.ErrResourceExhausted)
+	}
+	w.referenceCount += len(descs)
+
+	for _, desc := range descs {
 		children, err := handler.Handle(ctx, desc)
 		if err != nil {
 			if errors.Is(err, ErrSkipDesc) {
@@ -100,7 +128,7 @@ func Walk(ctx context.Context, handler Handler, descs ...ocispec.Descriptor) err
 		}
 
 		if len(children) > 0 {
-			if err := Walk(ctx, handler, children...); err != nil {
+			if err := w.walk(ctx, handler, children...); err != nil {
 				return err
 			}
 		}
@@ -146,45 +174,81 @@ func WalkNotEmpty(ctx context.Context, handler Handler, descs ...ocispec.Descrip
 // handler may return `ErrSkipDesc` to signal to the dispatcher to not traverse
 // any children.
 //
-// A concurrency limiter can be passed in to limit the number of concurrent
-// handlers running. When limiter is nil, there is no limit.
+// The limiter bounds concurrent handlers. When limiter is nil,
+// the limit is 32.
 //
 // Typically, this function will be used with `FetchHandler`, often composed
 // with other handlers.
 //
 // If any handler returns an error, the dispatch session will be canceled.
+//
+// Each call is limited to 10,000 references, including duplicates.
+// Exceeding this limit returns an error wrapping
+// [errdefs.ErrResourceExhausted].
 func Dispatch(ctx context.Context, handler Handler, limiter *semaphore.Weighted, descs ...ocispec.Descriptor) error {
-	eg, ctx2 := errgroup.WithContext(ctx)
-	for _, desc := range descs {
-		if limiter != nil {
-			if err := limiter.Acquire(ctx, 1); err != nil {
-				return err
-			}
-		}
-
-		eg.Go(func() error {
-			desc := desc
-
-			children, err := handler.Handle(ctx2, desc)
-			if limiter != nil {
-				limiter.Release(1)
-			}
-			if err != nil {
-				if errors.Is(err, ErrSkipDesc) {
-					return nil // don't traverse the children.
-				}
-				return err
-			}
-
-			if len(children) > 0 {
-				return Dispatch(ctx2, handler, limiter, children...)
-			}
-
-			return nil
-		})
+	if len(descs) == 0 {
+		return nil
+	}
+	if err := ctx.Err(); err != nil {
+		return err
 	}
 
-	return eg.Wait()
+	if limiter == nil {
+		limiter = semaphore.NewWeighted(defaultMaxConcurrency)
+	}
+
+	var (
+		mu             sync.Mutex
+		next           []ocispec.Descriptor
+		referenceCount int
+	)
+
+	admit := func(descs []ocispec.Descriptor) error {
+		mu.Lock()
+		defer mu.Unlock()
+		if referenceCount+len(descs) > maxReferences {
+			return fmt.Errorf("too many descriptors (limit %d): %w", maxReferences, errdefs.ErrResourceExhausted)
+		}
+		referenceCount += len(descs)
+		next = append(next, descs...)
+		return nil
+	}
+
+	if err := admit(descs); err != nil {
+		return err
+	}
+
+	for len(next) > 0 {
+		level := next
+		next = nil
+
+		eg, egCtx := errgroup.WithContext(ctx)
+		for _, desc := range level {
+			// Acquire here to bound the number of handler goroutines.
+			if err := limiter.Acquire(egCtx, 1); err != nil {
+				break
+			}
+			eg.Go(func() error {
+				defer limiter.Release(1)
+				children, err := handler.Handle(egCtx, desc)
+				if err != nil {
+					if errors.Is(err, ErrSkipDesc) {
+						return nil // don't traverse the children.
+					}
+					return err
+				}
+				return admit(children)
+			})
+		}
+		if err := eg.Wait(); err != nil {
+			return err
+		}
+		// Acquire can fail on cancellation even if every handler returns nil.
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
 // ChildrenHandler decodes well-known manifest types and returns their children.

core/images/handlers_test.go

@@ -0,0 +1,709 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package images
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/containerd/errdefs"
+	digest "github.com/opencontainers/go-digest"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+	"golang.org/x/sync/semaphore"
+)
+
+func newDesc(id string) ocispec.Descriptor {
+	return ocispec.Descriptor{
+		Digest:    digest.FromString(id),
+		MediaType: ocispec.MediaTypeImageLayer,
+		Size:      1,
+	}
+}
+
+func digestSet(digests []digest.Digest) map[digest.Digest]struct{} {
+	set := make(map[digest.Digest]struct{}, len(digests))
+	for _, d := range digests {
+		set[d] = struct{}{}
+	}
+	return set
+}
+
+func TestWalkTreeTraversal(t *testing.T) {
+	root := newDesc("root")
+	child1 := newDesc("child1")
+	child2 := newDesc("child2")
+	grandchild1 := newDesc("grandchild1")
+
+	graph := map[digest.Digest][]ocispec.Descriptor{
+		root.Digest:   {child1, child2},
+		child1.Digest: {grandchild1},
+	}
+
+	var visited []digest.Digest
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		visited = append(visited, desc.Digest)
+		return graph[desc.Digest], nil
+	})
+
+	err := Walk(context.Background(), handler, root)
+	require.NoError(t, err)
+
+	want := []digest.Digest{root.Digest, child1.Digest, grandchild1.Digest, child2.Digest}
+	assert.Equal(t, want, visited)
+}
+
+func TestWalkRepeatedReferences(t *testing.T) {
+	root := newDesc("root")
+	childA := newDesc("A")
+	childB := newDesc("B")
+	shared := newDesc("shared")
+
+	graph := map[digest.Digest][]ocispec.Descriptor{
+		root.Digest:   {childA, childB},
+		childA.Digest: {shared},
+		childB.Digest: {shared},
+	}
+
+	tests := []struct {
+		name            string
+		walk            func(context.Context, Handler, ...ocispec.Descriptor) error
+		wantTotalCalls  int
+		wantSharedCalls int
+	}{
+		{
+			name:            "Walk visits every reference",
+			walk:            Walk,
+			wantTotalCalls:  5,
+			wantSharedCalls: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var totalCalls, sharedCalls int
+			handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+				totalCalls++
+				if desc.Digest == shared.Digest {
+					sharedCalls++
+				}
+				return graph[desc.Digest], nil
+			})
+
+			err := tt.walk(context.Background(), handler, root)
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantTotalCalls, totalCalls)
+			assert.Equal(t, tt.wantSharedCalls, sharedCalls)
+		})
+	}
+}
+
+func TestWalkErrSkipDesc(t *testing.T) {
+	root := newDesc("root")
+	child1 := newDesc("child1")
+	child2 := newDesc("child2")
+	grandchild := newDesc("grandchild")
+
+	graph := map[digest.Digest][]ocispec.Descriptor{
+		root.Digest:   {child1, child2},
+		child1.Digest: {grandchild},
+	}
+
+	var visited []digest.Digest
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		visited = append(visited, desc.Digest)
+		if desc.Digest == child1.Digest {
+			return graph[desc.Digest], ErrSkipDesc
+		}
+		return graph[desc.Digest], nil
+	})
+
+	err := Walk(context.Background(), handler, root)
+	require.NoError(t, err)
+
+	set := digestSet(visited)
+	assert.Contains(t, set, root.Digest)
+	assert.Contains(t, set, child1.Digest)
+	assert.Contains(t, set, child2.Digest)
+	assert.NotContains(t, set, grandchild.Digest, "grandchild must not be visited when its parent returns ErrSkipDesc")
+}
+
+func TestWalkReferenceLimit(t *testing.T) {
+	root := newDesc("root")
+	handlerReturning := func(childCount int) Handler {
+		children := make([]ocispec.Descriptor, childCount)
+		for i := range children {
+			children[i] = newDesc(fmt.Sprintf("node-%d", i))
+		}
+		return HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+			if desc.Digest == root.Digest {
+				return children, nil
+			}
+			return nil, nil
+		})
+	}
+
+	require.NoError(t, Walk(context.Background(), handlerReturning(maxReferences-1), root))
+	require.ErrorIs(t, Walk(context.Background(), handlerReturning(maxReferences), root), errdefs.ErrResourceExhausted)
+}
+
+func TestWalkRejectsChildrenBeforeVisiting(t *testing.T) {
+	root := newDesc("root")
+	children := make([]ocispec.Descriptor, maxReferences)
+	for i := range children {
+		children[i] = newDesc(fmt.Sprintf("node-%d", i))
+	}
+	var childVisits int
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if desc.Digest == root.Digest {
+			return children, nil
+		}
+		childVisits++
+		return nil, nil
+	})
+
+	require.ErrorIs(t, Walk(context.Background(), handler, root), errdefs.ErrResourceExhausted)
+	assert.Equal(t, 0, childVisits, "children of an over-limit slice must not be visited")
+}
+
+func TestWalkErrorPropagation(t *testing.T) {
+	root := newDesc("root")
+	child1 := newDesc("child1")
+	child2 := newDesc("child2")
+
+	graph := map[digest.Digest][]ocispec.Descriptor{
+		root.Digest: {child1, child2},
+	}
+
+	errTest := errors.New("handler error")
+
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if desc.Digest == child1.Digest {
+			return nil, errTest
+		}
+		return graph[desc.Digest], nil
+	})
+
+	err := Walk(context.Background(), handler, root)
+	require.ErrorIs(t, err, errTest)
+}
+
+func TestWalkDuplicateRoots(t *testing.T) {
+	root := newDesc("root")
+
+	tests := []struct {
+		name      string
+		walk      func(context.Context, Handler, ...ocispec.Descriptor) error
+		wantCalls int
+	}{
+		{"Walk visits each root", Walk, 2},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var calls int
+			handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+				calls++
+				return nil, nil
+			})
+			require.NoError(t, tt.walk(context.Background(), handler, root, root))
+			assert.Equal(t, tt.wantCalls, calls)
+		})
+	}
+}
+
+func TestWalkCountsDuplicateReferences(t *testing.T) {
+	root := newDesc("root")
+	dup := newDesc("dup")
+	children := make([]ocispec.Descriptor, maxReferences+1)
+	for i := range children {
+		children[i] = dup
+	}
+	fanout := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if desc.Digest == root.Digest {
+			return children, nil
+		}
+		return nil, nil
+	})
+	require.ErrorIs(t, Walk(context.Background(), fanout, root), errdefs.ErrResourceExhausted)
+
+	a, b := newDesc("a"), newDesc("b")
+	cycle := map[digest.Digest][]ocispec.Descriptor{a.Digest: {b}, b.Digest: {a}}
+	cyclic := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		return cycle[desc.Digest], nil
+	})
+	require.ErrorIs(t, Walk(context.Background(), cyclic, a), errdefs.ErrResourceExhausted)
+}
+
+func TestWalkNotEmpty(t *testing.T) {
+	root := newDesc("root")
+	child := newDesc("child")
+
+	t.Run("returns nil when a handler produces children", func(t *testing.T) {
+		graph := map[digest.Digest][]ocispec.Descriptor{root.Digest: {child}}
+		handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+			return graph[desc.Digest], nil
+		})
+		require.NoError(t, WalkNotEmpty(context.Background(), handler, root))
+	})
+
+	t.Run("returns ErrEmptyWalk when everything is filtered out", func(t *testing.T) {
+		handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+			return nil, nil
+		})
+		require.ErrorIs(t, WalkNotEmpty(context.Background(), handler, root), ErrEmptyWalk)
+	})
+
+	t.Run("propagates the reference-cap error", func(t *testing.T) {
+		children := make([]ocispec.Descriptor, maxReferences+1)
+		for i := range children {
+			children[i] = newDesc(fmt.Sprintf("node-%d", i))
+		}
+		handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+			if desc.Digest == root.Digest {
+				return children, nil
+			}
+			return nil, nil
+		})
+		require.ErrorIs(t, WalkNotEmpty(context.Background(), handler, root), errdefs.ErrResourceExhausted)
+	})
+}
+
+func TestDispatchTreeTraversal(t *testing.T) {
+	root := newDesc("root")
+	child1 := newDesc("child1")
+	child2 := newDesc("child2")
+	grandchild1 := newDesc("grandchild1")
+
+	graph := map[digest.Digest][]ocispec.Descriptor{
+		root.Digest:   {child1, child2},
+		child1.Digest: {grandchild1},
+	}
+
+	var visited []digest.Digest
+	var mu sync.Mutex
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		mu.Lock()
+		visited = append(visited, desc.Digest)
+		mu.Unlock()
+		return graph[desc.Digest], nil
+	})
+
+	err := Dispatch(context.Background(), handler, nil, root)
+	require.NoError(t, err)
+
+	set := digestSet(visited)
+	assert.Contains(t, set, root.Digest)
+	assert.Contains(t, set, child1.Digest)
+	assert.Contains(t, set, child2.Digest)
+	assert.Contains(t, set, grandchild1.Digest)
+	assert.Len(t, visited, 4)
+}
+
+func TestDispatchRepeatedReferences(t *testing.T) {
+	root := newDesc("root")
+	childA := newDesc("A")
+	childB := newDesc("B")
+	shared := newDesc("shared")
+
+	graph := map[digest.Digest][]ocispec.Descriptor{
+		root.Digest:   {childA, childB},
+		childA.Digest: {shared},
+		childB.Digest: {shared},
+	}
+
+	tests := []struct {
+		name            string
+		dispatch        func(context.Context, Handler, *semaphore.Weighted, ...ocispec.Descriptor) error
+		wantTotalCalls  int32
+		wantSharedCalls int32
+	}{
+		{
+			name:            "Dispatch visits every reference",
+			dispatch:        Dispatch,
+			wantTotalCalls:  5,
+			wantSharedCalls: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var totalCalls atomic.Int32
+			var sharedCalls atomic.Int32
+			handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+				totalCalls.Add(1)
+				if desc.Digest == shared.Digest {
+					sharedCalls.Add(1)
+				}
+				return graph[desc.Digest], nil
+			})
+
+			err := tt.dispatch(context.Background(), handler, nil, root)
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantTotalCalls, totalCalls.Load())
+			assert.Equal(t, tt.wantSharedCalls, sharedCalls.Load())
+		})
+	}
+}
+
+func TestDispatchErrSkipDesc(t *testing.T) {
+	root := newDesc("root")
+	child1 := newDesc("child1")
+	child2 := newDesc("child2")
+	grandchild := newDesc("grandchild")
+
+	graph := map[digest.Digest][]ocispec.Descriptor{
+		root.Digest:   {child1, child2},
+		child1.Digest: {grandchild},
+	}
+
+	var visited []digest.Digest
+	var mu sync.Mutex
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		mu.Lock()
+		visited = append(visited, desc.Digest)
+		mu.Unlock()
+		if desc.Digest == child1.Digest {
+			return graph[desc.Digest], ErrSkipDesc
+		}
+		return graph[desc.Digest], nil
+	})
+
+	err := Dispatch(context.Background(), handler, nil, root)
+	require.NoError(t, err)
+
+	set := digestSet(visited)
+	assert.Contains(t, set, root.Digest)
+	assert.Contains(t, set, child1.Digest)
+	assert.Contains(t, set, child2.Digest)
+	assert.NotContains(t, set, grandchild.Digest, "grandchild must not be visited when its parent returns ErrSkipDesc")
+}
+
+func TestDispatchConcurrencyLimit(t *testing.T) {
+	tests := []struct {
+		name           string
+		limiter        *semaphore.Weighted
+		maxConcurrency int32
+		numChildren    int
+	}{
+		{
+			name:           "With explicit limiter",
+			limiter:        semaphore.NewWeighted(3),
+			maxConcurrency: 3,
+			numChildren:    3 * 5,
+		},
+		{
+			name:           "With default limiter",
+			limiter:        nil,
+			maxConcurrency: defaultMaxConcurrency,
+			numChildren:    defaultMaxConcurrency * 5,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			root := newDesc("root")
+			children := make([]ocispec.Descriptor, tt.numChildren)
+			for i := range tt.numChildren {
+				children[i] = newDesc(fmt.Sprintf("child-%d", i))
+			}
+
+			graph := map[digest.Digest][]ocispec.Descriptor{
+				root.Digest: children,
+			}
+
+			var concurrent, peak atomic.Int32
+			var exceeded atomic.Bool
+
+			handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+				n := concurrent.Add(1)
+				if n > tt.maxConcurrency {
+					exceeded.Store(true)
+				}
+				for {
+					old := peak.Load()
+					if n <= old || peak.CompareAndSwap(old, n) {
+						break
+					}
+				}
+				defer concurrent.Add(-1)
+
+				// Keep handlers running long enough to overlap.
+				time.Sleep(10 * time.Millisecond)
+				return graph[desc.Digest], nil
+			})
+
+			err := Dispatch(context.Background(), handler, tt.limiter, root)
+			require.NoError(t, err)
+
+			assert.False(t, exceeded.Load(), "concurrent handlers must not exceed capacity")
+			assert.Equal(t, tt.maxConcurrency, peak.Load(),
+				"concurrent handlers must reach the capacity the caller asked for")
+		})
+	}
+}
+
+func TestDispatchSharedLimiter(t *testing.T) {
+	const (
+		weight  = 4
+		callers = 4
+		perTree = 12
+	)
+
+	limiter := semaphore.NewWeighted(weight)
+
+	var concurrent atomic.Int32
+	var exceeded atomic.Bool
+	visited := make([]atomic.Int32, callers)
+
+	var eg errgroup.Group
+	for c := range callers {
+		eg.Go(func() error {
+			root := newDesc(fmt.Sprintf("root-%d", c))
+			children := make([]ocispec.Descriptor, perTree)
+			for i := range perTree {
+				children[i] = newDesc(fmt.Sprintf("caller-%d-child-%d", c, i))
+			}
+			graph := map[digest.Digest][]ocispec.Descriptor{root.Digest: children}
+
+			handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+				if concurrent.Add(1) > weight {
+					exceeded.Store(true)
+				}
+				defer concurrent.Add(-1)
+				visited[c].Add(1)
+
+				time.Sleep(2 * time.Millisecond)
+				return graph[desc.Digest], nil
+			})
+
+			return Dispatch(context.Background(), handler, limiter, root)
+		})
+	}
+
+	done := make(chan error, 1)
+	go func() { done <- eg.Wait() }()
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(30 * time.Second):
+		t.Fatal("shared-limiter traversals did not complete; a caller is starved or deadlocked")
+	}
+
+	assert.False(t, exceeded.Load(), "combined concurrent handlers must not exceed the shared limiter weight")
+	for c := range callers {
+		assert.Equal(t, int32(perTree+1), visited[c].Load(),
+			"caller %d must visit every descriptor in its own graph", c)
+	}
+}
+
+func TestDispatchReferenceLimit(t *testing.T) {
+	// Each root is below the cap on its own; their combined references exceed it.
+	rootA := newDesc("root-a")
+	rootB := newDesc("root-b")
+
+	childrenA := make([]ocispec.Descriptor, maxReferences/2)
+	childrenB := make([]ocispec.Descriptor, maxReferences/2)
+	for i := range childrenA {
+		childrenA[i] = newDesc(fmt.Sprintf("a-%d", i))
+		childrenB[i] = newDesc(fmt.Sprintf("b-%d", i))
+	}
+
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		switch desc.Digest {
+		case rootA.Digest:
+			return childrenA, nil
+		case rootB.Digest:
+			return childrenB, nil
+		}
+		return nil, nil
+	})
+
+	require.ErrorIs(t, Dispatch(context.Background(), handler, nil, rootA, rootB), errdefs.ErrResourceExhausted)
+
+	acceptRoot := newDesc("accept-root")
+	exact := make([]ocispec.Descriptor, maxReferences-1)
+	for i := range exact {
+		exact[i] = newDesc(fmt.Sprintf("accept-%d", i))
+	}
+	acceptHandler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if desc.Digest == acceptRoot.Digest {
+			return exact, nil
+		}
+		return nil, nil
+	})
+	require.NoError(t, Dispatch(context.Background(), acceptHandler, nil, acceptRoot))
+}
+
+func TestDispatchRejectsChildrenBeforeVisiting(t *testing.T) {
+	root := newDesc("root")
+	children := make([]ocispec.Descriptor, maxReferences)
+	for i := range children {
+		children[i] = newDesc(fmt.Sprintf("node-%d", i))
+	}
+	var childVisits atomic.Int32
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if desc.Digest == root.Digest {
+			return children, nil
+		}
+		childVisits.Add(1)
+		return nil, nil
+	})
+
+	require.ErrorIs(t, Dispatch(context.Background(), handler, nil, root), errdefs.ErrResourceExhausted)
+	assert.Equal(t, int32(0), childVisits.Load(), "children of an over-limit slice must not be visited")
+}
+
+func TestDispatchDuplicateRoots(t *testing.T) {
+	root := newDesc("root")
+
+	tests := []struct {
+		name      string
+		dispatch  func(context.Context, Handler, *semaphore.Weighted, ...ocispec.Descriptor) error
+		wantCalls int32
+	}{
+		{"Dispatch visits each root", Dispatch, 2},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var calls atomic.Int32
+			handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+				calls.Add(1)
+				return nil, nil
+			})
+			require.NoError(t, tt.dispatch(context.Background(), handler, nil, root, root))
+			assert.Equal(t, tt.wantCalls, calls.Load())
+		})
+	}
+}
+
+func TestDispatchCountsDuplicateReferences(t *testing.T) {
+	root := newDesc("root")
+	dup := newDesc("dup")
+	children := make([]ocispec.Descriptor, maxReferences+1)
+	for i := range children {
+		children[i] = dup
+	}
+	fanout := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if desc.Digest == root.Digest {
+			return children, nil
+		}
+		return nil, nil
+	})
+	require.ErrorIs(t, Dispatch(context.Background(), fanout, nil, root), errdefs.ErrResourceExhausted)
+
+	a, b := newDesc("a"), newDesc("b")
+	cycle := map[digest.Digest][]ocispec.Descriptor{a.Digest: {b}, b.Digest: {a}}
+	cyclic := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		return cycle[desc.Digest], nil
+	})
+	require.ErrorIs(t, Dispatch(context.Background(), cyclic, nil, a), errdefs.ErrResourceExhausted)
+}
+
+func TestDispatchErrorPropagation(t *testing.T) {
+	root := newDesc("root")
+	child1 := newDesc("child1")
+	child2 := newDesc("child2")
+
+	graph := map[digest.Digest][]ocispec.Descriptor{
+		root.Digest: {child1, child2},
+	}
+
+	errTest := errors.New("handler error")
+
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if desc.Digest == child1.Digest {
+			return nil, errTest
+		}
+		return graph[desc.Digest], nil
+	})
+
+	err := Dispatch(context.Background(), handler, nil, root)
+	require.ErrorIs(t, err, errTest)
+}
+
+func TestDispatchContextCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		return nil, nil
+	})
+
+	err := Dispatch(ctx, handler, nil, newDesc("root"))
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+func TestDispatchContextCancellationMidFlight(t *testing.T) {
+	// Use more descriptors than the limiter admits so cancellation also covers
+	// handlers waiting to start.
+	const n = defaultMaxConcurrency * 4
+	descs := make([]ocispec.Descriptor, n)
+	for i := range n {
+		descs[i] = newDesc(fmt.Sprintf("cancel-%d", i))
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	started := make(chan struct{}, n)
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		started <- struct{}{}
+		<-ctx.Done()
+		return nil, ctx.Err()
+	})
+
+	go func() {
+		for range defaultMaxConcurrency {
+			<-started
+		}
+		cancel()
+	}()
+
+	err := Dispatch(ctx, handler, nil, descs...)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+func TestDispatchCancellationWithNoHandlerError(t *testing.T) {
+	const n = defaultMaxConcurrency * 4
+	descs := make([]ocispec.Descriptor, n)
+	for i := range n {
+		descs[i] = newDesc(fmt.Sprintf("cancel-queued-%d", i))
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	handler := HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if desc.Digest == descs[0].Digest {
+			cancel()
+		}
+		return nil, nil
+	})
+
+	err := Dispatch(ctx, handler, nil, descs...)
+	require.ErrorIs(t, err, context.Canceled)
+}