1

pkg/utility/lockkey/lockkey.go

@@ -0,0 +1,40 @@
+package lockkey
+
+import (
+	"fmt"
+	"sync"
+)
+
+type LockByKey struct {
+	m map[string]*sync.RWMutex
+	l sync.Mutex
+}
+
+func NewLockByKey() *LockByKey {
+	return &LockByKey{m: make(map[string]*sync.RWMutex)}
+}
+
+func (lk *LockByKey) Lock(key string) {
+	lk.l.Lock()
+	defer lk.l.Unlock()
+
+	mu, ok := lk.m[key]
+	if !ok {
+		mu = &sync.RWMutex{}
+		lk.m[key] = mu
+	}
+
+	mu.Lock()
+}
+
+func (lk *LockByKey) Unlock(key string) {
+	lk.l.Lock()
+	defer lk.l.Unlock()
+
+	mu, ok := lk.m[key]
+	if !ok {
+		panic(fmt.Sprintf("unlock non-existent key: %v", key))
+	}
+
+	mu.Unlock()
+}

pkg/utility/lockkey/lockkey_test.go

@@ -0,0 +1,34 @@
+package lockkey
+
+import (
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestLock(t *testing.T) {
+	lk := NewLockByKey()
+
+	// Lock and unlock the same key multiple times.
+	for i := 0; i < 5; i++ {
+		go func() {
+			lk.Lock("key")
+			fmt.Println("locked")
+			lk.Unlock("key")
+			fmt.Println("unlocked")
+		}()
+	}
+
+	// Lock and unlock different keys.
+	lk.Lock("key1")
+	fmt.Println("locked key1")
+	lk.Unlock("key1")
+	fmt.Println("unlocked key1")
+
+	lk.Lock("key2")
+	fmt.Println("locked key2")
+	lk.Unlock("key2")
+	fmt.Println("unlocked key2")
+
+	time.Sleep(3 * time.Second)
+}