diff --git a/utils/counter/counter.go b/utils/counter/counter.go index eb18d07..ba5d80c 100644 --- a/utils/counter/counter.go +++ b/utils/counter/counter.go @@ -1,6 +1,7 @@ package counter import ( + "fmt" "github.com/kercylan98/minotaur/utils/generic" "sync" "time" @@ -30,15 +31,47 @@ type Counter[K comparable, V generic.Number] struct { lock sync.RWMutex sub bool dr map[K]int64 + drm map[string]int64 c map[K]any } +// AddWithMark 添加计数,根据 mark 判断是否重复计数 +func (slf *Counter[K, V]) AddWithMark(mark, key K, value V, expired time.Duration) { + slf.lock.Lock() + if expired > 0 { + now := time.Now().Unix() + mk := fmt.Sprintf("%v_%v", mark, key) + expiredTime, exist := slf.drm[mk] + if exist { + if expiredTime > now { + slf.lock.Unlock() + return + } + } + slf.drm[mk] = now + int64(expired/time.Second) + } + + v, exist := slf.c[key] + if !exist { + slf.c[key] = value + slf.lock.Unlock() + return + } + if v, ok := v.(V); !ok { + slf.lock.Unlock() + panic("counter value is sub counter") + } else { + slf.c[key] = v + value + } + slf.lock.Unlock() +} + // Add 添加计数 // - 当设置了 expired 时,在 expired 时间内,不会重复计数 // - 需要注意的是,当第一次设置了 expired,第二次未设置时,第二次的计数将生效 func (slf *Counter[K, V]) Add(key K, value V, expired ...time.Duration) { slf.lock.Lock() - if len(expired) > 0 { + if len(expired) > 0 && expired[0] > 0 { now := time.Now().Unix() expiredTime, exist := slf.dr[key] if exist { diff --git a/utils/counter/shadow.go b/utils/counter/shadow.go index e8bc4ec..95ccd58 100644 --- a/utils/counter/shadow.go +++ b/utils/counter/shadow.go @@ -6,25 +6,28 @@ import ( ) func newShadow[K comparable, V generic.Number](counter *Counter[K, V]) *Shadow[K, V] { - counter.lock.Lock() + counter.lock.RLock() shadow := &Shadow[K, V]{ - Sub: counter.sub, - DeduplicationRecord: hash.Copy(counter.dr), - Counter: counter.GetCounts(), + Sub: counter.sub, + DeduplicationRecord: hash.Copy(counter.dr), + DeduplicationRecordMark: hash.Copy(counter.drm), } - for k, c := range counter.GetSubCounters() { + counter.lock.RUnlock() + shadow.Counter = counter.GetCounts() + subs := counter.GetSubCounters() + for k, c := range subs { shadow.SubCounters[k] = newShadow(c) } - counter.lock.Unlock() return shadow } // Shadow 计数器的影子计数器 type Shadow[K comparable, V generic.Number] struct { - Sub bool // 是否为子计数器 - DeduplicationRecord map[K]int64 // 最后一次写入时间 - Counter map[K]V // 计数器 - SubCounters map[K]*Shadow[K, V] // 子计数器 + Sub bool // 是否为子计数器 + DeduplicationRecord map[K]int64 // 最后一次写入时间 + DeduplicationRecordMark map[string]int64 // 最后一次写入时间 + Counter map[K]V // 计数器 + SubCounters map[K]*Shadow[K, V] // 子计数器 } // ToCounter 将影子计数器转换为计数器 @@ -42,6 +45,7 @@ func (slf *Shadow[K, V]) toCounter(parent *Counter[K, V]) *Counter[K, V] { counter.dr = parent.dr } else { counter.dr = hash.Copy(slf.DeduplicationRecord) + counter.drm = hash.Copy(slf.DeduplicationRecordMark) } for k, v := range slf.Counter { counter.c[k] = v