//ref https://github.com/DeanThompson/syncmap/blob/master/syncmap.go
|
|
package shardmap
|
|
import (
|
"basic.com/pubsub/protomsg.git"
|
"sdkCompare/db"
|
"sync"
|
//"encoding/json"
|
//"fmt"
|
"time"
|
)
|
|
// var commonmux sync.Mutex
|
const (
|
BKDR_SEED = 131 // 31 131 1313 13131 131313 etc...
|
)
|
|
type shardItem struct {
|
sync.RWMutex
|
data map[string]interface{}
|
}
|
|
type ShardMap struct {
|
shardCnt uint8
|
shards []*shardItem
|
lasttime time.Time
|
}
|
|
var Count = make(chan int)
|
|
type wfOp func(a, b []float32) float32
|
|
/**
|
* @param uint8, shardCnt must be pow of two
|
*/
|
func New(shardCnt uint8) *ShardMap {
|
|
s := &ShardMap{
|
shardCnt: shardCnt,
|
shards: make([]*shardItem, shardCnt),
|
lasttime: time.Now(),
|
}
|
|
for i, _ := range s.shards {
|
s.shards[i] = &shardItem{
|
data: make(map[string]interface{}),
|
}
|
}
|
|
return s
|
}
|
|
func (s *ShardMap) Get(key string) (interface{}, bool) {
|
si := s.locate(key)
|
|
si.RLock()
|
value, ok := si.data[key]
|
si.RUnlock()
|
return value, ok
|
}
|
|
func (s *ShardMap) Gettime() time.Time {
|
tm := s.lasttime
|
return tm
|
}
|
|
func (s *ShardMap) Settime() {
|
s.lasttime = time.Now()
|
}
|
|
func (s *ShardMap) Set(key string, value interface{}) {
|
si := s.locate(key)
|
|
si.Lock()
|
si.data[key] = value
|
si.Unlock()
|
}
|
|
func (s *ShardMap) Del(key string) {
|
si := s.locate(key)
|
|
si.Lock()
|
delete(si.data, key)
|
si.Unlock()
|
}
|
|
type kvItem struct {
|
key string
|
value interface{}
|
}
|
|
func (s *ShardMap) Walk(wf wfOp, sourceFea []float32, baseScore float32) (targets []*protomsg.SdkCompareEach) {
|
var wg sync.WaitGroup
|
var lock sync.Mutex
|
|
for _, si := range s.shards {
|
var tempsi *shardItem = si // 保持对原始 shardItem 的指针引用
|
|
// 跳过空分片
|
if len(tempsi.data) == 0 {
|
continue
|
}
|
|
wg.Add(1)
|
go func(st *shardItem, fn wfOp, srcFeat []float32, baseSec float32) {
|
defer wg.Done()
|
st.RLock() // 锁定读取
|
defer st.RUnlock() // 确保读取完毕后解锁
|
|
for _, feature := range st.data {
|
// 读取操作在锁内进行，防止并发冲突
|
if item, ok := feature.(*db.FeatureCacheBase); ok {
|
score := fn(srcFeat, item.FaceFeature)
|
if score > 0 && score >= baseScore {
|
lock.Lock() // 保护目标切片的写入
|
targets = append(targets, &protomsg.SdkCompareEach{
|
Id: item.Id,
|
CompareScore: score,
|
Tableid: item.TableId,
|
})
|
lock.Unlock()
|
}
|
}
|
}
|
}(tempsi, wf, sourceFea, baseScore)
|
}
|
|
wg.Wait()
|
return targets
|
}
|
|
// print all
|
func (s *ShardMap) Printall() (infos []interface{}) {
|
var wg sync.WaitGroup
|
|
for _, si := range s.shards {
|
wg.Add(1)
|
go func(s *shardItem) {
|
defer wg.Done()
|
s.RLock()
|
for _, value := range s.data {
|
infos = append(infos, value)
|
}
|
|
s.RUnlock()
|
}(si)
|
}
|
wg.Wait()
|
return
|
}
|
|
func (s *ShardMap) GetLen() int {
|
var slen int
|
for i := 0; i < int(s.shardCnt); i++ {
|
slen += len(s.shards[i].data)
|
}
|
return slen
|
}
|
|
func (s *ShardMap) locate(key string) *shardItem {
|
i := bkdrHash(key) & uint32(s.shardCnt-1)
|
|
return s.shards[i]
|
}
|
|
// https://www.byvoid.com/blog/string-hash-compare/
|
func bkdrHash(str string) uint32 {
|
var h uint32
|
|
for _, c := range str {
|
h = h*BKDR_SEED + uint32(c)
|
}
|
|
return h
|
}
|
|
func strComp(source []byte, target string) float32 {
|
sourceStr := string(source)
|
if sourceStr == target {
|
return 100
|
} else {
|
return 0
|
}
|
}
|
