Bound Dispatch concurrency and references

Traverse breadth-first, one level at a time. Acquire the limiter before
starting each handler goroutine. Default to 32 concurrent handlers when
no limiter is provided. Children wait until every handler in the current
level finishes.

Preserve duplicate visits and return an error wrapping
ErrResourceExhausted when the traversal exceeds 10,000 references.
Count children as handlers return them to bound the descriptors retained
for the next level.

Assisted-by: Claude Code
Signed-off-by: Chris Henzie <chrishenzie@gmail.com>

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 a single Dispatch admits, duplicates
+	// included.
+	maxReferences = 10_000
+)
+
 var (
 	// ErrSkipDesc is used to skip processing of a descriptor and
 	// its descendants.
@@ -146,45 +157,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,469 @@
+/*
+   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 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)
+}