sig-pub/pkg/utils/collect/concurrent_map_test.go

package collect

import (
	"fmt"
	"math/rand"
	"strconv"
	"sync"
	"testing"
	"time"
)

func TestConcurrentMap(t *testing.T) {
	cm := NewConcurrentMap[string, string](19, func(k string) string { return k })
	concurrent := 1000

	wg := sync.WaitGroup{}
	wg.Add(concurrent)

	for i := 0; i < concurrent; i++ {
		go func(loop int) {
			for j := 0; j < concurrent; j++ {
				k := fmt.Sprintf("%d-%d", loop, j)
				cm.Store(k, k)
			}
			wg.Done()
		}(i)
	}
	wg.Wait()

	var sum int
	cm.RangeUpdate(func(key, value string) (bool, bool, string) {
		if key != value {
			panic(fmt.Errorf("value load error: %s:%s", key, value))
		}
		return true, false, value + "-update"
	})
	cm.Range(func(key, value string) bool {
		sum++
		if (key + "-update") != value {
			panic(fmt.Errorf("value update error: %s:%s", key, value))
		}
		return true
	})
	if sum != (concurrent * concurrent) {
		t.Errorf("count error: %d", sum)
	}

	wg.Add(concurrent)
	for i := 0; i < concurrent; i++ {
		go func() {
			r := rand.New(rand.NewSource(time.Now().UnixMilli()))
			for i := 0; i < concurrent; i++ {
				k := fmt.Sprintf("%d-%d", r.Intn(concurrent), r.Intn(concurrent))
				v, ok := cm.Load(k)
				if !ok || (k+"-update") != v {
					panic(fmt.Errorf("value load error: %s", k))
				}
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

type counter struct {
	c int
}

func (c *counter) increment() {
	c.c += 1
}

func TestComputeIfAbsent(t *testing.T) {
	cm := NewConcurrentMap[int, *counter](16, func(k int) string { return strconv.Itoa(k) })
	concurrent := 1000
	add := 10
	wg := &sync.WaitGroup{}
	wg.Add(concurrent)

	for i := 0; i < concurrent; i++ {
		go func() {
			defer wg.Done()

			r := rand.New(rand.NewSource(time.Now().UnixMilli()))
			v, mapped := cm.ComputeIfAbsent(r.Intn(10), func(k int) *counter {
				return &counter{}
			})
			if !mapped {
				return
			}
			for i := 0; i < add; i++ {
				v.increment()
			}
		}()
	}

	wg.Wait()

	cm.Range(func(k int, v *counter) bool {
		if v.c != add {
			t.Error("error value...")
		}
		return true
	})
}

func TestBenchmark(t *testing.T) {
	cm := NewConcurrentMap[string, int](16, func(s string) string { return s })
	for i := range 10000 {
		cm.Store(strconv.Itoa(i), i)
	}

	// sm := &sync.Map{}
	// for i := range 10000 {
	// 	sm.Store(strconv.Itoa(i), i)
	// }

	// mm := make(map[string]int, 10000)
	// for i := range 10000 {
	// 	mm[strconv.Itoa(i)] = i
	// }
	wg := &sync.WaitGroup{}
	for range 4 {
		wg.Add(1)
		go func() {
			defer wg.Done()

			now := time.Now()
			for range 10000 {
				for i := range 10000 {
					_, _ = cm.Load(strconv.Itoa(i))
					// _, _ = sm.Load(strconv.Itoa(i))
					// _, _ = mm[strconv.Itoa(i)]
				}
			}
			fmt.Printf("done: use %dms\n", time.Now().UnixMilli()-now.UnixMilli())
		}()
	}
	wg.Wait()
}