1+ package conn
2+
3+ import (
4+ "fmt"
5+ "github.com/stretchr/testify/assert"
6+ "net"
7+ "runtime"
8+ "sync"
9+ "sync/atomic"
10+ "testing"
11+ "time"
12+ )
13+
14+ func TestNewLimiter (t * testing.T ) {
15+ l := NewLimiter (60 , 100 * time .Millisecond , time .Minute )
16+
17+ ip1 := net .IPv4 (1 ,2 ,3 ,4 )
18+ ip2 := net .IPv4 (4 ,3 ,2 ,1 )
19+
20+ start := time .Now ()
21+
22+ // should accept all the queries as long as there are less than 60/100ms
23+ for time .Now ().Sub (start ) < 30 * time .Millisecond {
24+ assert .True (t , l .Accept (ip1 ))
25+ time .Sleep (time .Millisecond )
26+ }
27+ // should end up rejecting queries
28+ last := true
29+ for time .Now ().Sub (start ) < 90 * time .Millisecond {
30+ last = last && l .Accept (ip1 )
31+ }
32+ assert .False (t , last )
33+
34+ time .Sleep (30 * time .Millisecond )
35+ l .ClearOldEntries ()
36+
37+ // after some time, there is more room to get enough queries
38+ start = time .Now ()
39+ for i := 0 ; i < 10 ; i ++ {
40+ assert .True (t , l .Accept (ip1 ))
41+ assert .True (t , l .Accept (ip2 ))
42+ time .Sleep (time .Millisecond )
43+ }
44+ }
45+
46+ func TestTimedLRU (t * testing.T ) {
47+ now := time .Now ()
48+ clock = func () time.Time {
49+ return now
50+ }
51+ tlru := NewTimedLRU [string , string , int ](time .Second )
52+
53+ qu := func (i * string ) int {
54+ * i = * i + "-"
55+ return len (* i )
56+ }
57+ cr := func () (string , int ) {
58+ return "" , 0
59+ }
60+ assert .Equal (t , 0 , tlru .QueryUpdateOrCreate ("hi" , qu , cr ))
61+ assert .Equal (t , 0 , tlru .QueryUpdateOrCreate ("ho" , qu , cr ))
62+ now = now .Add (800 * time .Millisecond )
63+ assert .Equal (t , 0 , tlru .QueryUpdateOrCreate ("foo" , qu , cr ))
64+ assert .Equal (t , 1 , tlru .QueryUpdateOrCreate ("ho" , qu , cr ))
65+ now = now .Add (800 * time .Millisecond )
66+ tlru .RemoveOldItems ()
67+ // "hi" has been deleted then it returns a fresh one
68+ assert .Equal (t , 0 , tlru .QueryUpdateOrCreate ("hi" , qu , cr ))
69+ assert .Equal (t , 2 , tlru .QueryUpdateOrCreate ("ho" , qu , cr ))
70+ assert .Equal (t , 1 , tlru .QueryUpdateOrCreate ("foo" , qu , cr ))
71+ now = now .Add (10000 * time .Millisecond )
72+ tlru .RemoveOldItems ()
73+ // all entries have been deleted then it returns fresh instances
74+ assert .Equal (t , 0 , tlru .QueryUpdateOrCreate ("hi" , qu , cr ))
75+ assert .Equal (t , 0 , tlru .QueryUpdateOrCreate ("ho" , qu , cr ))
76+ assert .Equal (t , 0 , tlru .QueryUpdateOrCreate ("foo" , qu , cr ))
77+ }
78+
79+ func TestTimedLRU_ConcurrencyRaces (t * testing.T ) {
80+ timeBias := int64 (0 )
81+ clock = func () time.Time {
82+ return time .Now ().Add (time .Duration (atomic .LoadInt64 (& timeBias )))
83+ }
84+ tlru := NewTimedLRU [string , int64 , bool ](time .Second )
85+ wg := sync.WaitGroup {}
86+ wg .Add (4 )
87+ for i := 0 ; i < 4 ; i ++ {
88+ thread := i
89+ go func () {
90+ start := time .Now ()
91+ defer wg .Done ()
92+ for time .Now ().Sub (start ) < 100 * time .Millisecond {
93+ clk := clock ()
94+ key := fmt .Sprint (clk .UnixMilli ())
95+ tlru .QueryUpdateOrCreate (key , func (v * int64 ) bool {
96+ atomic .AddInt64 (v , 1 )
97+ return true
98+ }, func () (int64 , bool ) {
99+ return 0 , true
100+ })
101+ tlru .RemoveOldItems ()
102+ runtime .Gosched ()
103+ }
104+ if thread == 0 {
105+ atomic .StoreInt64 (& timeBias , int64 (950 * time .Millisecond ))
106+ tlru .RemoveOldItems ()
107+ } else {
108+ for time .Now ().Sub (start ) < 200 * time .Millisecond {
109+ clk := clock ()
110+ key := fmt .Sprint (clk .UnixMilli ())
111+ tlru .QueryUpdateOrCreate (key , func (v * int64 ) bool {
112+ atomic .AddInt64 (v , 1 )
113+ return true
114+ }, func () (int64 , bool ) {
115+ return 0 , true
116+ })
117+ tlru .RemoveOldItems ()
118+ runtime .Gosched ()
119+ }
120+ }
121+ }()
122+ }
123+ wg .Wait ()
124+ fmt .Println ("hoe!" )
125+ }
0 commit comments