From f352485485c69efacbf724a195488e2e7840c4ea Mon Sep 17 00:00:00 2001 From: Alan Shaw Date: Mon, 12 Oct 2026 14:27:31 +0100 Subject: [PATCH 1/3] feat: add keylock utility --- keylock/keylock.go | 43 +++++++++++++++++++++++++++++++++++++++++ keylock/keylock_test.go | 33 +++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+) create mode 100644 keylock/keylock.go create mode 100644 keylock/keylock_test.go diff --git a/keylock/keylock.go b/keylock/keylock.go new file mode 100644 index 0000000..11922a4 --- /dev/null +++ b/keylock/keylock.go @@ -0,0 +1,43 @@ +// Package keylock serializes work on a key within one process. +package keylock + +import "sync" + +// Locks holds one mutex per key in use. The zero value is ready to use. +type Locks struct { + mu sync.Mutex + locks map[string]*lock +} + +type lock struct { + sync.Mutex + // waiters counts the holder and everyone waiting; the lock is dropped + // from the map when it reaches zero. + waiters int +} + +// Lock locks key and returns the unlock. +func (l *Locks) Lock(key string) func() { + l.mu.Lock() + if l.locks == nil { + l.locks = map[string]*lock{} + } + k, ok := l.locks[key] + if !ok { + k = &lock{} + l.locks[key] = k + } + k.waiters++ + l.mu.Unlock() + + k.Lock() + return func() { + k.Unlock() + l.mu.Lock() + k.waiters-- + if k.waiters == 0 { + delete(l.locks, key) + } + l.mu.Unlock() + } +} diff --git a/keylock/keylock_test.go b/keylock/keylock_test.go new file mode 100644 index 0000000..c713227 --- /dev/null +++ b/keylock/keylock_test.go @@ -0,0 +1,33 @@ +package keylock + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestLocks(t *testing.T) { + var l Locks + var wg sync.WaitGroup + var mu sync.Mutex + held := map[string]int{} + for i := range 100 { + key := []string{"a", "b"}[i%2] + wg.Add(1) + go func() { + defer wg.Done() + defer l.Lock(key)() + mu.Lock() + held[key]++ + assert.Equal(t, 1, held[key], "one holder per key") + mu.Unlock() + mu.Lock() + held[key]-- + mu.Unlock() + }() + } + wg.Wait() + require.Empty(t, l.locks, "unused locks are dropped") +} From 27ffbfc145bcae0364c239040f804dd3527f7898 Mon Sep 17 00:00:00 2001 From: Alan Shaw Date: Mon, 12 Oct 2026 14:35:14 +0100 Subject: [PATCH 2/3] refactor: address review comments --- keylock/keylock.go | 26 +++++++++------ keylock/keylock_test.go | 71 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 9 deletions(-) diff --git a/keylock/keylock.go b/keylock/keylock.go index 11922a4..54d1837 100644 --- a/keylock/keylock.go +++ b/keylock/keylock.go @@ -3,7 +3,10 @@ package keylock import "sync" -// Locks holds one mutex per key in use. The zero value is ready to use. +// Locks holds one mutex per key in use. The zero value is ready to use. A +// Locks must not be copied after first use, as a sync.Mutex must not: a +// copy would share the map and the per-key records while guarding them +// with a mutex of its own. type Locks struct { mu sync.Mutex locks map[string]*lock @@ -16,7 +19,9 @@ type lock struct { waiters int } -// Lock locks key and returns the unlock. +// Lock locks key and returns the unlock. The unlock releases the key once; +// calling it again does nothing, so a deferred unlock beside an explicit +// one is safe. func (l *Locks) Lock(key string) func() { l.mu.Lock() if l.locks == nil { @@ -31,13 +36,16 @@ func (l *Locks) Lock(key string) func() { l.mu.Unlock() k.Lock() + var once sync.Once return func() { - k.Unlock() - l.mu.Lock() - k.waiters-- - if k.waiters == 0 { - delete(l.locks, key) - } - l.mu.Unlock() + once.Do(func() { + k.Unlock() + l.mu.Lock() + k.waiters-- + if k.waiters == 0 { + delete(l.locks, key) + } + l.mu.Unlock() + }) } } diff --git a/keylock/keylock_test.go b/keylock/keylock_test.go index c713227..38ca424 100644 --- a/keylock/keylock_test.go +++ b/keylock/keylock_test.go @@ -3,6 +3,7 @@ package keylock import ( "sync" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -31,3 +32,73 @@ func TestLocks(t *testing.T) { wg.Wait() require.Empty(t, l.locks, "unused locks are dropped") } + +// TestLocksAreKeyed pins the per-key contract: while a is held, a second +// a waits, b is free at once, and the waiting a proceeds only once a is +// released. +func TestLocksAreKeyed(t *testing.T) { + var l Locks + unlockA := l.Lock("a") + + acquiredA := make(chan struct{}) + go func() { + defer l.Lock("a")() + close(acquiredA) + }() + select { + case <-acquiredA: + t.Fatal("a second a was acquired while the first was held") + case <-time.After(50 * time.Millisecond): + } + + done := make(chan struct{}) + go func() { + defer l.Lock("b")() + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("b waited on a") + } + + unlockA() + select { + case <-acquiredA: + case <-time.After(time.Second): + t.Fatal("the second a never proceeded after the first was released") + } +} + +// TestUnlockIsIdempotent: a second call of an unlock changes nothing, so the +// key stays held by whoever holds it next and the records stay consistent. +func TestUnlockIsIdempotent(t *testing.T) { + var l Locks + unlock := l.Lock("k") + unlock() + unlock2 := l.Lock("k") + unlock() // a stale second call: must not release the new holder + + acquired := make(chan struct{}) + released := make(chan struct{}) + go func() { + unlock := l.Lock("k") + close(acquired) + unlock() + close(released) + }() + select { + case <-acquired: + t.Fatal("a stale unlock released the key") + case <-time.After(50 * time.Millisecond): + } + unlock2() + select { + case <-released: + case <-time.After(time.Second): + t.Fatal("the key was never released") + } + l.mu.Lock() + defer l.mu.Unlock() + require.Empty(t, l.locks, "unused locks are dropped") +} From 94db78b2895044e4c124ad9be978064e24137e71 Mon Sep 17 00:00:00 2001 From: Alan Shaw Date: Mon, 12 Oct 2026 15:03:38 +0100 Subject: [PATCH 3/3] refactor: rename --- keylock/keylock.go => syncutil/keyedmutex.go | 17 +++++++++-------- .../keyedmutex_test.go | 14 +++++++------- 2 files changed, 16 insertions(+), 15 deletions(-) rename keylock/keylock.go => syncutil/keyedmutex.go (61%) rename keylock/keylock_test.go => syncutil/keyedmutex_test.go (89%) diff --git a/keylock/keylock.go b/syncutil/keyedmutex.go similarity index 61% rename from keylock/keylock.go rename to syncutil/keyedmutex.go index 54d1837..d246a53 100644 --- a/keylock/keylock.go +++ b/syncutil/keyedmutex.go @@ -1,13 +1,14 @@ -// Package keylock serializes work on a key within one process. -package keylock +// Package syncutil holds synchronisation helpers beyond the standard +// library's. +package syncutil import "sync" -// Locks holds one mutex per key in use. The zero value is ready to use. A -// Locks must not be copied after first use, as a sync.Mutex must not: a -// copy would share the map and the per-key records while guarding them -// with a mutex of its own. -type Locks struct { +// KeyedMutex serializes work on a key within one process: one mutex per key +// in use. The zero value is ready to use. A KeyedMutex must not be copied +// after first use, as a sync.Mutex must not: a copy would share the map and +// the per-key records while guarding them with a mutex of its own. +type KeyedMutex struct { mu sync.Mutex locks map[string]*lock } @@ -22,7 +23,7 @@ type lock struct { // Lock locks key and returns the unlock. The unlock releases the key once; // calling it again does nothing, so a deferred unlock beside an explicit // one is safe. -func (l *Locks) Lock(key string) func() { +func (l *KeyedMutex) Lock(key string) func() { l.mu.Lock() if l.locks == nil { l.locks = map[string]*lock{} diff --git a/keylock/keylock_test.go b/syncutil/keyedmutex_test.go similarity index 89% rename from keylock/keylock_test.go rename to syncutil/keyedmutex_test.go index 38ca424..9f0e6a4 100644 --- a/keylock/keylock_test.go +++ b/syncutil/keyedmutex_test.go @@ -1,4 +1,4 @@ -package keylock +package syncutil import ( "sync" @@ -9,8 +9,8 @@ import ( "github.com/stretchr/testify/require" ) -func TestLocks(t *testing.T) { - var l Locks +func TestKeyedMutex(t *testing.T) { + var l KeyedMutex var wg sync.WaitGroup var mu sync.Mutex held := map[string]int{} @@ -33,11 +33,11 @@ func TestLocks(t *testing.T) { require.Empty(t, l.locks, "unused locks are dropped") } -// TestLocksAreKeyed pins the per-key contract: while a is held, a second +// TestKeyedMutexIsKeyed pins the per-key contract: while a is held, a second // a waits, b is free at once, and the waiting a proceeds only once a is // released. -func TestLocksAreKeyed(t *testing.T) { - var l Locks +func TestKeyedMutexIsKeyed(t *testing.T) { + var l KeyedMutex unlockA := l.Lock("a") acquiredA := make(chan struct{}) @@ -73,7 +73,7 @@ func TestLocksAreKeyed(t *testing.T) { // TestUnlockIsIdempotent: a second call of an unlock changes nothing, so the // key stays held by whoever holds it next and the records stay consistent. func TestUnlockIsIdempotent(t *testing.T) { - var l Locks + var l KeyedMutex unlock := l.Lock("k") unlock() unlock2 := l.Lock("k")