sonet/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.Range(func(key, value string) bool {
		sum++
		if key != value {
			t.Errorf("value load 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 || v != k {
					t.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
	})
}