github.com/chain5j/chain5j-pkg@v1.0.7/collection/maps/shardedmap/strmap.go (about)

     1  package shardedmap
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  type StrMap struct {
     8  	shardCount uint64
     9  	mutexes    []sync.RWMutex
    10  	maps       []map[string]interface{}
    11  }
    12  
    13  // NewStrMap ...
    14  func NewStrMap(shardCount int) *StrMap {
    15  	if shardCount <= 0 {
    16  		shardCount = defaultShards
    17  	}
    18  
    19  	sm := &StrMap{
    20  		shardCount: uint64(shardCount),
    21  		mutexes:    make([]sync.RWMutex, shardCount),
    22  		maps:       make([]map[string]interface{}, shardCount),
    23  	}
    24  
    25  	for i := range sm.maps {
    26  		sm.maps[i] = make(map[string]interface{})
    27  	}
    28  
    29  	return sm
    30  }
    31  
    32  func (sm *StrMap) pickShard(key string) uint64 {
    33  	return memHashString(key) % sm.shardCount
    34  }
    35  
    36  // Set ...
    37  func (sm *StrMap) Set(key string, value interface{}) {
    38  	shard := sm.pickShard(key)
    39  	sm.mutexes[shard].Lock()
    40  	sm.maps[shard][key] = value
    41  	sm.mutexes[shard].Unlock()
    42  }
    43  
    44  // Get ...
    45  func (sm *StrMap) Get(key string) (interface{}, bool) {
    46  	shard := sm.pickShard(key)
    47  	sm.mutexes[shard].RLock()
    48  	value, ok := sm.maps[shard][key]
    49  	sm.mutexes[shard].RUnlock()
    50  	return value, ok
    51  }
    52  
    53  // GetOrSet ...
    54  func (sm *StrMap) GetOrSet(key string, value interface{}) (actual interface{}, loaded bool) {
    55  	shard := sm.pickShard(key)
    56  	sm.mutexes[shard].RLock()
    57  	// Fast path assuming value has a somewhat high chance of already being
    58  	// there.
    59  	if actual, loaded = sm.maps[shard][key]; loaded {
    60  		sm.mutexes[shard].RUnlock()
    61  		return
    62  	}
    63  	sm.mutexes[shard].RUnlock()
    64  	// Gotta check again, unfortunately
    65  	sm.mutexes[shard].Lock()
    66  	if actual, loaded = sm.maps[shard][key]; loaded {
    67  		sm.mutexes[shard].Unlock()
    68  		return
    69  	}
    70  	sm.maps[shard][key] = value
    71  	sm.mutexes[shard].Unlock()
    72  	return value, loaded
    73  }
    74  
    75  // Delete ...
    76  func (sm *StrMap) Delete(key string) {
    77  	shard := sm.pickShard(key)
    78  	sm.mutexes[shard].Lock()
    79  	delete(sm.maps[shard], key)
    80  	sm.mutexes[shard].Unlock()
    81  }
    82  
    83  // Range is modeled after sync.Map.Range. It calls f sequentially for each key
    84  // and value present in each of the shards in the map. If f returns false, range
    85  // stops the iteration.
    86  //
    87  // No key will be visited more than once, but if any value is inserted
    88  // concurrently, Range may or may not visit it. Similarly, if a value is
    89  // modified concurrently, Range may visit the previous or newest version of said
    90  // value.
    91  func (sm *StrMap) Range(f func(key string, value interface{}) bool) {
    92  	for shard := range sm.mutexes {
    93  		sm.mutexes[shard].RLock()
    94  		for key, value := range sm.maps[shard] {
    95  			if !f(key, value) {
    96  				sm.mutexes[shard].RUnlock()
    97  				return
    98  			}
    99  		}
   100  		sm.mutexes[shard].RUnlock()
   101  	}
   102  }
   103  
   104  // ConcRange ranges concurrently over all the shards, calling f sequentially
   105  // over each shard's key and value. If f returns false, range stops the
   106  // iteration on that shard (but the other shards continue until completion).
   107  //
   108  // No key will be visited more than once, but if any value is inserted
   109  // concurrently, Range may or may not visit it. Similarly, if a value is
   110  // modified concurrently, Range may visit the previous or newest version of said
   111  // value.
   112  func (sm *StrMap) ConcRange(f func(key string, value interface{}) bool) {
   113  	var wg sync.WaitGroup
   114  	wg.Add(int(sm.shardCount))
   115  	for shard := range sm.mutexes {
   116  		go func(shard int) {
   117  			sm.mutexes[shard].RLock()
   118  			for key, value := range sm.maps[shard] {
   119  				if !f(key, value) {
   120  					sm.mutexes[shard].RUnlock()
   121  					wg.Done()
   122  					return
   123  				}
   124  			}
   125  			sm.mutexes[shard].RUnlock()
   126  			wg.Done()
   127  		}(shard)
   128  	}
   129  	wg.Wait()
   130  }
   131  
   132  // AsyncRange is exactly like ConcRange, but doesn't wait until all shards are
   133  // done. This is usually ok, although calls that appear to happen "sequentially"
   134  // on the same goroutine might get the before or after AsyncRange values, which
   135  // might be surprising behaviour. When that's not desirable, use ConcRange.
   136  func (sm *StrMap) AsyncRange(f func(key string, value interface{}) bool) {
   137  	for shard := range sm.mutexes {
   138  		go func(shard int) {
   139  			sm.mutexes[shard].RLock()
   140  			for key, value := range sm.maps[shard] {
   141  				if !f(key, value) {
   142  					sm.mutexes[shard].RUnlock()
   143  					return
   144  				}
   145  			}
   146  			sm.mutexes[shard].RUnlock()
   147  		}(shard)
   148  	}
   149  }
   150  
   151  // Len returns the number of values in map.
   152  func (sm *StrMap) Len() int {
   153  	var count int
   154  	for i := uint64(0); i < sm.shardCount; i++ {
   155  		sm.mutexes[i].RLock()
   156  		count += len(sm.maps[i])
   157  		sm.mutexes[i].RUnlock()
   158  	}
   159  	return count
   160  }