go-redis/datastruct/dict/sync_dict.go

// Author: simon (ynwdlxm@163.com)
// Date: 2025/9/15 17:07
// Desc: 线程安全的字典实现，基于 sync.Map

package dict

import (
	"math/rand"
	"sync"
)

// SyncDict 是基于 sync.Map 的线程安全字典实现
// 实现了 Dict 接口的所有方法
type SyncDict struct {
	m sync.Map
}

// Get 根据键获取值
// 参数 key 要查找的键
// 返回值 value 是与键关联的值，ok 表示键是否存在
func (s *SyncDict) Get(key string) (any, bool) {
	return s.m.Load(key)
}

// Len 返回字典中键值对的数量
// 通过遍历所有元素来计算长度
func (s *SyncDict) Len() int {
	length := 0
	s.m.Range(func(key, value any) bool {
		length++
		return true
	})
	return length
}

// Put 插入或更新键值对
// 参数 key 键
// 参数 value 值
// 返回值 result 表示操作结果：1表示新增，0表示更新
func (s *SyncDict) Put(key string, value any) (result int) {
	_, loaded := s.m.LoadOrStore(key, value)
	if loaded {
		// 如果键已存在，则更新它
		s.m.Store(key, value)
		return 0 // 更新
	}
	return 1 // 新增
}

// PutIfAbsent 当键不存在时才插入键值对
// 参数 key 键
// 参数 value 值
// 返回值 result 表示操作结果：1表示插入成功，0表示键已存在未插入
func (s *SyncDict) PutIfAbsent(key string, value any) (result int) {
	_, loaded := s.m.LoadOrStore(key, value)
	if loaded {
		return 0 // 键已存在
	}
	return 1 // 插入成功
}

// PutIfExists 当键存在时才更新键值对
// 参数 key 键
// 参数 value 值
// 返回值 result 表示操作结果：1表示更新成功，0表示键不存在未更新
func (s *SyncDict) PutIfExists(key string, value any) (result int) {
	_, ok := s.m.Load(key)
	if ok {
		s.m.Store(key, value)
		return 1 // 更新成功
	}
	return 0 // 键不存在
}

// Remove 根据键删除键值对
// 参数 key 要删除的键
// 返回值 result 表示操作结果：1表示删除成功，0表示键不存在未删除
func (s *SyncDict) Remove(key string) (result int) {
	_, ok := s.m.LoadAndDelete(key)
	if ok {
		return 1 // 删除成功
	}
	return 0 // 键不存在
}

// ForEach 遍历字典中的所有键值对
// 参数 consumer 是处理每个键值对的函数，当consumer返回false时停止遍历
func (s *SyncDict) ForEach(consumer Consumer) {
	s.m.Range(func(key, value any) bool {
		return consumer(key.(string), value)
	})
}

// Keys 返回字典中所有的键
// 返回包含所有键的字符串切片
func (s *SyncDict) Keys() []string {
	keys := make([]string, 0)
	s.m.Range(func(key, value any) bool {
		keys = append(keys, key.(string))
		return true
	})
	return keys
}

// RandomKeys 随机返回字典中的键，可能包含重复键
// 参数 limit 限制返回键的数量
func (s *SyncDict) RandomKeys(limit int) []string {
	keys := s.Keys()
	if len(keys) == 0 {
		return make([]string, limit)
	}

	result := make([]string, limit)
	for i := 0; i < limit; i++ {
		result[i] = keys[rand.Intn(len(keys))]
	}
	return result
}

// RandomDistinctKeys 随机返回字典中的不重复键
// 参数 limit 限制返回键的数量
func (s *SyncDict) RandomDistinctKeys(limit int) []string {
	keys := s.Keys()
	if len(keys) <= limit {
		return keys
	}

	// Fisher-Yates shuffle 算法的一部分，用于随机选择不重复的键
	result := make([]string, limit)
	for i := 0; i < limit; i++ {
		// 从剩余元素中随机选择一个
		j := i + rand.Intn(len(keys)-i)
		// 交换元素
		keys[i], keys[j] = keys[j], keys[i]
		result[i] = keys[i]
	}
	return result
}

// Clear 清空字典中的所有键值对
// 通过创建一个新的空 SyncDict 来实现清空操作
func (s *SyncDict) Clear() {
	*s = *NewSyncDict()
}

// NewSyncDict 创建并返回一个新的 SyncDict 实例
func NewSyncDict() *SyncDict {
	return &SyncDict{}
}