go/keyedmutex/keyedmutex_test.go

package keyedmutex


import (
	"sync"
	"testing"
	"time"
)

func TestKeyedMutex(t *testing.T) {
	checkState := func(t *testing.T, m KeyedMutex[string], keys ...string) {
		if len(m.waitList) != len(keys) {
			t.Fatal(m.waitList, keys)
		}

		for _, key := range keys {
			if _, ok := m.waitList[key]; !ok {
				t.Fatal(key)
			}
		}
	}

	m := New[string]()
	checkState(t, m)

	m.Lock("a")
	checkState(t, m, "a")
	m.Lock("b")
	checkState(t, m, "a", "b")
	m.Lock("c")
	checkState(t, m, "a", "b", "c")

	if m.TryLock("a") {
		t.Fatal("a")
	}
	if m.TryLock("b") {
		t.Fatal("b")
	}
	if m.TryLock("c") {
		t.Fatal("c")
	}

	if !m.TryLock("d") {
		t.Fatal("d")
	}

	checkState(t, m, "a", "b", "c", "d")

	if !m.TryLock("e") {
		t.Fatal("e")
	}
	checkState(t, m, "a", "b", "c", "d", "e")

	m.Unlock("c")
	checkState(t, m, "a", "b", "d", "e")
	m.Unlock("a")
	checkState(t, m, "b", "d", "e")
	m.Unlock("e")
	checkState(t, m, "b", "d")

	wg := sync.WaitGroup{}
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			m.Lock("b")
			m.Unlock("b")
		}()
	}

	time.Sleep(100 * time.Millisecond)
	m.Unlock("b")
	wg.Wait()

	checkState(t, m, "d")

	m.Unlock("d")
	checkState(t, m)
}

func TestKeyedMutex_unlockUnlocked(t *testing.T) {
	defer func() {
		if r := recover(); r == nil {
			t.Fatal(r)
		}
	}()

	m := New[string]()
	m.Unlock("aldkfj")
}

func BenchmarkUncontendedMutex(b *testing.B) {
	m := New[string]()
	key := "xyz"

	for i := 0; i < b.N; i++ {
		m.Lock(key)
		m.Unlock(key)
	}
}

func BenchmarkContendedMutex(b *testing.B) {
	m := New[string]()
	key := "xyz"

	m.Lock(key)

	wg := sync.WaitGroup{}
	for i := 0; i < b.N; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			m.Lock(key)
			m.Unlock(key)
		}()
	}

	time.Sleep(time.Second)

	b.ResetTimer()
	m.Unlock(key)
	wg.Wait()
}