Merge pull request #14 from RaduBerinde/sema-allow-large

fifo: allow large semaphore requests

Author: RaduBerinde

fifo/semaphore.go

@@ -16,7 +16,6 @@ package fifo
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"sync"
 
@@ -35,7 +34,8 @@ type Semaphore struct {
 		sync.Mutex
 
 		capacity int64
-		// outstanding can exceed capacity if the capacity is dynamically decreased.
+		// outstanding can exceed capacity if the capacity is dynamically decreased
+		// or if a single request exceeds the capacity.
 		outstanding int64
 
 		waiters Queue[semaWaiter]
@@ -64,44 +64,44 @@ func NewSemaphore(capacity int64) *Semaphore {
 
 var semaQueuePool = MakeQueueBackingPool[semaWaiter]()
 
-// ErrRequestExceedsCapacity is returned when an Acquire requests more than the
-// current capacity of the semaphore.
-var ErrRequestExceedsCapacity = errors.New("request exceeds semaphore capacity")
-
 // TryAcquire attempts to acquire n units from the semaphore without waiting. On
 // success, returns true and the caller must later Release the units.
 func (s *Semaphore) TryAcquire(n int64) bool {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	if s.numWaitersLocked() == 0 && s.mu.outstanding+n <= s.mu.capacity {
+	if s.numWaitersLocked() == 0 && s.canAcquireLocked(n) {
 		s.mu.outstanding += n
 		return true
 	}
 
 	return false
 }
 
+func (s *Semaphore) canAcquireLocked(n int64) bool {
+	// We allow a request larger than the capacity as long as there are no
+	// outstanding units.
+	return s.mu.outstanding+n <= s.mu.capacity || s.mu.outstanding == 0
+}
+
 // Acquire n units from the semaphore, waiting if necessary.
 //
 // If the context is canceled while we are waiting, returns the context error.
-// If n exceeds the current capacity, returns ErrRequestExceedsCapacity.
+//
+// If n exceeds the current capacity, the request will be allowed when there are
+// no other acquisitions (similar to n being equal to the capacity).
+//
 // On success, the caller must later Release the units.
 func (s *Semaphore) Acquire(ctx context.Context, n int64) error {
 	s.mu.Lock()
 
 	// Fast path.
-	if s.numWaitersLocked() == 0 && s.mu.outstanding+n <= s.mu.capacity {
+	if s.numWaitersLocked() == 0 && s.canAcquireLocked(n) {
 		s.mu.outstanding += n
 		s.mu.Unlock()
 		return nil
 	}
 
-	if n > s.mu.capacity {
-		s.mu.Unlock()
-		return ErrRequestExceedsCapacity
-	}
-
 	c := chanSyncPool.Get().(chan error)
 	defer chanSyncPool.Put(c)
 	w := s.mu.waiters.PushBack(semaWaiter{n: n, c: c})
@@ -177,7 +177,8 @@ type SemaphoreStats struct {
 	// Capacity is the current capacity of the semaphore.
 	Capacity int64
 	// Outstanding is the number of units that have been acquired. Note that this
-	// can exceed Capacity if the capacity was recently decreased.
+	// can exceed Capacity if the capacity was recently decreased or if a single
+	// request exceeded the capacity.
 	Outstanding int64
 	// NumHadToWait is the total number of Acquire calls (since the semaphore was
 	// created) that had to wait because the semaphore was exhausted. Useful for
@@ -219,16 +220,11 @@ func (s *Semaphore) processWaitersLocked() {
 				panic("negative numCanceled")
 			}
 
-		case s.mu.outstanding+w.n <= s.mu.capacity:
+		case s.canAcquireLocked(w.n):
 			// Request can be fulfilled.
 			s.mu.outstanding += w.n
 			w.c <- nil
 
-		case w.n > s.mu.capacity:
-			// Request must be failed. This can happen if the capacity was decreased
-			// while the element was queued.
-			w.c <- ErrRequestExceedsCapacity
-
 		default:
 			// Head of the queue needs to wait some more.
 			return

fifo/semaphore_test.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"math/rand"
 	"runtime"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -30,7 +31,6 @@ func TestSemaphoreAPI(t *testing.T) {
 	s := NewSemaphore(10)
 	require.Equal(t, s.TryAcquire(5), true)
 	require.Equal(t, s.TryAcquire(10), false)
-	require.Equal(t, s.Acquire(context.Background(), 20), ErrRequestExceedsCapacity)
 	require.Equal(t, "capacity: 10, outstanding: 5, num-had-to-wait: 0", s.Stats().String())
 
 	ch := make(chan struct{}, 10)
@@ -58,6 +58,7 @@ func TestSemaphoreAPI(t *testing.T) {
 	s.Release(8)
 	require.Recv(t, ch)
 
+	require.True(t, strings.Contains(s.Stats().String(), "capacity: 10, outstanding: 5"))
 	// Test UpdateCapacity.
 	go func() {
 		if err := s.Acquire(context.Background(), 8); err != nil {
@@ -68,15 +69,26 @@ func TestSemaphoreAPI(t *testing.T) {
 			t.Error(err)
 		}
 		ch <- struct{}{}
-		require.Equal(t, s.Acquire(context.Background(), 5), ErrRequestExceedsCapacity)
-		ch <- struct{}{}
 	}()
 	require.NoRecv(t, ch)
 	s.UpdateCapacity(15)
 	require.Recv(t, ch)
 	require.Recv(t, ch)
-	require.NoRecv(t, ch)
 	s.UpdateCapacity(2)
+	go func() {
+		// Request more than the capacity.
+		if err := s.Acquire(context.Background(), 5); err != nil {
+			t.Error(err)
+		}
+		ch <- struct{}{}
+	}()
+	require.NoRecv(t, ch)
+	s.Release(5)
+	require.NoRecv(t, ch)
+	s.Release(8)
+	require.NoRecv(t, ch)
+	s.Release(1)
+	// Last request should now be allowed, despite being larger than the capacity.
 	require.Recv(t, ch)
 }