Merge pull request #5 from corhere/version-two

Locker v2

Author: corhere

.github/workflows/test.yml

@@ -2,24 +2,27 @@ name: Run go test
 
 on:
   push:
-    branches: [ $default-branch ]
+    branches: [ 'main' ]
   pull_request:
-    branch: [ $default-branch ]
 
 jobs:
   test:
     name: Test
+    permissions:
+      contents: read
     strategy:
       matrix:
-        go-version: [1.14.x, 1.15.x]
+        go-version:
+          - stable
+          - oldstable
+          - 1.19.x
         os: [ubuntu-latest]
     runs-on: ${{ matrix.os }}
     steps:
+      - uses: actions/checkout@v4
       - name: Install Go
-        uses: actions/setup-go@v2
+        uses: actions/setup-go@v5
         with:
           go-version: ${{ matrix.go-version }}
-      - name: Checkout code
-        uses: actions/checkout@v2
       - name: Test
         run: go test -v ./...

README.md

@@ -5,8 +5,7 @@ locker provides a mechanism for creating finer-grained locking to help
 free up more global locks to handle other tasks.
 
 The implementation looks close to a sync.Mutex, however, the user must provide a
-reference to use to refer to the underlying lock when locking and unlocking,
-and unlock may generate an error.
+reference to use to refer to the underlying lock when locking and unlocking.
 
 If a lock with a given name does not exist when `Lock` is called, one is
 created.
@@ -23,11 +22,11 @@ import (
 	"sync"
 	"time"
 
-	"github.com/moby/locker"
+	"github.com/moby/locker/v2"
 )
 
 type important struct {
-	locks *locker.Locker
+	locks locker.MutexMap[string]
 	data  map[string]interface{}
 	mu    sync.Mutex
 }

go.mod

@@ -1,3 +1,3 @@
-module github.com/moby/locker
+module github.com/moby/locker/v2
 
-go 1.13
+go 1.19

locker.go

@@ -0,0 +1,89 @@
+/*
+Package locker provides a mechanism for creating finer-grained locking to help
+free up more global locks to handle other tasks.
+*/
+package locker
+
+import (
+	"sync"
+	"sync/atomic"
+)
+
+// MutexMap is a more convenient map[T]sync.Mutex. It automatically makes and
+// deletes mutexes as needed. Unlocked mutexes consume no memory.
+//
+// The zero value is a valid MutexMap.
+type MutexMap[T comparable] struct {
+	mu    sync.Mutex
+	locks map[T]*lockCtr
+}
+
+// lockCtr is used by Locker to represent a lock with a given key.
+type lockCtr struct {
+	sync.Mutex
+	waiters atomic.Int32 // Number of callers waiting to acquire the lock
+}
+
+var lockCtrPool = sync.Pool{New: func() any { return new(lockCtr) }}
+
+// Lock locks the mutex identified by key.
+func (l *MutexMap[T]) Lock(key T) {
+	l.mu.Lock()
+	if l.locks == nil {
+		l.locks = make(map[T]*lockCtr)
+	}
+
+	nameLock, exists := l.locks[key]
+	if !exists {
+		nameLock = lockCtrPool.Get().(*lockCtr)
+		l.locks[key] = nameLock
+	}
+
+	// Increment the nameLock waiters while inside the main mutex.
+	// This makes sure that the lock isn't deleted if `Lock` and `Unlock` are called concurrently.
+	nameLock.waiters.Add(1)
+	l.mu.Unlock()
+
+	// Lock the nameLock outside the main mutex so we don't block other operations.
+	// Once locked then we can decrement the number of waiters for this lock.
+	nameLock.Lock()
+	nameLock.waiters.Add(-1)
+}
+
+// Unlock unlocks the mutex identified by key.
+//
+// It is a run-time error if the mutex is not locked on entry to Unlock.
+func (l *MutexMap[T]) Unlock(key T) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	nameLock, exists := l.locks[key]
+	if !exists {
+		// Generate an un-recover()-able error without reaching into runtime internals.
+		(&sync.Mutex{}).Unlock()
+	}
+
+	if nameLock.waiters.Load() <= 0 {
+		delete(l.locks, key)
+		defer lockCtrPool.Put(nameLock)
+	}
+	nameLock.Unlock()
+}
+
+type nameLocker[T comparable] struct {
+	l   *MutexMap[T]
+	key T
+}
+
+// Locker returns a [sync.Locker] interface that implements
+// the [sync.Locker.Lock] and [sync.Locker.Unlock] methods
+// by calling l.Lock(key) and l.Unlock(key).
+func (l *MutexMap[T]) Locker(key T) sync.Locker {
+	return nameLocker[T]{l: l, key: key}
+}
+
+func (n nameLocker[T]) Lock() {
+	n.l.Lock(n.key)
+}
+func (n nameLocker[T]) Unlock() {
+	n.l.Unlock(n.key)
+}

mutexmap.go

@@ -8,27 +8,13 @@ import (
 	"time"
 )
 
-func TestLockCounter(t *testing.T) {
-	l := &lockCtr{}
-	l.inc()
-
-	if l.waiters != 1 {
-		t.Fatal("counter inc failed")
-	}
-
-	l.dec()
-	if l.waiters != 0 {
-		t.Fatal("counter dec failed")
-	}
-}
-
-func TestLockerLock(t *testing.T) {
-	l := New()
+func TestMutexMap_Lock(t *testing.T) {
+	var l MutexMap[string]
 	l.Lock("test")
 	ctr := l.locks["test"]
 
-	if ctr.count() != 0 {
-		t.Fatalf("expected waiters to be 0, got :%d", ctr.waiters)
+	if w := ctr.waiters.Load(); w != 0 {
+		t.Fatalf("expected waiters to be 0, got %d", w)
 	}
 
 	chDone := make(chan struct{})
@@ -40,7 +26,7 @@ func TestLockerLock(t *testing.T) {
 	chWaiting := make(chan struct{})
 	go func() {
 		for range time.Tick(1 * time.Millisecond) {
-			if ctr.count() == 1 {
+			if ctr.waiters.Load() == 1 {
 				close(chWaiting)
 				break
 			}
@@ -59,23 +45,21 @@ func TestLockerLock(t *testing.T) {
 	default:
 	}
 
-	if err := l.Unlock("test"); err != nil {
-		t.Fatal(err)
-	}
+	l.Unlock("test")
 
 	select {
 	case <-chDone:
 	case <-time.After(3 * time.Second):
 		t.Fatalf("lock should have completed")
 	}
 
-	if ctr.count() != 0 {
-		t.Fatalf("expected waiters to be 0, got: %d", ctr.count())
+	if w := ctr.waiters.Load(); w != 0 {
+		t.Fatalf("expected waiters to be 0, got %d", w)
 	}
 }
 
-func TestLockerUnlock(t *testing.T) {
-	l := New()
+func TestMutexMap_Unlock(t *testing.T) {
+	var l MutexMap[string]
 
 	l.Lock("test")
 	l.Unlock("test")
@@ -93,8 +77,8 @@ func TestLockerUnlock(t *testing.T) {
 	}
 }
 
-func TestLockerConcurrency(t *testing.T) {
-	l := New()
+func TestMutexMap_Concurrency(t *testing.T) {
+	var l MutexMap[string]
 
 	var wg sync.WaitGroup
 	for i := 0; i <= 10000; i++ {
@@ -125,16 +109,19 @@ func TestLockerConcurrency(t *testing.T) {
 	}
 }
 
-func BenchmarkLocker(b *testing.B) {
-	l := New()
+func BenchmarkMutexMap(b *testing.B) {
+	var l MutexMap[string]
+	b.ReportAllocs()
 	for i := 0; i < b.N; i++ {
 		l.Lock("test")
+		l.Lock(strconv.Itoa(i))
+		l.Unlock(strconv.Itoa(i))
 		l.Unlock("test")
 	}
 }
 
-func BenchmarkLockerParallel(b *testing.B) {
-	l := New()
+func BenchmarkMutexMap_Parallel(b *testing.B) {
+	var l MutexMap[string]
 	b.SetParallelism(128)
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
@@ -144,8 +131,8 @@ func BenchmarkLockerParallel(b *testing.B) {
 	})
 }
 
-func BenchmarkLockerMoreKeys(b *testing.B) {
-	l := New()
+func BenchmarkMutexMap_MoreKeys(b *testing.B) {
+	var l MutexMap[string]
 	var keys []string
 	for i := 0; i < 64; i++ {
 		keys = append(keys, strconv.Itoa(i))