|
|
@@ -1,9 +1,11 @@
|
|
|
package core
|
|
|
|
|
|
import (
|
|
|
+ "container/heap"
|
|
|
"time"
|
|
|
|
|
|
v2hash "github.com/v2ray/v2ray-core/hash"
|
|
|
+ "github.com/v2ray/v2ray-core/log"
|
|
|
)
|
|
|
|
|
|
const (
|
|
|
@@ -19,6 +21,7 @@ type UserSet interface {
|
|
|
type TimedUserSet struct {
|
|
|
validUserIds []ID
|
|
|
userHashes map[string]indexTimePair
|
|
|
+ hash2Remove hashEntrySet
|
|
|
}
|
|
|
|
|
|
type indexTimePair struct {
|
|
|
@@ -31,22 +34,62 @@ type hashEntry struct {
|
|
|
timeSec int64
|
|
|
}
|
|
|
|
|
|
+type hashEntrySet []*hashEntry
|
|
|
+
|
|
|
+func (set hashEntrySet) Len() int {
|
|
|
+ return len(set)
|
|
|
+}
|
|
|
+
|
|
|
+func (set hashEntrySet) Less(i, j int) bool {
|
|
|
+ return set[i].timeSec < set[j].timeSec
|
|
|
+}
|
|
|
+
|
|
|
+func (set hashEntrySet) Swap(i, j int) {
|
|
|
+ tmp := set[i]
|
|
|
+ set[i] = set[j]
|
|
|
+ set[j] = tmp
|
|
|
+}
|
|
|
+
|
|
|
+func (set *hashEntrySet) Push(value interface{}) {
|
|
|
+ entry := value.(*hashEntry)
|
|
|
+ *set = append(*set, entry)
|
|
|
+}
|
|
|
+
|
|
|
+func (set *hashEntrySet) Pop() interface{} {
|
|
|
+ old := *set
|
|
|
+ n := len(old)
|
|
|
+ v := old[n-1]
|
|
|
+ *set = old[:n-1]
|
|
|
+ return v
|
|
|
+}
|
|
|
+
|
|
|
func NewTimedUserSet() UserSet {
|
|
|
vuSet := new(TimedUserSet)
|
|
|
vuSet.validUserIds = make([]ID, 0, 16)
|
|
|
vuSet.userHashes = make(map[string]indexTimePair)
|
|
|
+ vuSet.hash2Remove = make(hashEntrySet, 0, cacheDurationSec*10)
|
|
|
|
|
|
go vuSet.updateUserHash(time.Tick(updateIntervalSec * time.Second))
|
|
|
return vuSet
|
|
|
}
|
|
|
|
|
|
+func (us *TimedUserSet) generateNewHashes(lastSec, nowSec int64, idx int, id ID) {
|
|
|
+ idHash := v2hash.NewTimeHash(v2hash.HMACHash{})
|
|
|
+ for lastSec < nowSec+cacheDurationSec {
|
|
|
+
|
|
|
+ idHash := idHash.Hash(id.Bytes, lastSec)
|
|
|
+ log.Debug("Valid User Hash: %v", idHash)
|
|
|
+ heap.Push(&us.hash2Remove, &hashEntry{string(idHash), lastSec})
|
|
|
+ us.userHashes[string(idHash)] = indexTimePair{idx, lastSec}
|
|
|
+ lastSec++
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
func (us *TimedUserSet) updateUserHash(tick <-chan time.Time) {
|
|
|
now := time.Now().UTC()
|
|
|
- lastSec := now.Unix() - cacheDurationSec
|
|
|
-
|
|
|
- hash2Remove := make(chan hashEntry, cacheDurationSec*3*len(us.validUserIds))
|
|
|
+ lastSec := now.Unix()
|
|
|
lastSec2Remove := now.Unix()
|
|
|
- idHash := v2hash.NewTimeHash(v2hash.HMACHash{})
|
|
|
+
|
|
|
for {
|
|
|
now := <-tick
|
|
|
nowSec := now.UTC().Unix()
|
|
|
@@ -54,26 +97,27 @@ func (us *TimedUserSet) updateUserHash(tick <-chan time.Time) {
|
|
|
remove2Sec := nowSec - cacheDurationSec
|
|
|
if remove2Sec > lastSec2Remove {
|
|
|
for lastSec2Remove+1 < remove2Sec {
|
|
|
- entry := <-hash2Remove
|
|
|
+ front := heap.Pop(&us.hash2Remove)
|
|
|
+ entry := front.(*hashEntry)
|
|
|
lastSec2Remove = entry.timeSec
|
|
|
delete(us.userHashes, entry.hash)
|
|
|
}
|
|
|
}
|
|
|
-
|
|
|
- for lastSec < nowSec+cacheDurationSec {
|
|
|
- for idx, id := range us.validUserIds {
|
|
|
- idHash := idHash.Hash(id.Bytes, lastSec)
|
|
|
- hash2Remove <- hashEntry{string(idHash), lastSec}
|
|
|
- us.userHashes[string(idHash)] = indexTimePair{idx, lastSec}
|
|
|
- }
|
|
|
- lastSec++
|
|
|
+ for idx, id := range us.validUserIds {
|
|
|
+ us.generateNewHashes(lastSec, nowSec, idx, id)
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
func (us *TimedUserSet) AddUser(user User) error {
|
|
|
id := user.Id
|
|
|
+ idx := len(us.validUserIds)
|
|
|
us.validUserIds = append(us.validUserIds, id)
|
|
|
+
|
|
|
+ nowSec := time.Now().UTC().Unix()
|
|
|
+ lastSec := nowSec - cacheDurationSec
|
|
|
+ us.generateNewHashes(lastSec, nowSec, idx, id)
|
|
|
+
|
|
|
return nil
|
|
|
}
|
|
|
|