mirror of
https://github.com/v2fly/v2ray-core.git
synced 2024-11-04 09:17:32 -05:00
rewrite id hash queue
This commit is contained in:
v2ray
parent
8f1c4b7f71
commit
ea080b2ebf
24
common/collect/sized_queue.go
Normal file
24
common/collect/sized_queue.go
Normal file
@ -0,0 +1,24 @@
|
||||
package collect
|
||||
|
||||
type SizedQueue struct {
|
||||
elements []interface{}
|
||||
nextPos int
|
||||
}
|
||||
|
||||
func NewSizedQueue(size int) *SizedQueue {
|
||||
return &SizedQueue{
|
||||
elements: make([]interface{}, size),
|
||||
nextPos: 0,
|
||||
}
|
||||
}
|
||||
|
||||
// Put puts a new element into the queue and pop out the first element if queue is full.
|
||||
func (this *SizedQueue) Put(element interface{}) interface{} {
|
||||
res := this.elements[this.nextPos]
|
||||
this.elements[this.nextPos] = element
|
||||
this.nextPos++
|
||||
if this.nextPos == len(this.elements) {
|
||||
this.nextPos = 0
|
||||
}
|
||||
return res
|
||||
}
|
||||
@ -1,105 +0,0 @@
|
||||
package collect
|
||||
|
||||
import (
|
||||
"container/heap"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type timedQueueEntry struct {
|
||||
timeSec int64
|
||||
value interface{}
|
||||
}
|
||||
|
||||
type timedQueueImpl []*timedQueueEntry
|
||||
|
||||
func (queue timedQueueImpl) Len() int {
|
||||
return len(queue)
|
||||
}
|
||||
|
||||
func (queue timedQueueImpl) Less(i, j int) bool {
|
||||
return queue[i].timeSec < queue[j].timeSec
|
||||
}
|
||||
|
||||
func (queue timedQueueImpl) Swap(i, j int) {
|
||||
queue[i], queue[j] = queue[j], queue[i]
|
||||
}
|
||||
|
||||
func (queue *timedQueueImpl) Push(value interface{}) {
|
||||
entry := value.(*timedQueueEntry)
|
||||
*queue = append(*queue, entry)
|
||||
}
|
||||
|
||||
func (queue *timedQueueImpl) Pop() interface{} {
|
||||
old := *queue
|
||||
n := len(old)
|
||||
v := old[n-1]
|
||||
old[n-1] = nil
|
||||
*queue = old[:n-1]
|
||||
return v
|
||||
}
|
||||
|
||||
// TimedQueue is a priority queue that entries with oldest timestamp get removed first.
|
||||
type TimedQueue struct {
|
||||
queue timedQueueImpl
|
||||
access sync.Mutex
|
||||
removedCallback func(interface{})
|
||||
}
|
||||
|
||||
func NewTimedQueue(updateInterval int, removedCallback func(interface{})) *TimedQueue {
|
||||
queue := &TimedQueue{
|
||||
queue: make([]*timedQueueEntry, 0, 256),
|
||||
removedCallback: removedCallback,
|
||||
access: sync.Mutex{},
|
||||
}
|
||||
go queue.cleanup(time.Tick(time.Duration(updateInterval) * time.Second))
|
||||
return queue
|
||||
}
|
||||
|
||||
func (queue *TimedQueue) Add(value interface{}, time2Remove int64) {
|
||||
newEntry := &timedQueueEntry{
|
||||
timeSec: time2Remove,
|
||||
value: value,
|
||||
}
|
||||
var removedEntry *timedQueueEntry
|
||||
queue.access.Lock()
|
||||
nowSec := time.Now().Unix()
|
||||
if queue.queue.Len() > 0 && queue.queue[0].timeSec < nowSec {
|
||||
removedEntry = queue.queue[0]
|
||||
queue.queue[0] = newEntry
|
||||
heap.Fix(&queue.queue, 0)
|
||||
} else {
|
||||
heap.Push(&queue.queue, newEntry)
|
||||
}
|
||||
queue.access.Unlock()
|
||||
if removedEntry != nil {
|
||||
queue.removedCallback(removedEntry.value)
|
||||
}
|
||||
}
|
||||
|
||||
func (queue *TimedQueue) cleanup(tick <-chan time.Time) {
|
||||
for now := range tick {
|
||||
nowSec := now.Unix()
|
||||
removedEntries := make([]*timedQueueEntry, 0, 128)
|
||||
queue.access.Lock()
|
||||
changed := false
|
||||
for i := 0; i < queue.queue.Len(); i++ {
|
||||
entry := queue.queue[i]
|
||||
if entry.timeSec < nowSec {
|
||||
removedEntries = append(removedEntries, entry)
|
||||
queue.queue.Swap(i, queue.queue.Len()-1)
|
||||
queue.queue.Pop()
|
||||
changed = true
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
if changed {
|
||||
heap.Init(&queue.queue)
|
||||
}
|
||||
queue.access.Unlock()
|
||||
for _, entry := range removedEntries {
|
||||
queue.removedCallback(entry.value)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,56 +0,0 @@
|
||||
package collect
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
v2testing "github.com/v2ray/v2ray-core/testing"
|
||||
"github.com/v2ray/v2ray-core/testing/assert"
|
||||
)
|
||||
|
||||
func TestTimedQueue(t *testing.T) {
|
||||
v2testing.Current(t)
|
||||
|
||||
removed := make(map[string]bool)
|
||||
|
||||
nowSec := time.Now().Unix()
|
||||
q := NewTimedQueue(2, func(v interface{}) {
|
||||
removed[v.(string)] = true
|
||||
})
|
||||
|
||||
q.Add("Value1", nowSec)
|
||||
q.Add("Value2", nowSec+5)
|
||||
|
||||
v1, ok := removed["Value1"]
|
||||
assert.Bool(ok).IsFalse()
|
||||
|
||||
v2, ok := removed["Value2"]
|
||||
assert.Bool(ok).IsFalse()
|
||||
|
||||
tick := time.Tick(4 * time.Second)
|
||||
<-tick
|
||||
|
||||
v1, ok = removed["Value1"]
|
||||
assert.Bool(ok).IsTrue()
|
||||
assert.Bool(v1).IsTrue()
|
||||
removed["Value1"] = false
|
||||
|
||||
v2, ok = removed["Value2"]
|
||||
assert.Bool(ok).IsFalse()
|
||||
|
||||
<-tick
|
||||
v2, ok = removed["Value2"]
|
||||
assert.Bool(ok).IsTrue()
|
||||
assert.Bool(v2).IsTrue()
|
||||
removed["Value2"] = false
|
||||
|
||||
<-tick
|
||||
assert.Bool(removed["Values"]).IsFalse()
|
||||
|
||||
q.Add("Value1", time.Now().Unix()+10)
|
||||
|
||||
<-tick
|
||||
v1, ok = removed["Value1"]
|
||||
assert.Bool(ok).IsTrue()
|
||||
assert.Bool(v1).IsFalse()
|
||||
}
|
||||
@ -13,16 +13,23 @@ const (
|
||||
cacheDurationSec = 120
|
||||
)
|
||||
|
||||
type idEntry struct {
|
||||
id *vmess.ID
|
||||
userIdx int
|
||||
lastSec int64
|
||||
hashes *collect.SizedQueue
|
||||
}
|
||||
|
||||
type UserSet interface {
|
||||
AddUser(user vmess.User) error
|
||||
GetUser(timeHash []byte) (vmess.User, int64, bool)
|
||||
}
|
||||
|
||||
type TimedUserSet struct {
|
||||
validUsers []vmess.User
|
||||
userHash map[string]indexTimePair
|
||||
userHashDeleteQueue *collect.TimedQueue
|
||||
access sync.RWMutex
|
||||
validUsers []vmess.User
|
||||
userHash map[string]indexTimePair
|
||||
ids []*idEntry
|
||||
access sync.RWMutex
|
||||
}
|
||||
|
||||
type indexTimePair struct {
|
||||
@ -35,55 +42,63 @@ func NewTimedUserSet() UserSet {
|
||||
validUsers: make([]vmess.User, 0, 16),
|
||||
userHash: make(map[string]indexTimePair, 512),
|
||||
access: sync.RWMutex{},
|
||||
ids: make([]*idEntry, 0, 512),
|
||||
}
|
||||
tus.userHashDeleteQueue = collect.NewTimedQueue(updateIntervalSec, tus.removeEntry)
|
||||
go tus.updateUserHash(time.Tick(updateIntervalSec * time.Second))
|
||||
return tus
|
||||
}
|
||||
|
||||
func (us *TimedUserSet) removeEntry(entry interface{}) {
|
||||
us.access.Lock()
|
||||
delete(us.userHash, entry.(string))
|
||||
us.access.Unlock()
|
||||
}
|
||||
|
||||
func (us *TimedUserSet) generateNewHashes(lastSec, nowSec int64, idx int, id *vmess.ID) {
|
||||
func (us *TimedUserSet) generateNewHashes(nowSec int64, idx int, entry *idEntry) {
|
||||
idHash := NewTimeHash(HMACHash{})
|
||||
for lastSec < nowSec {
|
||||
idHash := idHash.Hash(id.Bytes(), lastSec)
|
||||
for entry.lastSec <= nowSec {
|
||||
idHashSlice := idHash.Hash(entry.id.Bytes(), entry.lastSec)
|
||||
hashValue := string(idHashSlice)
|
||||
us.access.Lock()
|
||||
us.userHash[string(idHash)] = indexTimePair{idx, lastSec}
|
||||
us.userHash[hashValue] = indexTimePair{idx, entry.lastSec}
|
||||
us.access.Unlock()
|
||||
us.userHashDeleteQueue.Add(string(idHash), lastSec+2*cacheDurationSec)
|
||||
lastSec++
|
||||
|
||||
hash2Remove := entry.hashes.Put(hashValue)
|
||||
if hash2Remove != nil {
|
||||
us.access.Lock()
|
||||
delete(us.userHash, hash2Remove.(string))
|
||||
us.access.Unlock()
|
||||
}
|
||||
entry.lastSec++
|
||||
}
|
||||
}
|
||||
|
||||
func (us *TimedUserSet) updateUserHash(tick <-chan time.Time) {
|
||||
lastSec := time.Now().Unix() - cacheDurationSec
|
||||
|
||||
for now := range tick {
|
||||
nowSec := now.Unix() + cacheDurationSec
|
||||
for idx, user := range us.validUsers {
|
||||
us.generateNewHashes(lastSec, nowSec, idx, user.ID())
|
||||
for _, alterid := range user.AlterIDs() {
|
||||
us.generateNewHashes(lastSec, nowSec, idx, alterid)
|
||||
}
|
||||
for _, entry := range us.ids {
|
||||
us.generateNewHashes(nowSec, entry.userIdx, entry)
|
||||
}
|
||||
lastSec = nowSec
|
||||
}
|
||||
}
|
||||
|
||||
func (us *TimedUserSet) AddUser(user vmess.User) error {
|
||||
id := user.ID()
|
||||
idx := len(us.validUsers)
|
||||
us.validUsers = append(us.validUsers, user)
|
||||
|
||||
nowSec := time.Now().Unix()
|
||||
lastSec := nowSec - cacheDurationSec
|
||||
us.generateNewHashes(lastSec, nowSec+cacheDurationSec, idx, id)
|
||||
|
||||
entry := &idEntry{
|
||||
id: user.ID(),
|
||||
userIdx: idx,
|
||||
lastSec: nowSec - cacheDurationSec,
|
||||
hashes: collect.NewSizedQueue(2*cacheDurationSec + 1),
|
||||
}
|
||||
us.generateNewHashes(nowSec+cacheDurationSec, idx, entry)
|
||||
us.ids = append(us.ids, entry)
|
||||
for _, alterid := range user.AlterIDs() {
|
||||
us.generateNewHashes(lastSec, nowSec+cacheDurationSec, idx, alterid)
|
||||
entry := &idEntry{
|
||||
id: alterid,
|
||||
userIdx: idx,
|
||||
lastSec: nowSec - cacheDurationSec,
|
||||
hashes: collect.NewSizedQueue(2*cacheDurationSec + 1),
|
||||
}
|
||||
us.generateNewHashes(nowSec+cacheDurationSec, idx, entry)
|
||||
us.ids = append(us.ids, entry)
|
||||
}
|
||||
|
||||
return nil
|
||||
|
||||
Loading…
Reference in New Issue
Block a user