feat: implement priority lock (#33)

README.md

@@ -17,7 +17,8 @@ Async is a synchronization and asynchronous computation package for Go.
 * **Once** - An object similar to sync.Once having the Do method taking `f func() (T, error)` and returning `(T, error)`.
 * **Value** - An object similar to atomic.Value, but without the consistent type constraint.
 * **WaitGroupContext** - A WaitGroup with the `context.Context` support for graceful unblocking.
-* **ReentrantLock** - A Mutex that allows goroutines to enter into the lock on a resource more than once.
+* **ReentrantLock** - A mutex that allows goroutines to enter into the lock on a resource more than once.
+* **PriorityLock** - A non-reentrant mutex that allows for the specification of lock acquisition priority.
 
 ## Examples
 Can be found in the examples directory/tests.

priority_lock.go

@@ -0,0 +1,80 @@
+package async
+
+import (
+	"fmt"
+	"sync"
+)
+
+const priorityLimit = 1024
+
+// PriorityLock is a non-reentrant mutex that allows specifying a priority
+// level when acquiring the lock. It extends the standard sync.Locker interface
+// with an additional locking method, LockP, which takes a priority level as an
+// argument.
+//
+// The current implementation may cause starvation for lower priority
+// lock requests.
+type PriorityLock struct {
+	sem []chan struct{}
+	max int
+}
+
+var _ sync.Locker = (*PriorityLock)(nil)
+
+// NewPriorityLock instantiates and returns a new PriorityLock, specifying the
+// maximum priority level that can be used in the LockP method. It panics if
+// the maximum priority level is non-positive or exceeds the hard limit.
+func NewPriorityLock(maxPriority int) *PriorityLock {
+	if maxPriority < 1 {
+		panic(fmt.Errorf("nonpositive maximum priority: %d", maxPriority))
+	}
+	if maxPriority > priorityLimit {
+		panic(fmt.Errorf("maximum priority %d exceeds hard limit of %d",
+			maxPriority, priorityLimit))
+	}
+	sem := make([]chan struct{}, maxPriority+1)
+	sem[0] = make(chan struct{}, 1)
+	sem[0] <- struct{}{}
+	for i := 1; i <= maxPriority; i++ {
+		sem[i] = make(chan struct{})
+	}
+	return &PriorityLock{
+		sem: sem,
+		max: maxPriority,
+	}
+}
+
+// Lock will block the calling goroutine until it acquires the lock, using
+// the highest available priority.
+func (pl *PriorityLock) Lock() {
+	pl.LockP(pl.max)
+}
+
+// LockP blocks the calling goroutine until it acquires the lock. Requests with
+// higher priorities acquire the lock first. If the provided priority is
+// outside the valid range, it will be assigned the boundary value.
+func (pl *PriorityLock) LockP(priority int) {
+	switch {
+	case priority < 1:
+		priority = 1
+	case priority > pl.max:
+		priority = pl.max
+	}
+	select {
+	case <-pl.sem[priority]:
+	case <-pl.sem[0]:
+	}
+}
+
+// Unlock releases the previously acquired lock.
+// It will panic if the lock is already unlocked.
+func (pl *PriorityLock) Unlock() {
+	for i := pl.max; i >= 0; i-- {
+		select {
+		case pl.sem[i] <- struct{}{}:
+			return
+		default:
+		}
+	}
+	panic("async: unlock of unlocked PriorityLock")
+}

priority_lock_test.go

@@ -0,0 +1,66 @@
+package async
+
+import (
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/reugn/async/internal/assert"
+)
+
+func TestPriorityLock(t *testing.T) {
+	p := NewPriorityLock(5)
+	var b strings.Builder
+
+	p.Lock() // acquire first to make the result predictable
+	go func() {
+		time.Sleep(time.Millisecond)
+		p.Unlock()
+	}()
+	for i := 0; i < 10; i++ {
+		for j := 5; j > 0; j-- {
+			go func(n int) {
+				p.LockP(n)
+				time.Sleep(time.Microsecond)
+				b.WriteString(strconv.Itoa(n))
+				p.Unlock()
+			}(j)
+		}
+	}
+	time.Sleep(20 * time.Millisecond)
+
+	p.Lock()
+	result := b.String()
+	p.Unlock()
+	var expected strings.Builder
+	for i := 5; i > 0; i-- {
+		expected.WriteString(strings.Repeat(strconv.Itoa(i), 10))
+	}
+	assert.Equal(t, result, expected.String())
+}
+
+func TestPriorityLock_LockRange(t *testing.T) {
+	p := NewPriorityLock(2)
+	var b strings.Builder
+	p.LockP(-1)
+	b.WriteRune('1')
+	p.Unlock()
+	p.LockP(2048)
+	b.WriteRune('1')
+	p.Unlock()
+	assert.Equal(t, b.String(), "11")
+}
+
+func TestPriorityLock_Panic(t *testing.T) {
+	p := NewPriorityLock(2)
+	p.Lock()
+	time.Sleep(time.Nanosecond) // to silence empty critical section warning
+	p.Unlock()
+	assert.PanicMsgContains(t, func() { p.Unlock() }, "unlock of unlocked PriorityLock")
+}
+
+func TestPriorityLock_Validation(t *testing.T) {
+	assert.PanicMsgContains(t, func() { NewPriorityLock(-1) }, "nonpositive maximum priority")
+	assert.PanicMsgContains(t, func() { NewPriorityLock(2048) }, "exceeds hard limit")
+}