github.com/benz9527/toy-box/algo@v0.0.0-20240221120937-66c0c6bd5abd/kv/thread_safe_map.go (about)

     1  package kv
     2  
     3  import (
     4  	"io"
     5  	"log/slog"
     6  	"reflect"
     7  	"strings"
     8  	"sync"
     9  )
    10  
    11  type threadSafeMap[T any] struct {
    12  	lock         sync.RWMutex
    13  	items        map[string]T
    14  	isCloserItem bool
    15  }
    16  
    17  func (t *threadSafeMap[T]) AddOrUpdate(key string, obj T) {
    18  	t.lock.Lock()
    19  	defer t.lock.Unlock()
    20  	t.items[key] = obj
    21  }
    22  
    23  func (t *threadSafeMap[T]) Replace(items map[string]T) {
    24  	t.lock.Lock()
    25  	defer t.lock.Unlock()
    26  	t.items = items
    27  }
    28  
    29  func (t *threadSafeMap[T]) Delete(key string) {
    30  	t.lock.Lock()
    31  	defer t.lock.Unlock()
    32  	if _, exists := t.items[key]; exists {
    33  		delete(t.items, key)
    34  	}
    35  }
    36  
    37  func (t *threadSafeMap[T]) Get(key string) (item T, exists bool) {
    38  	t.lock.RLock()
    39  	defer t.lock.RUnlock()
    40  	item, exists = t.items[key]
    41  	return
    42  }
    43  
    44  func (t *threadSafeMap[T]) ListKeys(filters ...SafeStoreKeyFilterFunc) []string {
    45  
    46  	realFilters := make([]SafeStoreKeyFilterFunc, 0, len(filters))
    47  	for _, filter := range filters {
    48  		if filter != nil {
    49  			realFilters = append(realFilters, filter)
    50  		}
    51  	}
    52  	if len(realFilters) == 0 {
    53  		realFilters = append(realFilters, defaultAllKeysFilter)
    54  	}
    55  
    56  	t.lock.RLock()
    57  	defer t.lock.RUnlock()
    58  
    59  	keys := make([]string, 0, len(t.items))
    60  	for key := range t.items {
    61  		for _, filter := range realFilters {
    62  			if filter(key) {
    63  				keys = append(keys, key)
    64  				break
    65  			}
    66  		}
    67  	}
    68  	return keys
    69  }
    70  
    71  func (t *threadSafeMap[T]) ListValues(keys ...string) (items []T) {
    72  	realKeys := make([]string, 0, len(keys))
    73  	for _, key := range keys {
    74  		if len(strings.TrimSpace(key)) > 0 {
    75  			realKeys = append(realKeys, key)
    76  		}
    77  	}
    78  
    79  	contains := func(keys []string, key string) bool {
    80  		for _, k := range keys {
    81  			if k == key {
    82  				return true
    83  			}
    84  		}
    85  		return false
    86  	}
    87  
    88  	t.lock.RLock()
    89  	defer t.lock.RUnlock()
    90  	values := make([]T, 0, len(t.items))
    91  	for key, item := range t.items {
    92  		i := item
    93  		if len(realKeys) > 0 && contains(realKeys, key) {
    94  			values = append(values, i)
    95  		} else if len(realKeys) == 0 {
    96  			values = append(values, i)
    97  		}
    98  	}
    99  	return values
   100  }
   101  
   102  func (t *threadSafeMap[T]) Close() error {
   103  	t.lock.Lock()
   104  	defer t.lock.Unlock()
   105  
   106  	if t.isCloserItem {
   107  		for _, item := range t.items {
   108  			if reflect.ValueOf(item).IsNil() {
   109  				continue
   110  			}
   111  			typ := reflect.TypeOf(item)
   112  			if typ.Implements(reflect.TypeOf((*io.Closer)(nil)).Elem()) {
   113  				vals := reflect.ValueOf(item).MethodByName("Close").Call([]reflect.Value{})
   114  				if len(vals) > 0 && !vals[0].IsNil() {
   115  					intf := vals[0].Elem().Interface()
   116  					switch intf.(type) {
   117  					case error:
   118  						err := intf.(error)
   119  						slog.Error("Close info", "error", err)
   120  					}
   121  				}
   122  			}
   123  		}
   124  	}
   125  
   126  	t.items = nil
   127  	return nil
   128  }
   129  
   130  func NewThreadSafeMap[T any]() IThreadSafeStore[T] {
   131  	isCloserItem := false
   132  	nilT := *new(T)
   133  	if reflect.TypeOf(nilT).Implements(reflect.TypeOf((*io.Closer)(nil)).Elem()) {
   134  		isCloserItem = true
   135  	}
   136  	return &threadSafeMap[T]{items: make(map[string]T), isCloserItem: isCloserItem}
   137  }