From ab723fa0b586532e516a2d097d445016cff5f669 Mon Sep 17 00:00:00 2001
From: V2Ray <admin@v2ray.com>
Date: Sat, 19 Sep 2015 21:56:00 +0200
Subject: [PATCH] Generate user hash on demand

---
 userset.go | 70 ++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 57 insertions(+), 13 deletions(-)

diff --git a/userset.go b/userset.go
index fc55e5c8..520105d0 100644
--- a/userset.go
+++ b/userset.go
@@ -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
 }