Panic on flowcontrol deadlock

flowcontrol/fixed.go

@@ -3,35 +3,38 @@ package flowcontrol
 import (
 	"context"
 	"fmt"
-	"math"
 
 	"golang.org/x/sync/semaphore"
 )
 
 // Returns a FlowLimiter that enforces a fixed limit on the total size of
 // outstanding messages.
 func NewFixedLimiter(size int64) FlowLimiter {
-	return (*fixedLimiter)(semaphore.NewWeighted(size))
+	return &fixedLimiter{
+		size: size,
+		sem:  semaphore.NewWeighted(size),
+	}
 }
 
-type fixedLimiter semaphore.Weighted
+type fixedLimiter struct {
+	size int64
+	sem  *semaphore.Weighted
+}
 
 func (fl *fixedLimiter) StartMessage(ctx context.Context, size uint64) (gotResponse func(), err error) {
-	if size > math.MaxInt64 {
-		// semaphore.Weighted expects an int64, so we need to check the bounds.
-		return nil, fmt.Errorf(
-			"StartMessage(): limit %v is too large (max %v)",
-			size, int64(math.MaxInt64),
-		)
+	// HACK:  avoid dead-locking if the size of the message exceeds the maximum
+	//        reservation on the semaphore. We can't return an error because it
+	//        is currently ignored by the caller.
+	if int64(size) > fl.size {
+		const s = "StartMessage(): message size %d is too large (max %d)"
+		panic(fmt.Sprintf(s, size, fl.size))
 	}
-	w := (*semaphore.Weighted)(fl)
-	err = w.Acquire(ctx, int64(size))
-	if err != nil {
-		return nil, err
+
+	if err = fl.sem.Acquire(ctx, int64(size)); err == nil {
+		gotResponse = func() { fl.sem.Release(int64(size)) }
 	}
-	return func() {
-		w.Release(int64(size))
-	}, nil
+
+	return
 }
 
 func (fixedLimiter) Release() {}

flowcontrol/fixed_test.go

@@ -6,32 +6,34 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestFixed(t *testing.T) {
+	t.Parallel()
+
 	ctx := context.Background()
 	lim := NewFixedLimiter(10)
 
 	// Start a couple messages:
 	got4, err := lim.StartMessage(ctx, 4)
-	assert.Nil(t, err, "Limiter returned an error")
+	require.NoError(t, err, "Limiter returned an error")
 	got6, err := lim.StartMessage(ctx, 6)
-	assert.Nil(t, err, "Limiter returned an error")
+	require.NoError(t, err, "Limiter returned an error")
 
 	// We're now exactly at the limit, so if we try again it should block:
 	func() {
 		ctxTimeout, cancel := context.WithTimeout(ctx, 10*time.Millisecond)
 		defer cancel()
 
 		_, err = lim.StartMessage(ctxTimeout, 1)
-		assert.NotNil(t, err, "Limiter didn't return an error")
-		assert.Equal(t, err, ctxTimeout.Err(), "Error wasn't from the context")
+		assert.ErrorIs(t, err, context.DeadlineExceeded, "should return context error")
 	}()
 
 	// Ok, finish one of them and then it should go through again:
 	got4()
 	got1, err := lim.StartMessage(ctx, 1)
-	assert.Nil(t, err, "Limiter returned an error")
+	require.NoError(t, err, "Limiter returned an error")
 
 	// There are 10 - (6 + 1) = 3 bytes remaining. It should therefore block
 	// if we ask for four:
@@ -40,10 +42,16 @@ func TestFixed(t *testing.T) {
 		defer cancel()
 
 		_, err = lim.StartMessage(ctxTimeout, 4)
-		assert.NotNil(t, err, "Limiter didn't return an error")
-		assert.Equal(t, err, ctxTimeout.Err(), "Error wasn't from the context")
+		assert.ErrorIs(t, err, context.DeadlineExceeded, "should return context error")
 	}()
 	got6()
 	got1()
+}
+
+func TestFixeLimiterPanics(t *testing.T) {
+	t.Parallel()
 
+	assert.Panics(t, func() {
+		NewFixedLimiter(1024).StartMessage(context.Background(), 1025)
+	}, "should panic if reservation would cause deadlock")
 }