``` feat(pet): 重构宠物繁殖系统，添加蛋孵化功能
This commit is contained in:
286
common/utils/concurrent-swiss-map/concurrent_swiss_map.go
Normal file
286
common/utils/concurrent-swiss-map/concurrent_swiss_map.go
Normal file
@@ -0,0 +1,286 @@
|
||||
package csmap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"sync"
|
||||
|
||||
"github.com/mhmtszr/concurrent-swiss-map/maphash"
|
||||
"github.com/panjf2000/ants/v2"
|
||||
|
||||
"github.com/mhmtszr/concurrent-swiss-map/swiss"
|
||||
)
|
||||
|
||||
type CsMap[K comparable, V any] struct {
|
||||
hasher func(key K) uint64
|
||||
shards []shard[K, V]
|
||||
shardCount uint64
|
||||
size uint64
|
||||
}
|
||||
|
||||
type HashShardPair[K comparable, V any] struct {
|
||||
shard shard[K, V]
|
||||
hash uint64
|
||||
}
|
||||
|
||||
type shard[K comparable, V any] struct {
|
||||
items *swiss.Map[K, V]
|
||||
*sync.RWMutex
|
||||
}
|
||||
|
||||
// OptFunc is a type that is used in New function for passing options.
|
||||
type OptFunc[K comparable, V any] func(o *CsMap[K, V])
|
||||
|
||||
// New function creates *CsMap[K, V].
|
||||
func New[K comparable, V any](options ...OptFunc[K, V]) *CsMap[K, V] {
|
||||
m := CsMap[K, V]{
|
||||
hasher: maphash.NewHasher[K]().Hash,
|
||||
shardCount: 32,
|
||||
}
|
||||
for _, option := range options {
|
||||
option(&m)
|
||||
}
|
||||
|
||||
m.shards = make([]shard[K, V], m.shardCount)
|
||||
|
||||
for i := 0; i < int(m.shardCount); i++ {
|
||||
m.shards[i] = shard[K, V]{items: swiss.NewMap[K, V](uint32((m.size / m.shardCount) + 1)), RWMutex: &sync.RWMutex{}}
|
||||
}
|
||||
return &m
|
||||
}
|
||||
|
||||
// Create creates *CsMap.
|
||||
//
|
||||
// Deprecated: New function should be used instead.
|
||||
func Create[K comparable, V any](options ...func(options *CsMap[K, V])) *CsMap[K, V] {
|
||||
m := CsMap[K, V]{
|
||||
hasher: maphash.NewHasher[K]().Hash,
|
||||
shardCount: 32,
|
||||
}
|
||||
for _, option := range options {
|
||||
option(&m)
|
||||
}
|
||||
|
||||
m.shards = make([]shard[K, V], m.shardCount)
|
||||
|
||||
for i := 0; i < int(m.shardCount); i++ {
|
||||
m.shards[i] = shard[K, V]{items: swiss.NewMap[K, V](uint32((m.size / m.shardCount) + 1)), RWMutex: &sync.RWMutex{}}
|
||||
}
|
||||
return &m
|
||||
}
|
||||
|
||||
func WithShardCount[K comparable, V any](count uint64) func(csMap *CsMap[K, V]) {
|
||||
return func(csMap *CsMap[K, V]) {
|
||||
csMap.shardCount = count
|
||||
}
|
||||
}
|
||||
|
||||
func WithCustomHasher[K comparable, V any](h func(key K) uint64) func(csMap *CsMap[K, V]) {
|
||||
return func(csMap *CsMap[K, V]) {
|
||||
csMap.hasher = h
|
||||
}
|
||||
}
|
||||
|
||||
func WithSize[K comparable, V any](size uint64) func(csMap *CsMap[K, V]) {
|
||||
return func(csMap *CsMap[K, V]) {
|
||||
csMap.size = size
|
||||
}
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) getShard(key K) HashShardPair[K, V] {
|
||||
u := m.hasher(key)
|
||||
return HashShardPair[K, V]{
|
||||
hash: u,
|
||||
shard: m.shards[u%m.shardCount],
|
||||
}
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) Store(key K, value V) {
|
||||
hashShardPair := m.getShard(key)
|
||||
shard := hashShardPair.shard
|
||||
shard.Lock()
|
||||
shard.items.PutWithHash(key, value, hashShardPair.hash)
|
||||
shard.Unlock()
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) Delete(key K) bool {
|
||||
hashShardPair := m.getShard(key)
|
||||
shard := hashShardPair.shard
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
return shard.items.DeleteWithHash(key, hashShardPair.hash)
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) DeleteIf(key K, condition func(value V) bool) bool {
|
||||
hashShardPair := m.getShard(key)
|
||||
shard := hashShardPair.shard
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
value, ok := shard.items.GetWithHash(key, hashShardPair.hash)
|
||||
if ok && condition(value) {
|
||||
return shard.items.DeleteWithHash(key, hashShardPair.hash)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) Load(key K) (V, bool) {
|
||||
hashShardPair := m.getShard(key)
|
||||
shard := hashShardPair.shard
|
||||
shard.RLock()
|
||||
defer shard.RUnlock()
|
||||
return shard.items.GetWithHash(key, hashShardPair.hash)
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) Has(key K) bool {
|
||||
hashShardPair := m.getShard(key)
|
||||
shard := hashShardPair.shard
|
||||
shard.RLock()
|
||||
defer shard.RUnlock()
|
||||
return shard.items.HasWithHash(key, hashShardPair.hash)
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) Clear() {
|
||||
for i := range m.shards {
|
||||
shard := m.shards[i]
|
||||
|
||||
shard.Lock()
|
||||
shard.items.Clear()
|
||||
shard.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) Count() int {
|
||||
count := 0
|
||||
for i := range m.shards {
|
||||
shard := m.shards[i]
|
||||
shard.RLock()
|
||||
count += shard.items.Count()
|
||||
shard.RUnlock()
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) SetIfAbsent(key K, value V) {
|
||||
hashShardPair := m.getShard(key)
|
||||
shard := hashShardPair.shard
|
||||
shard.Lock()
|
||||
_, ok := shard.items.GetWithHash(key, hashShardPair.hash)
|
||||
if !ok {
|
||||
shard.items.PutWithHash(key, value, hashShardPair.hash)
|
||||
}
|
||||
shard.Unlock()
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) SetIf(key K, conditionFn func(previousVale V, previousFound bool) (value V, set bool)) {
|
||||
hashShardPair := m.getShard(key)
|
||||
shard := hashShardPair.shard
|
||||
shard.Lock()
|
||||
value, found := shard.items.GetWithHash(key, hashShardPair.hash)
|
||||
value, ok := conditionFn(value, found)
|
||||
if ok {
|
||||
shard.items.PutWithHash(key, value, hashShardPair.hash)
|
||||
}
|
||||
shard.Unlock()
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) SetIfPresent(key K, value V) {
|
||||
hashShardPair := m.getShard(key)
|
||||
shard := hashShardPair.shard
|
||||
shard.Lock()
|
||||
_, ok := shard.items.GetWithHash(key, hashShardPair.hash)
|
||||
if ok {
|
||||
shard.items.PutWithHash(key, value, hashShardPair.hash)
|
||||
}
|
||||
shard.Unlock()
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) IsEmpty() bool {
|
||||
return m.Count() == 0
|
||||
}
|
||||
|
||||
type Tuple[K comparable, V any] struct {
|
||||
Key K
|
||||
Val V
|
||||
}
|
||||
|
||||
// Range If the callback function returns true iteration will stop.
|
||||
func (m *CsMap[K, V]) Range(f func(key K, value V) (stop bool)) {
|
||||
ch := make(chan Tuple[K, V], m.Count())
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
listenCompleted := m.listen(f, ch)
|
||||
m.produce(ctx, ch)
|
||||
listenCompleted.Wait()
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) MarshalJSON() ([]byte, error) {
|
||||
tmp := make(map[K]V, m.Count())
|
||||
m.Range(func(key K, value V) (stop bool) {
|
||||
tmp[key] = value
|
||||
return false
|
||||
})
|
||||
return json.Marshal(tmp)
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) UnmarshalJSON(b []byte) error {
|
||||
tmp := make(map[K]V, m.Count())
|
||||
|
||||
if err := json.Unmarshal(b, &tmp); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for key, val := range tmp {
|
||||
m.Store(key, val)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *CsMap[K, V]) produce(ctx context.Context, ch chan Tuple[K, V]) {
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(len(m.shards))
|
||||
|
||||
var producepool, _ = ants.NewPoolWithFuncGeneric(-1, func(i int) {
|
||||
defer wg.Done()
|
||||
|
||||
shard := m.shards[i]
|
||||
shard.RLock()
|
||||
shard.items.Iter(func(k K, v V) (stop bool) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return true
|
||||
default:
|
||||
ch <- Tuple[K, V]{Key: k, Val: v}
|
||||
}
|
||||
return false
|
||||
})
|
||||
shard.RUnlock()
|
||||
})
|
||||
|
||||
for i := range m.shards {
|
||||
producepool.Invoke(i)
|
||||
}
|
||||
|
||||
pool.Submit(func() {
|
||||
wg.Wait()
|
||||
close(ch)
|
||||
})
|
||||
}
|
||||
|
||||
var pool, _ = ants.NewPool(-1)
|
||||
|
||||
func (m *CsMap[K, V]) listen(f func(key K, value V) (stop bool), ch chan Tuple[K, V]) *sync.WaitGroup {
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
pool.Submit(func() {
|
||||
defer wg.Done()
|
||||
for t := range ch {
|
||||
if stop := f(t.Key, t.Val); stop {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
return &wg
|
||||
}
|
||||
Reference in New Issue
Block a user