feat: add TryKeyedLocker

concurrency/keyed_locker.go

@@ -191,3 +191,60 @@ func (l *RWKeyedLocker[K]) release(entry *rwLockEntry, rawKey K) {
 		entry.timer.Store(timer)
 	}
 }
+
+// TryKeyedLocker is a non-blocking version of KeyedLocker.
+// It allows for trying to acquire a lock without blocking if the lock is already held.
+type TryKeyedLocker[K comparable] struct {
+	mu    sync.Mutex
+	locks map[K]*casMutex
+}
+
+// NewTryKeyedLocker creates a new TryKeyedLocker.
+func NewTryKeyedLocker[K comparable]() *TryKeyedLocker[K] {
+	return &TryKeyedLocker[K]{locks: make(map[K]*casMutex)}
+}
+
+// TryLock tries to acquire a lock for the specified key.
+// It returns true if the lock was acquired, false otherwise.
+func (l *TryKeyedLocker[K]) TryLock(key K) bool {
+	l.mu.Lock()
+
+	lock, ok := l.locks[key]
+	if !ok {
+		lock = &casMutex{}
+		l.locks[key] = lock
+	}
+	l.mu.Unlock()
+
+	return lock.TryLock()
+}
+
+// Unlock releases the lock for the specified key.
+func (l *TryKeyedLocker[K]) Unlock(key K) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	lock, ok := l.locks[key]
+	if ok {
+		lock.Unlock()
+		if lock.lock == 0 {
+			delete(l.locks, key)
+		}
+	}
+}
+
+// casMutex is a simple mutex that uses atomic operations to provide a non-blocking lock.
+type casMutex struct {
+	lock int32
+}
+
+// TryLock tries to acquire the lock without blocking.
+// It returns true if the lock was acquired, false otherwise.
+func (m *casMutex) TryLock() bool {
+	return atomic.CompareAndSwapInt32(&m.lock, 0, 1)
+}
+
+// Unlock releases the lock.
+func (m *casMutex) Unlock() {
+	atomic.StoreInt32(&m.lock, 0)
+}

concurrency/keyed_locker_test.go

@@ -175,3 +175,56 @@ func TestRWKeyedLocker_LockTimeout(t *testing.T) {
 
 	assert.IsNotNil(err)
 }
+
+func TestTryKeyedLocker_SimpleLockUnlock(t *testing.T) {
+	t.Parallel()
+	assert := internal.NewAssert(t, "TestTryKeyedLocker_SimpleLockUnlock")
+
+	locker := NewTryKeyedLocker[string]()
+
+	ok := locker.TryLock("key1")
+	assert.Equal(true, ok)
+
+	ok = locker.TryLock("key1")
+	assert.Equal(false, ok)
+
+	locker.Unlock("key1")
+
+	ok = locker.TryLock("key1")
+	assert.Equal(true, ok)
+
+	locker.Unlock("key1")
+}
+
+func TestTryKeyedLocker_ParallelTry(t *testing.T) {
+	t.Parallel()
+	assert := internal.NewAssert(t, "TestTryKeyedLocker_ParallelTry")
+
+	locker := NewTryKeyedLocker[string]()
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var count int
+
+	for i := 0; i < 5; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			ok := locker.TryLock("key" + strconv.Itoa(i))
+			mu.Lock()
+			if ok {
+				count++
+			}
+			mu.Unlock()
+			time.Sleep(10 * time.Millisecond)
+			if ok {
+				locker.Unlock("key" + strconv.Itoa(i))
+			}
+		}(i)
+	}
+
+	wg.Wait()
+
+	assert.Equal(5, count)
+	assert.Equal(0, len(locker.locks))
+}