base-golang/pkg/utils/map.go

package utils

import (
	"errors"
	"fmt"
	"reflect"
	"sync"
	"sync/atomic"
)

var _ ConcurrentMap = (*concurrentMap)(nil)

type ConcurrentMap interface {
	i()
	Store(key, value any)
	Delete(key any)
	Range(f func(key, value any) bool)
	Load(key any) (value any, ok bool)
	LoadOrStore(key, value any) (actual any, loaded bool)
}

type concurrentMap struct {
	m         sync.Map
	keyType   reflect.Type
	valueType reflect.Type
	counter   int64
}

func NewConcurrentMap(keyType, valueType reflect.Type) (ConcurrentMap, error) {
	if keyType == nil {
		return nil, errors.New("nil key type")
	}
	if !keyType.Comparable() {
		return nil, fmt.Errorf("incomparable key type: %s", keyType)
	}
	if valueType == nil {
		return nil, errors.New("nil value type")
	}
	cMap := &concurrentMap{
		keyType:   keyType,
		valueType: valueType,
		counter:   0,
	}
	return cMap, nil
}

func (cMap *concurrentMap) i() {}

func (cMap *concurrentMap) Delete(key any) {
	if reflect.TypeOf(key) != cMap.keyType {
		return
	}
	cMap.m.Delete(key)
	atomic.AddInt64(&cMap.counter, -1)
}

func (cMap *concurrentMap) Len() int64 {
	return cMap.counter
}

func (cMap *concurrentMap) Load(key any) (value any, ok bool) {
	if reflect.TypeOf(key) != cMap.keyType {
		return
	}
	return cMap.m.Load(key)
}

func (cMap *concurrentMap) LoadOrStore(key, value any) (actual any, loaded bool) {
	if reflect.TypeOf(key) != cMap.keyType {
		panic(fmt.Errorf("wrong key type: %v", reflect.TypeOf(key)))
	}
	if reflect.TypeOf(value) != cMap.valueType {
		panic(fmt.Errorf("wrong value type: %v", reflect.TypeOf(value)))
	}
	actual, loaded = cMap.m.LoadOrStore(key, value)
	return
}

func (cMap *concurrentMap) Range(f func(key, value any) bool) {
	cMap.m.Range(f)
}

func (cMap *concurrentMap) Store(key, value any) {
	if reflect.TypeOf(key) != cMap.keyType {
		panic(fmt.Errorf("wrong key type: %v", reflect.TypeOf(key)))
	}
	if reflect.TypeOf(value) != cMap.valueType {
		panic(fmt.Errorf("wrong value type: %v", reflect.TypeOf(value)))
	}
	if _, ok := cMap.m.Load(key); ok {
		cMap.Delete(key)
	}
	cMap.m.Store(key, value)
	atomic.AddInt64(&cMap.counter, 1)
}