github.com/dtroyer-salad/og2/v2@v2.0.0-20240412154159-c47231610877/internal/syncutil/limit.go (about) 1 /* 2 Copyright The ORAS Authors. 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 */ 15 16 package syncutil 17 18 import ( 19 "context" 20 "sync/atomic" 21 22 "golang.org/x/sync/errgroup" 23 "golang.org/x/sync/semaphore" 24 ) 25 26 // LimitedRegion provides a way to bound concurrent access to a code block. 27 type LimitedRegion struct { 28 ctx context.Context 29 limiter *semaphore.Weighted 30 ended bool 31 } 32 33 // LimitRegion creates a new LimitedRegion. 34 func LimitRegion(ctx context.Context, limiter *semaphore.Weighted) *LimitedRegion { 35 if limiter == nil { 36 return nil 37 } 38 return &LimitedRegion{ 39 ctx: ctx, 40 limiter: limiter, 41 ended: true, 42 } 43 } 44 45 // Start starts the region with concurrency limit. 46 func (lr *LimitedRegion) Start() error { 47 if lr == nil || !lr.ended { 48 return nil 49 } 50 if err := lr.limiter.Acquire(lr.ctx, 1); err != nil { 51 return err 52 } 53 lr.ended = false 54 return nil 55 } 56 57 // End ends the region with concurrency limit. 58 func (lr *LimitedRegion) End() { 59 if lr == nil || lr.ended { 60 return 61 } 62 lr.limiter.Release(1) 63 lr.ended = true 64 } 65 66 // GoFunc represents a function that can be invoked by Go. 67 type GoFunc[T any] func(ctx context.Context, region *LimitedRegion, t T) error 68 69 // Go concurrently invokes fn on items. 70 func Go[T any](ctx context.Context, limiter *semaphore.Weighted, fn GoFunc[T], items ...T) error { 71 eg, egCtx := errgroup.WithContext(ctx) 72 var egErr atomic.Value 73 for _, item := range items { 74 region := LimitRegion(egCtx, limiter) 75 if err := region.Start(); err != nil { 76 if egErr, ok := egErr.Load().(error); ok && egErr != nil { 77 return egErr 78 } 79 return err 80 } 81 eg.Go(func(t T) func() error { 82 return func() error { 83 defer region.End() 84 err := fn(egCtx, region, t) 85 if err != nil { 86 egErr.CompareAndSwap(nil, err) 87 return err 88 } 89 return nil 90 } 91 }(item)) 92 } 93 return eg.Wait() 94 }