feat: add RWKeyedLocker

concurrency/keyed_locker.go

@@ -58,6 +58,7 @@ func (l *KeyedLocker[K]) Do(ctx context.Context, key K, fn func()) error {
 	}
 }
 
+// acquire tries to acquire a lock for the specified key.
 func (l *KeyedLocker[K]) acquire(key K) *lockEntry {
 	lock, _ := l.locks.LoadOrStore(key, &lockEntry{})
 	entry := lock.(*lockEntry)
@@ -70,6 +71,7 @@ func (l *KeyedLocker[K]) acquire(key K) *lockEntry {
 	return entry
 }
 
+// release releases the lock for the specified key.
 func (l *KeyedLocker[K]) release(key K, entry *lockEntry, rawKey K) {
 	if atomic.AddInt32(&entry.ref, -1) == 0 {
 		entry.mu.Lock()
@@ -91,3 +93,101 @@ func (l *KeyedLocker[K]) release(key K, entry *lockEntry, rawKey K) {
 		}
 	}
 }
+
+// RWKeyedLocker is a read-write version of KeyedLocker.
+type RWKeyedLocker[K comparable] struct {
+	locks sync.Map
+	ttl   time.Duration
+}
+
+type rwLockEntry struct {
+	mu    sync.RWMutex
+	ref   int32
+	timer atomic.Pointer[time.Timer]
+}
+
+// NewRWKeyedLocker creates a new RWKeyedLocker with the specified TTL for lock expiration.
+// The TTL is used to automatically release locks that are no longer held.
+func NewRWKeyedLocker[K comparable](ttl time.Duration) *RWKeyedLocker[K] {
+	return &RWKeyedLocker[K]{ttl: ttl}
+}
+
+// RLock acquires a read lock for the specified key and executes the provided function.
+// It returns an error if the context is canceled before the function completes.
+func (l *RWKeyedLocker[K]) RLock(ctx context.Context, key K, fn func()) error {
+	entry := l.acquire(key)
+	defer l.release(entry, key)
+
+	done := make(chan struct{})
+
+	go func() {
+		entry.mu.RLock()
+		defer entry.mu.RUnlock()
+
+		select {
+		case <-ctx.Done():
+		default:
+			fn()
+		}
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-done:
+		return nil
+	}
+}
+
+// Lock acquires a write lock for the specified key and executes the provided function.
+// It returns an error if the context is canceled before the function completes.
+func (l *RWKeyedLocker[K]) Lock(ctx context.Context, key K, fn func()) error {
+	entry := l.acquire(key)
+	defer l.release(entry, key)
+
+	done := make(chan struct{})
+
+	go func() {
+		entry.mu.Lock()
+		defer entry.mu.Unlock()
+
+		select {
+		case <-ctx.Done():
+		default:
+			fn()
+		}
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-done:
+		return nil
+	}
+}
+
+// acquire tries to acquire a read lock for the specified key.
+func (l *RWKeyedLocker[K]) acquire(key K) *rwLockEntry {
+	actual, _ := l.locks.LoadOrStore(key, &rwLockEntry{})
+	entry := actual.(*rwLockEntry)
+	atomic.AddInt32(&entry.ref, 1)
+
+	if t := entry.timer.Swap(nil); t != nil {
+		t.Stop()
+	}
+	return entry
+}
+
+// release releases the lock for the specified key.
+func (l *RWKeyedLocker[K]) release(entry *rwLockEntry, rawKey K) {
+	if atomic.AddInt32(&entry.ref, -1) == 0 {
+		timer := time.AfterFunc(l.ttl, func() {
+			if atomic.LoadInt32(&entry.ref) == 0 {
+				l.locks.Delete(rawKey)
+			}
+		})
+		entry.timer.Store(timer)
+	}
+}

concurrency/keyed_locker_test.go

@@ -104,3 +104,74 @@ func TestKeyedLocker_LockReleaseAfterTTL(t *testing.T) {
 	err = locker.Do(context.Background(), "ttl-key", func() {})
 	assert.IsNil(err)
 }
+
+func TestRWKeyedLocker_LockAndUnlock(t *testing.T) {
+	t.Parallel()
+	assert := internal.NewAssert(t, "TestKeyedLocker_LockReleaseAfterTTL")
+
+	locker := NewRWKeyedLocker[string](500 * time.Millisecond)
+
+	var locked bool
+	err := locker.Lock(context.Background(), "key1", func() {
+		locked = true
+	})
+
+	assert.IsNil(err)
+	assert.Equal(true, locked)
+}
+
+func TestRWKeyedLocker_RLockParallel(t *testing.T) {
+	t.Parallel()
+	assert := internal.NewAssert(t, "TestKeyedLocker_LockReleaseAfterTTL")
+
+	locker := NewRWKeyedLocker[string](1 * time.Second)
+
+	var mu sync.Mutex
+	var count int
+	wg := sync.WaitGroup{}
+
+	for i := 0; i < 5; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			err := locker.RLock(context.Background(), "shared-key", func() {
+				time.Sleep(10 * time.Millisecond)
+				mu.Lock()
+				count++
+				mu.Unlock()
+			})
+
+			assert.IsNil(err)
+		}()
+	}
+
+	wg.Wait()
+
+	assert.Equal(5, count)
+}
+
+func TestRWKeyedLocker_LockTimeout(t *testing.T) {
+	t.Parallel()
+	assert := internal.NewAssert(t, "TestRWKeyedLocker_LockTimeout")
+
+	locker := NewRWKeyedLocker[string](1 * time.Second)
+
+	start := make(chan struct{})
+
+	go func() {
+		locker.Lock(context.Background(), "key-timeout", func() {
+			close(start)
+			time.Sleep(200 * time.Millisecond)
+		})
+	}()
+
+	<-start
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	err := locker.Lock(ctx, "key-timeout", func() {
+		t.Error("should not reach here")
+	})
+
+	assert.IsNotNil(err)
+}