fix: fix concurrency bugs causing panics and race conditions in internal/syncutil (#921)

1. Remove the use of `atomic.Value`
2. Cancel all go-routines before releasing semaphore permits
3. Wait for schedule go-routines to release the semaphore permits before
returning
4. Add unit tests

Fixes #916, fixes #908
Signed-off-by: Lixia (Sylvia) Lei <[email protected]>

Author: Wwwsylvia

internal/syncutil/limit.go

@@ -17,7 +17,6 @@ package syncutil
 
 import (
 	"context"
-	"sync/atomic"
 
 	"golang.org/x/sync/errgroup"
 	"golang.org/x/sync/semaphore"
@@ -68,27 +67,41 @@ type GoFunc[T any] func(ctx context.Context, region *LimitedRegion, t T) error
 
 // Go concurrently invokes fn on items.
 func Go[T any](ctx context.Context, limiter *semaphore.Weighted, fn GoFunc[T], items ...T) error {
+	ctx, cancel := context.WithCancelCause(ctx)
+	defer cancel(nil)
+
 	eg, egCtx := errgroup.WithContext(ctx)
-	var egErr atomic.Value
 	for _, item := range items {
 		region := LimitRegion(egCtx, limiter)
 		if err := region.Start(); err != nil {
-			if egErr, ok := egErr.Load().(error); ok && egErr != nil {
-				return egErr
-			}
-			return err
+			cancel(err)
+			// break loop instead of returning to allow previously scheduled
+			// goroutines to finish their deferred region.End() calls
+			break
 		}
-		eg.Go(func(t T) func() error {
+
+		eg.Go(func(t T, lr *LimitedRegion) func() error {
 			return func() error {
-				defer region.End()
-				err := fn(egCtx, region, t)
-				if err != nil {
-					egErr.CompareAndSwap(nil, err)
+				defer lr.End()
+
+				select {
+				case <-egCtx.Done():
+					// skip the task if the context is already cancelled
+					return nil
+				default:
+				}
+
+				if err := fn(egCtx, lr, t); err != nil {
+					cancel(err)
 					return err
 				}
 				return nil
 			}
-		}(item))
+		}(item, region))
+	}
+
+	if err := eg.Wait(); err != nil {
+		cancel(err)
 	}
-	return eg.Wait()
+	return context.Cause(ctx)
 }

internal/syncutil/limit_test.go

@@ -0,0 +1,88 @@
+/*
+Copyright The ORAS 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 syncutil
+
+import (
+	"context"
+	"errors"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"golang.org/x/sync/semaphore"
+)
+
+func TestLimitedRegion_Success(t *testing.T) {
+	limiter := semaphore.NewWeighted(2)
+	ctx := context.Background()
+	var counter int32
+
+	err := Go(ctx, limiter, func(ctx context.Context, region *LimitedRegion, i int) error {
+		// just sleeps a little and increments counter to simulate task
+		time.Sleep(10 * time.Millisecond)
+		atomic.AddInt32(&counter, 1)
+		return nil
+	}, 1, 2, 3, 4, 5)
+
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	// after everything finishes, we expect counter to be 5
+	if want := 5; atomic.LoadInt32(&counter) != int32(want) {
+		t.Errorf("expected counter == %v, got %v", want, counter)
+	}
+
+	// when all work is done the semaphore should have all permits available
+	if !limiter.TryAcquire(2) {
+		t.Error("semaphore permits were not fully released at the end")
+	}
+	limiter.Release(2)
+}
+
+func TestLimitedRegion_Cancellation(t *testing.T) {
+	limiter := semaphore.NewWeighted(2)
+	ctx := context.Background()
+	var counter int32
+
+	errTest := errors.New("test error")
+	err := Go(ctx, limiter, func(ctx context.Context, region *LimitedRegion, i int) error {
+		if i < 0 {
+			// trigger an error on negative values
+			return errTest
+		}
+		// just sleeps a little and increments counter to simulate task
+		time.Sleep(50 * time.Millisecond)
+		atomic.AddInt32(&counter, 1)
+		return nil
+	}, 1, -1, 2, 0, -2)
+
+	// we expect the returned error to be errTest.
+	if !errors.Is(err, errTest) {
+		t.Fatalf("expected error %v; got %v", errTest, err)
+	}
+
+	// after everything finishes, we expect counter to be smaller than 5.
+	if max := 5; atomic.LoadInt32(&counter) >= int32(max) {
+		t.Errorf("expected counter < %v, got %v", max, counter)
+	}
+
+	// when all work is done the semaphore should have all permits available
+	if !limiter.TryAcquire(2) {
+		t.Error("semaphore permit was not released after error cancellation")
+	}
+	limiter.Release(2)
+}

internal/syncutil/limitgroup.go

@@ -21,7 +21,7 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-// A LimitedGroup is a collection of goroutines working on subtasks that are part of
+// LimitedGroup is a collection of goroutines working on subtasks that are part of
 // the same overall task.
 type LimitedGroup struct {
 	grp *errgroup.Group

internal/syncutil/limitgroup_test.go

@@ -0,0 +1,120 @@
+/*
+Copyright The ORAS 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 syncutil
+
+import (
+	"context"
+	"errors"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestLimitedGroup_Success(t *testing.T) {
+	ctx := context.Background()
+	numTasks := 5
+
+	// Create a limited group with a concurrency limit of 2.
+	lg, _ := LimitGroup(ctx, 2)
+	var counter int32
+
+	for range numTasks {
+		lg.Go(func() error {
+			// simulate some work.
+			time.Sleep(10 * time.Millisecond)
+			atomic.AddInt32(&counter, 1)
+			return nil
+		})
+	}
+
+	if err := lg.Wait(); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if got := atomic.LoadInt32(&counter); got != int32(numTasks) {
+		t.Errorf("expected counter %d, got %d", numTasks, got)
+	}
+}
+
+func TestLimitedGroup_Error(t *testing.T) {
+	ctx := context.Background()
+	lg, _ := LimitGroup(ctx, 2)
+	errTest := errors.New("test error")
+	var executed int32
+
+	lg.Go(func() error {
+		// delay a bit so that other tasks are scheduled.
+		time.Sleep(20 * time.Millisecond)
+		atomic.AddInt32(&executed, 1)
+		return errTest
+	})
+
+	// simulates a slower, normal task.
+	lg.Go(func() error {
+		// wait until cancellation is (hopefully) in effect.
+		time.Sleep(50 * time.Millisecond)
+		atomic.AddInt32(&executed, 1)
+		return nil
+	})
+
+	err := lg.Wait()
+	if !errors.Is(err, errTest) {
+		t.Fatalf("expected error %v, got %v", errTest, err)
+	}
+
+	if atomic.LoadInt32(&executed) < 1 {
+		t.Errorf("expected at least one task executed, got %d", executed)
+	}
+}
+
+func TestLimitedGroup_Limit(t *testing.T) {
+	ctx := context.Background()
+	limit := 2
+	lg, _ := LimitGroup(ctx, limit)
+	var concurrent, maxConcurrent int32
+	numTasks := 10
+
+	for range numTasks {
+		lg.Go(func() error {
+			// increment the count of concurrently active tasks.
+			cur := atomic.AddInt32(&concurrent, 1)
+			// update the max concurrent tasks if needed.
+			for {
+				prevMax := atomic.LoadInt32(&maxConcurrent)
+				if cur > prevMax {
+					if atomic.CompareAndSwapInt32(&maxConcurrent, prevMax, cur) {
+						break
+					}
+				} else {
+					break
+				}
+			}
+
+			// simulate a short task.
+			time.Sleep(20 * time.Millisecond)
+			atomic.AddInt32(&concurrent, -1)
+			return nil
+		})
+	}
+
+	if err := lg.Wait(); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if maxConcurrent > int32(limit) {
+		t.Errorf("expected max concurrent tasks <= %d, got %d", limit, maxConcurrent)
+	}
+}