github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/sync/map_int.go (about)

     1  // Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package syncutils
     4  
     5  import (
     6  	"runtime"
     7  	"sync"
     8  )
     9  
    10  type KType interface {
    11  	int | int16 | int32 | int64 | uint | uint16 | uint32 | uint64 | uintptr
    12  }
    13  
    14  type VType interface {
    15  	any
    16  }
    17  
    18  type IntMap[K KType, V VType] struct {
    19  	count   int
    20  	m       []map[K]V
    21  	lockers []*sync.RWMutex
    22  }
    23  
    24  func NewIntMap[K KType, V VType]() *IntMap[K, V] {
    25  	var count = runtime.NumCPU() * 8
    26  	if count <= 0 {
    27  		count = 32
    28  	}
    29  
    30  	var m = []map[K]V{}
    31  	var lockers = []*sync.RWMutex{}
    32  	for i := 0; i < count; i++ {
    33  		m = append(m, map[K]V{})
    34  		lockers = append(lockers, &sync.RWMutex{})
    35  	}
    36  
    37  	return &IntMap[K, V]{
    38  		count:   count,
    39  		m:       m,
    40  		lockers: lockers,
    41  	}
    42  }
    43  
    44  func (this *IntMap[K, V]) Put(k K, v V) {
    45  	var index = this.index(k)
    46  	this.lockers[index].Lock()
    47  	this.m[index][k] = v
    48  	this.lockers[index].Unlock()
    49  }
    50  
    51  func (this *IntMap[K, V]) PutCompact(k K, v V, compactFunc func(oldV V, newV V) V) {
    52  	var index = this.index(k)
    53  	this.lockers[index].Lock()
    54  	// 再次检查是否已经存在，如果已经存在则合并
    55  	oldV, ok := this.m[index][k]
    56  	if ok {
    57  		this.m[index][k] = compactFunc(oldV, v)
    58  	} else {
    59  		this.m[index][k] = v
    60  	}
    61  	this.lockers[index].Unlock()
    62  }
    63  
    64  func (this *IntMap[K, V]) Has(k K) bool {
    65  	var index = this.index(k)
    66  	this.lockers[index].RLock()
    67  	_, ok := this.m[index][k]
    68  	this.lockers[index].RUnlock()
    69  	return ok
    70  }
    71  
    72  func (this *IntMap[K, V]) Get(k K) (value V) {
    73  	var index = this.index(k)
    74  	this.lockers[index].RLock()
    75  	value = this.m[index][k]
    76  	this.lockers[index].RUnlock()
    77  	return
    78  }
    79  
    80  func (this *IntMap[K, V]) GetOk(k K) (value V, ok bool) {
    81  	var index = this.index(k)
    82  	this.lockers[index].RLock()
    83  	value, ok = this.m[index][k]
    84  	this.lockers[index].RUnlock()
    85  	return
    86  }
    87  
    88  func (this *IntMap[K, V]) Delete(k K) {
    89  	var index = this.index(k)
    90  	this.lockers[index].Lock()
    91  	delete(this.m[index], k)
    92  	this.lockers[index].Unlock()
    93  }
    94  
    95  func (this *IntMap[K, V]) DeleteUnsafe(k K) {
    96  	var index = this.index(k)
    97  	delete(this.m[index], k)
    98  }
    99  
   100  func (this *IntMap[K, V]) Len() int {
   101  	var l int
   102  	for i := 0; i < this.count; i++ {
   103  		this.lockers[i].RLock()
   104  		l += len(this.m[i])
   105  		this.lockers[i].RUnlock()
   106  	}
   107  	return l
   108  }
   109  
   110  func (this *IntMap[K, V]) ForEachRead(iterator func(k K, v V)) {
   111  	for i := 0; i < this.count; i++ {
   112  		this.lockers[i].RLock()
   113  		for k, v := range this.m[i] {
   114  			iterator(k, v)
   115  		}
   116  		this.lockers[i].RUnlock()
   117  	}
   118  }
   119  
   120  func (this *IntMap[K, V]) ForEachWrite(iterator func(k K, v V)) {
   121  	for i := 0; i < this.count; i++ {
   122  		this.lockers[i].Lock()
   123  		for k, v := range this.m[i] {
   124  			iterator(k, v)
   125  		}
   126  		this.lockers[i].Unlock()
   127  	}
   128  }
   129  
   130  func (this *IntMap[K, V]) index(k K) int {
   131  	var index = int(k % K(this.count))
   132  	if index < 0 {
   133  		index = -index
   134  	}
   135  	return index
   136  }