github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/core/container/rcu_map.go (about)

     1  package container
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  )
     7  
     8  type RCUMapElement[Key comparable, Val any] struct {
     9  	Key   Key
    10  	Value Val
    11  }
    12  
    13  type RCUDeleteNode[Val any] struct {
    14  	Value Val
    15  	Ok    bool
    16  }
    17  
    18  // RCUMap 这个Map的实现只适合少量key-value, 或者几乎无写的场景
    19  // 在大量key-value时拷贝数据的开销很大
    20  type RCUMap[Key comparable, Val any] struct {
    21  	mu      sync.Mutex // 串行写入操作, 读取操作不需要上锁
    22  	pointer atomic.Pointer[map[Key]Val]
    23  }
    24  
    25  func NewRCUMap[K comparable, V any](size int) *RCUMap[K, V] {
    26  	m := new(RCUMap[K, V])
    27  	tmp := make(map[K]V, size)
    28  	m.pointer.Store(&tmp)
    29  	return m
    30  }
    31  
    32  func (R *RCUMap[Key, Val]) LoadOk(key Key) (Val, bool) {
    33  	snapshot := R.pointer.Load()
    34  	val, ok := (*snapshot)[key]
    35  	return val, ok
    36  }
    37  
    38  func (R *RCUMap[Key, Val]) Range(fn func(key Key, val Val) bool) {
    39  	snapshot := R.pointer.Load()
    40  	for k, v := range *snapshot {
    41  		if !fn(k, v) {
    42  			break
    43  		}
    44  	}
    45  }
    46  
    47  func (R *RCUMap[Key, Val]) Store(key Key, val Val) {
    48  	R.StoreMulti([]RCUMapElement[Key, Val]{{Key: key, Value: val}})
    49  }
    50  
    51  func (R *RCUMap[Key, Val]) StoreMulti(kvs []RCUMapElement[Key, Val]) {
    52  	R.StoreAndDeleteMulti(kvs, nil)
    53  }
    54  
    55  func (R *RCUMap[Key, Val]) DeleteOk(key Key) (Val, bool) {
    56  	val := R.DeleteMulti([]Key{key})
    57  	if val == nil || len(val) == 0 {
    58  		return *new(Val), false
    59  	}
    60  	return val[0].Value, val[0].Ok
    61  }
    62  
    63  func (R *RCUMap[Key, Val]) Delete(key Key) {
    64  	R.DeleteOk(key)
    65  }
    66  
    67  func (R *RCUMap[Key, Val]) DeleteMulti(keys []Key) []RCUDeleteNode[Val] {
    68  	return R.StoreAndDeleteMulti(nil, keys)
    69  }
    70  
    71  func (R *RCUMap[Key, Val]) StoreAndDeleteMulti(kvs []RCUMapElement[Key, Val], ks []Key) (delRes []RCUDeleteNode[Val]) {
    72  	if len(kvs) == 0 && len(ks) == 0 {
    73  		return
    74  	}
    75  	R.mu.Lock()
    76  	defer R.mu.Unlock()
    77  	copyMap := R.copy()
    78  	for _, kv := range kvs {
    79  		copyMap[kv.Key] = kv.Value
    80  	}
    81  	if len(kvs) > 0 && len(ks) == 0 {
    82  		R.pointer.Store(&copyMap)
    83  		return
    84  	}
    85  	values := make([]RCUDeleteNode[Val], 0, len(ks))
    86  	for _, key := range ks {
    87  		val, ok := copyMap[key]
    88  		values = append(values, RCUDeleteNode[Val]{
    89  			Value: val,
    90  			Ok:    ok,
    91  		})
    92  		if !ok {
    93  			continue
    94  		}
    95  		delete(copyMap, key)
    96  	}
    97  	R.pointer.Store(&copyMap)
    98  	return values
    99  }
   100  
   101  func (R *RCUMap[Key, Val]) Len() int {
   102  	return len(*R.pointer.Load())
   103  }
   104  
   105  func (R *RCUMap[Key, Val]) copy() map[Key]Val {
   106  	snapshot := *R.pointer.Load()
   107  	copyMap := make(map[Key]Val, len(snapshot))
   108  	for k, v := range snapshot {
   109  		copyMap[k] = v
   110  	}
   111  	return copyMap
   112  }