map.go 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. package utils
  2. import (
  3. "errors"
  4. "fmt"
  5. "reflect"
  6. "sync"
  7. "sync/atomic"
  8. )
  9. var _ ConcurrentMap = (*concurrentMap)(nil)
  10. type ConcurrentMap interface {
  11. i()
  12. Store(key, value interface{})
  13. Delete(key interface{})
  14. Range(f func(key, value interface{}) bool)
  15. Load(key interface{}) (value interface{}, ok bool)
  16. LoadOrStore(key, value interface{}) (actual interface{}, loaded bool)
  17. }
  18. type concurrentMap struct {
  19. m sync.Map
  20. keyType reflect.Type
  21. valueType reflect.Type
  22. counter int64
  23. }
  24. func NewConcurrentMap(keyType, valueType reflect.Type) (ConcurrentMap, error) {
  25. if keyType == nil {
  26. return nil, errors.New("nil key type")
  27. }
  28. if !keyType.Comparable() {
  29. return nil, fmt.Errorf("incomparable key type: %s", keyType)
  30. }
  31. if valueType == nil {
  32. return nil, errors.New("nil value type")
  33. }
  34. cMap := &concurrentMap{
  35. keyType: keyType,
  36. valueType: valueType,
  37. counter: 0,
  38. }
  39. return cMap, nil
  40. }
  41. func (cMap *concurrentMap) i() {}
  42. func (cMap *concurrentMap) Delete(key interface{}) {
  43. if reflect.TypeOf(key) != cMap.keyType {
  44. return
  45. }
  46. cMap.m.Delete(key)
  47. atomic.AddInt64(&cMap.counter, -1)
  48. }
  49. func (cMap *concurrentMap) Len() int64 {
  50. return cMap.counter
  51. }
  52. func (cMap *concurrentMap) Load(key interface{}) (value interface{}, ok bool) {
  53. if reflect.TypeOf(key) != cMap.keyType {
  54. return
  55. }
  56. return cMap.m.Load(key)
  57. }
  58. func (cMap *concurrentMap) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) {
  59. if reflect.TypeOf(key) != cMap.keyType {
  60. panic(fmt.Errorf("wrong key type: %v", reflect.TypeOf(key)))
  61. }
  62. if reflect.TypeOf(value) != cMap.valueType {
  63. panic(fmt.Errorf("wrong value type: %v", reflect.TypeOf(value)))
  64. }
  65. actual, loaded = cMap.m.LoadOrStore(key, value)
  66. return
  67. }
  68. func (cMap *concurrentMap) Range(f func(key, value interface{}) bool) {
  69. cMap.m.Range(f)
  70. }
  71. func (cMap *concurrentMap) Store(key, value interface{}) {
  72. if reflect.TypeOf(key) != cMap.keyType {
  73. panic(fmt.Errorf("wrong key type: %v", reflect.TypeOf(key)))
  74. }
  75. if reflect.TypeOf(value) != cMap.valueType {
  76. panic(fmt.Errorf("wrong value type: %v", reflect.TypeOf(value)))
  77. }
  78. if _, ok := cMap.m.Load(key); ok {
  79. cMap.Delete(key)
  80. }
  81. cMap.m.Store(key, value)
  82. atomic.AddInt64(&cMap.counter, 1)
  83. }