Add force acquire implementation

Author: aertje

README.md

@@ -15,6 +15,7 @@ The implementation does not interfere with Go's runtime semaphore. It is opt-in
 - **Priority-based scheduling**: Tasks are executed based on their priority. High priority tasks are started before low priority tasks once the maximum concurrency limit is reached.
 - **Configurable maximum concurrency limit**: The number of concurrent tasks is configurable, defaulting to [`GOMAXPROCS`](https://pkg.go.dev/runtime#GOMAXPROCS).
 - **Context cancellation**: Waiting tasks can optionally be cancelled using a context.
+- **Force acquire**: Tasks can bypass the maximum concurrency limit using force acquire. These tasks will execute immediately but still count towards the concurrency limit for regular tasks. This ensures critical tasks are never blocked while maintaining backpressure on non-critical tasks.
 
 ## Installation
 
@@ -64,68 +65,64 @@ Then, for each task to be prioritised:
 
 Note the importance of calling the `Release` method to signal the completion of the task. This is necessary to allow other tasks to be executed by the semaphore.
 
-If the context needs to be taken into account in order to support cancellation, the `AcquireContext` method can be used instead.
+If the context needs to be taken into account in order to support cancellation, the `AcquireContext` method can be used instead. If a highly critical task needs to be executed, the `ForceAcquire` method can be used to bypass the maximum concurrency limit.
 
-## Example use case: Prioritizing `/alive` endpoint
+## Example use case: Prioritizing critical endpoints
 
-In this example, we will create a semaphore that prioritises an `/alive` endpoint over other endpoints. This is useful in scenarios where the `/alive` endpoint is critical and needs to be executed before other endpoints.
+This example demonstrates the key features of the semaphore:
 
-It also demonstrates use of the `AcquireContext` method to support context cancellation. This is useful in scenarios where the client cancels the request, and the server should dispose of the task.
+- Priority-based task execution
+- Force acquire for critical tasks
+- Context cancellation support
 
 ```go
 package main
 
 import (
     "context"
-    "errors"
+    "log"
     "net/http"
     "time"
 
     "github.com/aertje/semaphore/semaphore"
 )
 
 func main() {
-    // Create a new semaphore with a maximum concurrency limit of 10.
-    s := semaphore.NewPrioritized(semaphore.WithMaxConcurrency(10))
-
-    http.HandleFunc("/alive", func(w http.ResponseWriter, r *http.Request) {
-        // Register a task with the semaphore with a higher priority of 1.
-        err := s.AcquireContext(r.Context(), 1)
-        if err != nil {
-            if errors.Is(err, context.Canceled) {
-                http.Error(w, context.Cause(r.Context()).Error(), 499)
-                return
-            }
-
-            http.Error(w, err.Error(), http.StatusInternalServerError)
+    // Create a semaphore with max 2 concurrent tasks
+    s := semaphore.NewPrioritized(semaphore.WithMaxConcurrency(2))
+
+    // Critical health check - uses force acquire to bypass limits
+    http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
+        s.ForceAcquire()
+        defer s.Release()
+        w.Write([]byte("OK"))
+    })
+
+    // High priority endpoint - uses priority 1
+    http.HandleFunc("/users", func(w http.ResponseWriter, r *http.Request) {
+        if err := s.AcquireContext(r.Context(), 1); err != nil {
+            http.Error(w, "Request cancelled", 499)
             return
         }
         defer s.Release()
 
-        w.Write([]byte("I'm alive!"))
+        time.Sleep(100 * time.Millisecond) // Simulate work
+        w.Write([]byte("User data"))
     })
 
+    // Low priority endpoint - uses priority 2
     http.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
-        // Register a task with the semaphore with a lower priority of 2.
-        err := s.AcquireContext(r.Context(), 2)
-        if err != nil {
-            if errors.Is(err, context.Canceled) {
-                http.Error(w, context.Cause(r.Context()).Error(), 499)
-                return
-            }
-
-            http.Error(w, err.Error(), http.StatusInternalServerError)
+        if err := s.AcquireContext(r.Context(), 2); err != nil {
+            http.Error(w, "Request cancelled", 499)
             return
         }
-
         defer s.Release()
 
-        time.Sleep(1 * time.Second)
-
-        w.Write([]byte("Metrics are here!"))
+        time.Sleep(500 * time.Millisecond) // Simulate heavy work
+        w.Write([]byte("Metrics data"))
     })
 
-    http.ListenAndServe(":8080", nil)
+    log.Fatal(http.ListenAndServe(":8080", nil))
 }
 ```
 

queue/queue.go

@@ -1,48 +1,60 @@
 package queue
 
-type Item[T any] struct {
+import "container/heap"
+
+type item[T any] struct {
 	value    T
 	priority int
 	index    int
 }
 
-func NewItem[T any](priority int, value T) *Item[T] {
-	return &Item[T]{value: value, priority: priority}
-}
-
-func (item *Item[T]) Value() T {
-	return item.value
-}
-
-type Q[T any] []*Item[T]
+type Q[T any] []*item[T]
 
-func (pq Q[T]) Len() int {
-	return len(pq)
+// Len implements heap.Interface.
+func (q Q[T]) Len() int {
+	return len(q)
 }
 
-func (pq Q[T]) Less(i, j int) bool {
-	return pq[i].priority < pq[j].priority
+// Less implements heap.Interface.
+func (q Q[T]) Less(i, j int) bool {
+	return q[i].priority < q[j].priority
 }
 
-func (pq Q[T]) Swap(i, j int) {
-	pq[i], pq[j] = pq[j], pq[i]
-	pq[i].index = i
-	pq[j].index = j
+// Swap implements heap.Interface, do not use this method directly.
+func (q Q[T]) Swap(i, j int) {
+	q[i], q[j] = q[j], q[i]
+	q[i].index = i
+	q[j].index = j
 }
 
-func (pq *Q[T]) Push(x any) {
-	n := len(*pq)
-	item := x.(*Item[T])
+// Push implements heap.Interface, do not use this method directly.
+func (q *Q[T]) Push(x any) {
+	n := len(*q)
+	item := x.(*item[T])
 	item.index = n
-	*pq = append(*pq, item)
+	*q = append(*q, item)
 }
 
-func (pq *Q[T]) Pop() any {
-	old := *pq
+// Pop implements heap.Interface, do not use this method directly.
+func (q *Q[T]) Pop() any {
+	old := *q
 	n := len(old)
 	item := old[n-1]
 	old[n-1] = nil  // don't stop the GC from reclaiming the item eventually
 	item.index = -1 // for safety
-	*pq = old[0 : n-1]
+	*q = old[0 : n-1]
 	return item
 }
+
+func (q *Q[T]) PushItem(priority int, value T) {
+	item := &item[T]{value: value, priority: priority}
+	heap.Push(q, item)
+}
+
+func (q *Q[T]) PopItem() T {
+	if q.Len() == 0 {
+		var zero T
+		return zero
+	}
+	return heap.Pop(q).(*item[T]).value
+}

queue/queue_test.go

@@ -12,18 +12,18 @@ func TestOrder(t *testing.T) {
 	q := new(queue.Q[string])
 
 	heap.Init(q)
-	heap.Push(q, queue.NewItem(2, "b"))
-	heap.Push(q, queue.NewItem(3, "c"))
-	heap.Push(q, queue.NewItem(1, "a"))
-	heap.Push(q, queue.NewItem(2, "b"))
-	heap.Push(q, queue.NewItem(1, "a"))
-	heap.Push(q, queue.NewItem(3, "c"))
+
+	q.PushItem(2, "b")
+	q.PushItem(3, "c")
+	q.PushItem(1, "a")
+	q.PushItem(2, "b")
+	q.PushItem(1, "a")
+	q.PushItem(3, "c")
 
 	assert.Equal(t, 6, q.Len())
 
-	for _, want := range []string{"a", "a", "b", "b", "c", "c"} {
-		item := heap.Pop(q).(*queue.Item[string])
-		assert.Equal(t, want, item.Value())
+	for _, want := range []string{"a", "a", "b", "b", "c", "c", ""} {
+		assert.Equal(t, want, q.PopItem())
 	}
 
 	assert.Equal(t, 0, q.Len())

semaphore/semaphore.go

@@ -58,7 +58,7 @@ func (s *Prioritized) assessEntries() {
 		if s.entries.Len() == 0 {
 			return
 		}
-		entry := heap.Pop(s.entries).(*queue.Item[entry]).Value()
+		entry := s.entries.PopItem()
 
 		select {
 		case <-entry.cancelChan:
@@ -71,7 +71,7 @@ func (s *Prioritized) assessEntries() {
 	}
 }
 
-func (s *Prioritized) AcquireContext(ctx context.Context, priority int) error {
+func (s *Prioritized) acquireInternal(ctx context.Context, priority int, force bool) error {
 	waitChan := make(chan struct{})
 	cancelChan := make(chan struct{})
 
@@ -81,7 +81,7 @@ func (s *Prioritized) AcquireContext(ctx context.Context, priority int) error {
 	}
 
 	s.lock.Lock()
-	heap.Push(s.entries, queue.NewItem(priority, entry))
+	s.entries.PushItem(priority, entry)
 	s.lock.Unlock()
 
 	go func() {
@@ -97,8 +97,18 @@ func (s *Prioritized) AcquireContext(ctx context.Context, priority int) error {
 	}
 }
 
+func (s *Prioritized) AcquireContext(ctx context.Context, priority int) error {
+	return s.acquireInternal(ctx, priority, false)
+}
+
 func (s *Prioritized) Acquire(priority int) {
-	s.AcquireContext(context.Background(), priority)
+	_ = s.acquireInternal(context.Background(), priority, false)
+}
+
+func (s *Prioritized) ForceAcquire() {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.concurrency++
 }
 
 func (s *Prioritized) Release() {

semaphore/semaphore_test.go

@@ -56,10 +56,11 @@ func testOrderForConcurrency(maxConcurrency int, totalTasks int) []int {
 
 	// Saturate the scheduler otherwise subsequent tasks will be executed
 	// immediately in undefined order.
-	for i := 0; i < maxConcurrency; i++ {
+	for range maxConcurrency {
 		go func() {
 			s.Acquire(0)
 			defer s.Release()
+
 			time.Sleep(10 * time.Millisecond)
 		}()
 	}
@@ -72,7 +73,7 @@ func testOrderForConcurrency(maxConcurrency int, totalTasks int) []int {
 	var wg sync.WaitGroup
 
 	for i := totalTasks / maxConcurrency; i > 0; i-- {
-		for j := 0; j < maxConcurrency; j++ {
+		for range maxConcurrency {
 			priority := i
 			wg.Add(1)
 			go func() {
@@ -102,8 +103,9 @@ func TestCancel(t *testing.T) {
 	// immediately without waiting.
 	go func() {
 		s.Acquire(0)
+		defer s.Release()
+
 		time.Sleep(10 * time.Millisecond)
-		s.Release()
 	}()
 
 	// Give the scheduler some time to start the goroutine.
@@ -118,3 +120,64 @@ func TestCancel(t *testing.T) {
 	require.Error(t, err)
 	assert.Equal(t, context.DeadlineExceeded, err)
 }
+
+func TestForceAcquire(t *testing.T) {
+	s := semaphore.NewPrioritized(semaphore.WithMaxConcurrency(1))
+
+	// Saturate the scheduler otherwise subsequent tasks will be executed
+	// immediately in undefined order.
+	go func() {
+		s.Acquire(0)
+		defer s.Release()
+
+		time.Sleep(10 * time.Millisecond)
+	}()
+
+	// Give the scheduler some time to start the goroutine.
+	time.Sleep(1 * time.Millisecond)
+
+	results := make([]bool, 0)
+	var lock sync.Mutex
+	var wg sync.WaitGroup
+
+	for range 5 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			s.Acquire(0)
+			defer s.Release()
+
+			time.Sleep(10 * time.Millisecond)
+			lock.Lock()
+			defer lock.Unlock()
+			results = append(results, false)
+		}()
+	}
+	// Give the scheduler some time to start the goroutine.
+	time.Sleep(1 * time.Millisecond)
+
+	// These tasks start later, but they should finish earlier as they're
+	// prioritized by the force acquire - while the normal prioritized tasks are
+	// waiting for the concurrency to be available.
+	for range 5 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			s.ForceAcquire()
+			defer s.Release()
+
+			lock.Lock()
+			time.Sleep(10 * time.Millisecond)
+			defer lock.Unlock()
+			results = append(results, true)
+		}()
+	}
+
+	wg.Wait()
+
+	expected := []bool{true, true, true, true, true, false, false, false, false, false}
+	require.Len(t, results, len(expected))
+	assert.Equal(t, expected, results)
+}