[BREAKING CHANGE] API adjustment to add context.Context

Author: muir

claude_test.go

@@ -3,6 +3,7 @@ package simultaneous_test
 // This generated with Claude 3.7 Sonnet
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -16,17 +17,17 @@ func TestTimeoutRejection(t *testing.T) {
 	limit := simultaneous.New[any](1)
 
 	// Take the only available slot
-	done := limit.Forever()
+	done := limit.Forever(context.Background())
 	defer done.Done()
 
 	// This should time out immediately
-	_, err := limit.Timeout(0)
+	_, err := limit.Timeout(context.Background(), 0)
 	assert.Error(t, err)
 	t.Log("Timeout(0) correctly rejected when limit is full")
 
 	// This should time out after a short wait
 	start := time.Now()
-	_, err = limit.Timeout(50 * time.Millisecond)
+	_, err = limit.Timeout(context.Background(), 50*time.Millisecond)
 	duration := time.Since(start)
 
 	assert.Error(t, err)

limit.go

@@ -5,6 +5,7 @@ need a limit.
 package simultaneous
 
 import (
+	"context"
 	"time"
 
 	"github.com/memsql/errors"
@@ -28,8 +29,8 @@ type Enforced[T any] interface {
 // contract.
 type Limit[T any] struct {
 	queue           chan struct{}
-	stuckCallback   func()
-	unstuckCallback func()
+	stuckCallback   func(context.Context)
+	unstuckCallback func(context.Context)
 	stuckTimeout    time.Duration
 }
 
@@ -49,25 +50,28 @@ func Unlimited[T any]() Enforced[T] {
 }
 
 // Forever waits until there is space in the Limit for another
-// simultaneous runner. It will wait forever. The Done() method
+// simultaneous runner. It will wait until the context is
+// cancelled.  The Done() method
 // must be called to release the space.
 //
 //	defer limit.Forever().Done()
-func (l *Limit[T]) Forever() Limited[T] {
+func (l *Limit[T]) Forever(ctx context.Context) Limited[T] {
 	if l.stuckTimeout == 0 {
 		l.queue <- struct{}{}
 	} else {
 		timer := time.NewTimer(l.stuckTimeout)
 		select {
 		case l.queue <- struct{}{}:
 			timer.Stop()
+		case <-ctx.Done():
+			timer.Stop()
 		case <-timer.C:
 			if l.stuckCallback != nil {
-				l.stuckCallback()
+				l.stuckCallback(ctx)
 			}
 			l.queue <- struct{}{}
 			if l.unstuckCallback != nil {
-				l.unstuckCallback()
+				l.unstuckCallback(ctx)
 			}
 		}
 	}
@@ -82,13 +86,15 @@ var ErrTimeout errors.String = "could not get permission to run before timeout"
 // simultaneous runner. In the case of a timeout, ErrTimeout is returned
 // and the Done method is a no-op. If there is room, the Done method must
 // be invoked to make room for another runner.
-func (l *Limit[T]) Timeout(timeout time.Duration) (Limited[T], error) {
+func (l *Limit[T]) Timeout(ctx context.Context, timeout time.Duration) (Limited[T], error) {
 	if timeout <= 0 {
 		select {
 		case l.queue <- struct{}{}:
 			return limited[T](func() {
 				<-l.queue
 			}), nil
+		case <-ctx.Done():
+			return limited[T](nil), ErrTimeout.Errorf("context cancelled before any simultaneous runner (of %d) became available", cap(l.queue))
 		default:
 			return limited[T](nil), ErrTimeout.Errorf("timeout (%s) expired before any simultaneous runner (of %d) became available", timeout, cap(l.queue))
 		}
@@ -100,6 +106,8 @@ func (l *Limit[T]) Timeout(timeout time.Duration) (Limited[T], error) {
 		return limited[T](func() {
 			<-l.queue
 		}), nil
+	case <-ctx.Done():
+		return limited[T](nil), ErrTimeout.Errorf("context cancelled before any simultaneous runner (of %d) became available", cap(l.queue))
 	case <-timer.C:
 		return limited[T](nil), ErrTimeout.Errorf("timeout (%s) expired before any simultaneous runner (of %d) became available", timeout, cap(l.queue))
 	}
@@ -108,7 +116,7 @@ func (l *Limit[T]) Timeout(timeout time.Duration) (Limited[T], error) {
 // SetForeverMessaging returns a modified Limit that changes the behavior of Forever() so that
 // it will call stuckCallback() (if set) after waiting for stuckTimeout duration. If past that duration,
 // and it will call unstuckCallback() (if set) when it finally gets a limit.
-func (l Limit[T]) SetForeverMessaging(stuckTimeout time.Duration, stuckCallback func(), unstuckCallback func()) *Limit[T] {
+func (l Limit[T]) SetForeverMessaging(stuckTimeout time.Duration, stuckCallback func(context.Context), unstuckCallback func(context.Context)) *Limit[T] {
 	l.stuckTimeout = stuckTimeout
 	l.stuckCallback = stuckCallback
 	l.unstuckCallback = unstuckCallback

limit_test.go

@@ -1,6 +1,7 @@
 package simultaneous_test
 
 import (
+	"context"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -38,12 +39,12 @@ func testLimit(t *testing.T, withStuck bool) {
 
 	if withStuck {
 		limit = limit.SetForeverMessaging(time.Millisecond,
-			func() {
+			func(context.Context) {
 				if stuckCalled.Add(1) == 1 {
 					close(someUnstuck)
 				}
 			},
-			func() {
+			func(context.Context) {
 				unstuckCalled.Add(1)
 			},
 		)
@@ -58,12 +59,12 @@ func testLimit(t *testing.T, withStuck bool) {
 		go func() {
 			defer wg.Done()
 			var done simultaneous.Limited[any]
-			switch i % 3 {
+			switch i % 4 {
 			case 0:
-				done = limit.Forever()
+				done = limit.Forever(context.Background())
 			case 1:
 				var err error
-				done, err = limit.Timeout(0)
+				done, err = limit.Timeout(context.Background(), 0)
 				if err != nil {
 					fail.Add(1)
 					return
@@ -72,7 +73,15 @@ func testLimit(t *testing.T, withStuck bool) {
 				}
 			case 2:
 				var err error
-				done, err = limit.Timeout(time.Second * 2)
+				done, err = limit.Timeout(context.Background(), time.Second*2)
+				if !assert.NoError(t, err) {
+					return
+				}
+			case 3:
+				var err error
+				ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
+				t.Cleanup(cancel)
+				done, err = limit.Timeout(ctx, time.Hour)
 				if !assert.NoError(t, err) {
 					return
 				}