diff --git a/core/images/handlers.go b/core/images/handlers.go index 245b9ac43c..7d1ec35cb4 100644 --- a/core/images/handlers.go +++ b/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. diff --git a/core/images/handlers_test.go b/core/images/handlers_test.go new file mode 100644 index 0000000000..acad816d79 --- /dev/null +++ b/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) +}