cmap_test.go

package cmap

import (
	"sync"
	"sync/atomic"
	"testing"
)

type (
	Key   int // any comparable type
	Value int // any type
	Map   map[Key]Value
)

type ConcurrentMap interface {
	Get(key Key) (Value, bool)
	Set(key Key, value Value)
	Delete(key Key)
}

// concurrent map guarded by a read-write mutex
type RWMMap struct {
	m  Map
	mu sync.RWMutex
}

func NewRWMMap() *RWMMap {
	return &RWMMap{m: make(Map)}
}

func (m *RWMMap) Get(key Key) (Value, bool) {
	m.mu.RLock()
	v, ok := m.m[key]
	m.mu.RUnlock()
	return v, ok
}

func (m *RWMMap) Set(key Key, value Value) {
	m.mu.Lock()
	m.m[key] = value
	m.mu.Unlock()
}

func (m *RWMMap) Delete(key Key) {
	m.mu.Lock()
	delete(m.m, key)
	m.mu.Unlock()
}

// concurrent map guarded by atomic.Value with a mutex for writers
type AVMap struct {
	val atomic.Value
	mu  sync.Mutex
}

func NewAVMap() *AVMap {
	var m AVMap
	m.val.Store(make(Map))
	return &m
}

func (m *AVMap) Get(key Key) (Value, bool) {
	v, ok := m.val.Load().(Map)[key]
	return v, ok
}

func (m *AVMap) Set(key Key, value Value) {
	m.mu.Lock()

	m1 := m.val.Load().(Map)
	m2 := make(Map, len(m1))
	for k, v := range m1 {
		m2[k] = v
	}
	m2[key] = value
	m.val.Store(m2)

	m.mu.Unlock()
}

func (m *AVMap) Delete(key Key) {
	m.mu.Lock()

	m1 := m.val.Load().(Map)
	m2 := make(Map, len(m1))
	for k, v := range m1 {
		if k != key {
			m2[k] = v
		}
	}
	m.val.Store(m2)

	m.mu.Unlock()
}

const (
	readers = 10000
	writers = 1
	max     = 1000
)

func BenchmarkAVMap(b *testing.B) {
	bench(b, NewAVMap())
}

func BenchmarkRWMMap(b *testing.B) {
	bench(b, NewRWMMap())
}

func bench(b *testing.B, m ConcurrentMap) {
	n := b.N
	wg := new(sync.WaitGroup)

	rw := readers + writers
	ready := make(chan struct{}, rw)
	start := make(chan struct{})

	wg.Add(rw)

	for i := 0; i < readers; i++ {
		go func(i int) {
			k := i
			ready <- struct{}{}
			<-start
			for j := 0; j < n; j++ {
				k = (k + 1) % max
				_, _ = m.Get(Key(k))
			}
			wg.Done()
		}(i)
	}

	for i := 0; i < writers; i++ {
		go func(k int) {
			ready <- struct{}{}
			<-start
			mode := 0
			for j := 0; j < n; j++ {
				k = (k + 1) % max
				switch mode {
				case 0:
					m.Set(Key(k), Value(k))
					mode = 1
				case 1:
					m.Delete(Key(k))
					mode = 0
				}
			}
			wg.Done()
		}(i)
	}

	for i := 0; i < rw; i++ {
		<-ready
	}

	b.ResetTimer()
	close(start)
	wg.Wait()
}