zeromicro/go-zero
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

initial import

Browse Source
This commit is contained in:
kevin
2020-07-26 17:09:05 +08:00
commit 7e3a369a8f
647 changed files with 54754 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

275
core/collection/cache.go Normal file
View File

@@ -0,0 +1,275 @@
package collection
import (
"container/list"
"sync"
"sync/atomic"
"time"
"zero/core/logx"
"zero/core/mathx"
"zero/core/syncx"
)
const (
defaultCacheName = "proc"
slots = 300
statInterval = time.Minute
// make the expiry unstable to avoid lots of cached items expire at the same time
// make the unstable expiry to be [0.95, 1.05] * seconds
expiryDeviation = 0.05
)
var emptyLruCache = emptyLru{}
type (
CacheOption func(cache *Cache)
Cache struct {
name string
lock sync.Mutex
data map[string]interface{}
evicts *list.List
expire time.Duration
timingWheel *TimingWheel
lruCache lru
barrier syncx.SharedCalls
unstableExpiry mathx.Unstable
stats *cacheStat
}
)
func NewCache(expire time.Duration, opts ...CacheOption) (*Cache, error) {
cache := &Cache{
data: make(map[string]interface{}),
expire: expire,
lruCache: emptyLruCache,
barrier: syncx.NewSharedCalls(),
unstableExpiry: mathx.NewUnstable(expiryDeviation),
}
for _, opt := range opts {
opt(cache)
}
if len(cache.name) == 0 {
cache.name = defaultCacheName
}
cache.stats = newCacheStat(cache.name, cache.size)
timingWheel, err := NewTimingWheel(time.Second, slots, func(k, v interface{}) {
key, ok := k.(string)
if !ok {
return
}
cache.Del(key)
})
if err != nil {
return nil, err
}
cache.timingWheel = timingWheel
return cache, nil
}
func (c *Cache) Del(key string) {
c.lock.Lock()
delete(c.data, key)
c.lruCache.remove(key)
c.lock.Unlock()
c.timingWheel.RemoveTimer(key)
}
func (c *Cache) Get(key string) (interface{}, bool) {
c.lock.Lock()
value, ok := c.data[key]
if ok {
c.lruCache.add(key)
}
c.lock.Unlock()
if ok {
c.stats.IncrementHit()
} else {
c.stats.IncrementMiss()
}
return value, ok
}
func (c *Cache) Set(key string, value interface{}) {
c.lock.Lock()
_, ok := c.data[key]
c.data[key] = value
c.lruCache.add(key)
c.lock.Unlock()
expiry := c.unstableExpiry.AroundDuration(c.expire)
if ok {
c.timingWheel.MoveTimer(key, expiry)
} else {
c.timingWheel.SetTimer(key, value, expiry)
}
}
func (c *Cache) Take(key string, fetch func() (interface{}, error)) (interface{}, error) {
val, fresh, err := c.barrier.DoEx(key, func() (interface{}, error) {
v, e := fetch()
if e != nil {
return nil, e
}
c.Set(key, v)
return v, nil
})
if err != nil {
return nil, err
}
if fresh {
c.stats.IncrementMiss()
return val, nil
} else {
// got the result from previous ongoing query
c.stats.IncrementHit()
}
return val, nil
}
func (c *Cache) onEvict(key string) {
// already locked
delete(c.data, key)
c.timingWheel.RemoveTimer(key)
}
func (c *Cache) size() int {
c.lock.Lock()
defer c.lock.Unlock()
return len(c.data)
}
func WithLimit(limit int) CacheOption {
return func(cache *Cache) {
if limit > 0 {
cache.lruCache = newKeyLru(limit, cache.onEvict)
}
}
}
func WithName(name string) CacheOption {
return func(cache *Cache) {
cache.name = name
}
}
type (
lru interface {
add(key string)
remove(key string)
}
emptyLru struct{}
keyLru struct {
limit int
evicts *list.List
elements map[string]*list.Element
onEvict func(key string)
}
)
func (elru emptyLru) add(string) {
}
func (elru emptyLru) remove(string) {
}
func newKeyLru(limit int, onEvict func(key string)) *keyLru {
return &keyLru{
limit: limit,
evicts: list.New(),
elements: make(map[string]*list.Element),
onEvict: onEvict,
}
}
func (klru *keyLru) add(key string) {
if elem, ok := klru.elements[key]; ok {
klru.evicts.MoveToFront(elem)
return
}
// Add new item
elem := klru.evicts.PushFront(key)
klru.elements[key] = elem
// Verify size not exceeded
if klru.evicts.Len() > klru.limit {
klru.removeOldest()
}
}
func (klru *keyLru) remove(key string) {
if elem, ok := klru.elements[key]; ok {
klru.removeElement(elem)
}
}
func (klru *keyLru) removeOldest() {
elem := klru.evicts.Back()
if elem != nil {
klru.removeElement(elem)
}
}
func (klru *keyLru) removeElement(e *list.Element) {
klru.evicts.Remove(e)
key := e.Value.(string)
delete(klru.elements, key)
klru.onEvict(key)
}
type cacheStat struct {
name string
hit uint64
miss uint64
sizeCallback func() int
}
func newCacheStat(name string, sizeCallback func() int) *cacheStat {
st := &cacheStat{
name: name,
sizeCallback: sizeCallback,
}
go st.statLoop()
return st
}
func (cs *cacheStat) IncrementHit() {
atomic.AddUint64(&cs.hit, 1)
}
func (cs *cacheStat) IncrementMiss() {
atomic.AddUint64(&cs.miss, 1)
}
func (cs *cacheStat) statLoop() {
ticker := time.NewTicker(statInterval)
defer ticker.Stop()
for {
select {
case <-ticker.C:
hit := atomic.SwapUint64(&cs.hit, 0)
miss := atomic.SwapUint64(&cs.miss, 0)
total := hit + miss
if total == 0 {
continue
}
percent := 100 * float32(hit) / float32(total)
logx.Statf("cache(%s) - qpm: %d, hit_ratio: %.1f%%, elements: %d, hit: %d, miss: %d",
cs.name, total, percent, cs.sizeCallback(), hit, miss)
}
}
}

139
core/collection/cache_test.go Normal file
View File

@@ -0,0 +1,139 @@
package collection
import (
"strconv"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestCacheSet(t *testing.T) {
cache, err := NewCache(time.Second*2, WithName("any"))
assert.Nil(t, err)
cache.Set("first", "first element")
cache.Set("second", "second element")
value, ok := cache.Get("first")
assert.True(t, ok)
assert.Equal(t, "first element", value)
value, ok = cache.Get("second")
assert.True(t, ok)
assert.Equal(t, "second element", value)
}
func TestCacheDel(t *testing.T) {
cache, err := NewCache(time.Second * 2)
assert.Nil(t, err)
cache.Set("first", "first element")
cache.Set("second", "second element")
cache.Del("first")
_, ok := cache.Get("first")
assert.False(t, ok)
value, ok := cache.Get("second")
assert.True(t, ok)
assert.Equal(t, "second element", value)
}
func TestCacheTake(t *testing.T) {
cache, err := NewCache(time.Second * 2)
assert.Nil(t, err)
var count int32
var wg sync.WaitGroup
for i := 0; i < 100; i++ {
wg.Add(1)
go func() {
cache.Take("first", func() (interface{}, error) {
atomic.AddInt32(&count, 1)
time.Sleep(time.Millisecond * 100)
return "first element", nil
})
wg.Done()
}()
}
wg.Wait()
assert.Equal(t, 1, cache.size())
assert.Equal(t, int32(1), atomic.LoadInt32(&count))
}
func TestCacheWithLruEvicts(t *testing.T) {
cache, err := NewCache(time.Minute, WithLimit(3))
assert.Nil(t, err)
cache.Set("first", "first element")
cache.Set("second", "second element")
cache.Set("third", "third element")
cache.Set("fourth", "fourth element")
value, ok := cache.Get("first")
assert.False(t, ok)
value, ok = cache.Get("second")
assert.True(t, ok)
assert.Equal(t, "second element", value)
value, ok = cache.Get("third")
assert.True(t, ok)
assert.Equal(t, "third element", value)
value, ok = cache.Get("fourth")
assert.True(t, ok)
assert.Equal(t, "fourth element", value)
}
func TestCacheWithLruEvicted(t *testing.T) {
cache, err := NewCache(time.Minute, WithLimit(3))
assert.Nil(t, err)
cache.Set("first", "first element")
cache.Set("second", "second element")
cache.Set("third", "third element")
cache.Set("fourth", "fourth element")
value, ok := cache.Get("first")
assert.False(t, ok)
value, ok = cache.Get("second")
assert.True(t, ok)
assert.Equal(t, "second element", value)
cache.Set("fifth", "fifth element")
cache.Set("sixth", "sixth element")
_, ok = cache.Get("third")
assert.False(t, ok)
_, ok = cache.Get("fourth")
assert.False(t, ok)
value, ok = cache.Get("second")
assert.True(t, ok)
assert.Equal(t, "second element", value)
}
func BenchmarkCache(b *testing.B) {
cache, err := NewCache(time.Second*5, WithLimit(100000))
if err != nil {
b.Fatal(err)
}
for i := 0; i < 10000; i++ {
for j := 0; j < 10; j++ {
index := strconv.Itoa(i*10000 + j)
cache.Set("key:"+index, "value:"+index)
}
}
time.Sleep(time.Second * 5)
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
for i := 0; i < b.N; i++ {
index := strconv.Itoa(i % 10000)
cache.Get("key:" + index)
if i%100 == 0 {
cache.Set("key1:"+index, "value1:"+index)
}
}
}
})
}

60
core/collection/fifo.go Normal file
View File

@@ -0,0 +1,60 @@
package collection
import "sync"
type Queue struct {
lock sync.Mutex
elements []interface{}
size int
head int
tail int
count int
}
func NewQueue(size int) *Queue {
return &Queue{
elements: make([]interface{}, size),
size: size,
}
}
func (q *Queue) Empty() bool {
q.lock.Lock()
empty := q.count == 0
q.lock.Unlock()
return empty
}
func (q *Queue) Put(element interface{}) {
q.lock.Lock()
defer q.lock.Unlock()
if q.head == q.tail && q.count > 0 {
nodes := make([]interface{}, len(q.elements)+q.size)
copy(nodes, q.elements[q.head:])
copy(nodes[len(q.elements)-q.head:], q.elements[:q.head])
q.head = 0
q.tail = len(q.elements)
q.elements = nodes
}
q.elements[q.tail] = element
q.tail = (q.tail + 1) % len(q.elements)
q.count++
}
func (q *Queue) Take() (interface{}, bool) {
q.lock.Lock()
defer q.lock.Unlock()
if q.count == 0 {
return nil, false
}
element := q.elements[q.head]
q.head = (q.head + 1) % len(q.elements)
q.count--
return element, true
}

63
core/collection/fifo_test.go Normal file
View File

@@ -0,0 +1,63 @@
package collection
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestFifo(t *testing.T) {
elements := [][]byte{
[]byte("hello"),
[]byte("world"),
[]byte("again"),
}
queue := NewQueue(8)
for i := range elements {
queue.Put(elements[i])
}
for _, element := range elements {
body, ok := queue.Take()
assert.True(t, ok)
assert.Equal(t, string(element), string(body.([]byte)))
}
}
func TestTakeTooMany(t *testing.T) {
elements := [][]byte{
[]byte("hello"),
[]byte("world"),
[]byte("again"),
}
queue := NewQueue(8)
for i := range elements {
queue.Put(elements[i])
}
for range elements {
queue.Take()
}
assert.True(t, queue.Empty())
_, ok := queue.Take()
assert.False(t, ok)
}
func TestPutMore(t *testing.T) {
elements := [][]byte{
[]byte("hello"),
[]byte("world"),
[]byte("again"),
}
queue := NewQueue(2)
for i := range elements {
queue.Put(elements[i])
}
for _, element := range elements {
body, ok := queue.Take()
assert.True(t, ok)
assert.Equal(t, string(element), string(body.([]byte)))
}
}

35
core/collection/ring.go Normal file
View File

@@ -0,0 +1,35 @@
package collection
type Ring struct {
elements []interface{}
index int
}
func NewRing(n int) *Ring {
return &Ring{
elements: make([]interface{}, n),
}
}
func (r *Ring) Add(v interface{}) {
r.elements[r.index%len(r.elements)] = v
r.index++
}
func (r *Ring) Take() []interface{} {
var size int
var start int
if r.index > len(r.elements) {
size = len(r.elements)
start = r.index % len(r.elements)
} else {
size = r.index
}
elements := make([]interface{}, size)
for i := 0; i < size; i++ {
elements[i] = r.elements[(start+i)%len(r.elements)]
}
return elements
}

25
core/collection/ring_test.go Normal file
View File

@@ -0,0 +1,25 @@
package collection
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestRingLess(t *testing.T) {
ring := NewRing(5)
for i := 0; i < 3; i++ {
ring.Add(i)
}
elements := ring.Take()
assert.ElementsMatch(t, []interface{}{0, 1, 2}, elements)
}
func TestRingMore(t *testing.T) {
ring := NewRing(5)
for i := 0; i < 11; i++ {
ring.Add(i)
}
elements := ring.Take()
assert.ElementsMatch(t, []interface{}{6, 7, 8, 9, 10}, elements)
}

145
core/collection/rollingwindow.go Normal file
View File

@@ -0,0 +1,145 @@
package collection
import (
"sync"
"time"
"zero/core/timex"
)
type (
RollingWindowOption func(rollingWindow *RollingWindow)
RollingWindow struct {
lock sync.RWMutex
size int
win *window
interval time.Duration
offset int
ignoreCurrent bool
lastTime time.Duration
}
)
func NewRollingWindow(size int, interval time.Duration, opts ...RollingWindowOption) *RollingWindow {
w := &RollingWindow{
size: size,
win: newWindow(size),
interval: interval,
lastTime: timex.Now(),
}
for _, opt := range opts {
opt(w)
}
return w
}
func (rw *RollingWindow) Add(v float64) {
rw.lock.Lock()
defer rw.lock.Unlock()
rw.updateOffset()
rw.win.add(rw.offset, v)
}
func (rw *RollingWindow) Reduce(fn func(b *Bucket)) {
rw.lock.RLock()
defer rw.lock.RUnlock()
var diff int
span := rw.span()
// ignore current bucket, because of partial data
if span == 0 && rw.ignoreCurrent {
diff = rw.size - 1
} else {
diff = rw.size - span
}
if diff > 0 {
offset := (rw.offset + span + 1) % rw.size
rw.win.reduce(offset, diff, fn)
}
}
func (rw *RollingWindow) span() int {
offset := int(timex.Since(rw.lastTime) / rw.interval)
if 0 <= offset && offset < rw.size {
return offset
} else {
return rw.size
}
}
func (rw *RollingWindow) updateOffset() {
span := rw.span()
if span > 0 {
offset := rw.offset
// reset expired buckets
start := offset + 1
steps := start + span
var remainder int
if steps > rw.size {
remainder = steps - rw.size
steps = rw.size
}
for i := start; i < steps; i++ {
rw.win.resetBucket(i)
offset = i
}
for i := 0; i < remainder; i++ {
rw.win.resetBucket(i)
offset = i
}
rw.offset = offset
rw.lastTime = timex.Now()
}
}
type Bucket struct {
Sum float64
Count int64
}
func (b *Bucket) add(v float64) {
b.Sum += v
b.Count++
}
func (b *Bucket) reset() {
b.Sum = 0
b.Count = 0
}
type window struct {
buckets []*Bucket
size int
}
func newWindow(size int) *window {
var buckets []*Bucket
for i := 0; i < size; i++ {
buckets = append(buckets, new(Bucket))
}
return &window{
buckets: buckets,
size: size,
}
}
func (w *window) add(offset int, v float64) {
w.buckets[offset%w.size].add(v)
}
func (w *window) reduce(start, count int, fn func(b *Bucket)) {
for i := 0; i < count; i++ {
fn(w.buckets[(start+i)%len(w.buckets)])
}
}
func (w *window) resetBucket(offset int) {
w.buckets[offset].reset()
}
func IgnoreCurrentBucket() RollingWindowOption {
return func(w *RollingWindow) {
w.ignoreCurrent = true
}
}

133
core/collection/rollingwindow_test.go Normal file
View File

@@ -0,0 +1,133 @@
package collection
import (
"math/rand"
"testing"
"time"
"zero/core/stringx"
"github.com/stretchr/testify/assert"
)
const duration = time.Millisecond * 50
func TestRollingWindowAdd(t *testing.T) {
const size = 3
r := NewRollingWindow(size, duration)
listBuckets := func() []float64 {
var buckets []float64
r.Reduce(func(b *Bucket) {
buckets = append(buckets, b.Sum)
})
return buckets
}
assert.Equal(t, []float64{0, 0, 0}, listBuckets())
r.Add(1)
assert.Equal(t, []float64{0, 0, 1}, listBuckets())
elapse()
r.Add(2)
r.Add(3)
assert.Equal(t, []float64{0, 1, 5}, listBuckets())
elapse()
r.Add(4)
r.Add(5)
r.Add(6)
assert.Equal(t, []float64{1, 5, 15}, listBuckets())
elapse()
r.Add(7)
assert.Equal(t, []float64{5, 15, 7}, listBuckets())
}
func TestRollingWindowReset(t *testing.T) {
const size = 3
r := NewRollingWindow(size, duration, IgnoreCurrentBucket())
listBuckets := func() []float64 {
var buckets []float64
r.Reduce(func(b *Bucket) {
buckets = append(buckets, b.Sum)
})
return buckets
}
r.Add(1)
elapse()
assert.Equal(t, []float64{0, 1}, listBuckets())
elapse()
assert.Equal(t, []float64{1}, listBuckets())
elapse()
assert.Nil(t, listBuckets())
// cross window
r.Add(1)
time.Sleep(duration * 10)
assert.Nil(t, listBuckets())
}
func TestRollingWindowReduce(t *testing.T) {
const size = 4
tests := []struct {
win *RollingWindow
expect float64
}{
{
win: NewRollingWindow(size, duration),
expect: 10,
},
{
win: NewRollingWindow(size, duration, IgnoreCurrentBucket()),
expect: 4,
},
}
for _, test := range tests {
t.Run(stringx.Rand(), func(t *testing.T) {
r := test.win
for x := 0; x < size; x = x + 1 {
for i := 0; i <= x; i++ {
r.Add(float64(i))
}
if x < size-1 {
elapse()
}
}
var result float64
r.Reduce(func(b *Bucket) {
result += b.Sum
})
assert.Equal(t, test.expect, result)
})
}
}
func TestRollingWindowDataRace(t *testing.T) {
const size = 3
r := NewRollingWindow(size, duration)
var stop = make(chan bool)
go func() {
for {
select {
case <-stop:
return
default:
r.Add(float64(rand.Int63()))
time.Sleep(duration / 2)
}
}
}()
go func() {
for {
select {
case <-stop:
return
default:
r.Reduce(func(b *Bucket) {})
}
}
}()
time.Sleep(duration * 5)
close(stop)
}
func elapse() {
time.Sleep(duration)
}

91
core/collection/safemap.go Normal file
View File

@@ -0,0 +1,91 @@
package collection
import "sync"
const (
copyThreshold = 1000
maxDeletion = 10000
)
// SafeMap provides a map alternative to avoid memory leak.
// This implementation is not needed until issue below fixed.
// https://github.com/golang/go/issues/20135
type SafeMap struct {
lock sync.RWMutex
deletionOld int
deletionNew int
dirtyOld map[interface{}]interface{}
dirtyNew map[interface{}]interface{}
}
func NewSafeMap() *SafeMap {
return &SafeMap{
dirtyOld: make(map[interface{}]interface{}),
dirtyNew: make(map[interface{}]interface{}),
}
}
func (m *SafeMap) Del(key interface{}) {
m.lock.Lock()
if _, ok := m.dirtyOld[key]; ok {
delete(m.dirtyOld, key)
m.deletionOld++
} else if _, ok := m.dirtyNew[key]; ok {
delete(m.dirtyNew, key)
m.deletionNew++
}
if m.deletionOld >= maxDeletion && len(m.dirtyOld) < copyThreshold {
for k, v := range m.dirtyOld {
m.dirtyNew[k] = v
}
m.dirtyOld = m.dirtyNew
m.deletionOld = m.deletionNew
m.dirtyNew = make(map[interface{}]interface{})
m.deletionNew = 0
}
if m.deletionNew >= maxDeletion && len(m.dirtyNew) < copyThreshold {
for k, v := range m.dirtyNew {
m.dirtyOld[k] = v
}
m.dirtyNew = make(map[interface{}]interface{})
m.deletionNew = 0
}
m.lock.Unlock()
}
func (m *SafeMap) Get(key interface{}) (interface{}, bool) {
m.lock.RLock()
defer m.lock.RUnlock()
if val, ok := m.dirtyOld[key]; ok {
return val, true
} else {
val, ok := m.dirtyNew[key]
return val, ok
}
}
func (m *SafeMap) Set(key, value interface{}) {
m.lock.Lock()
if m.deletionOld <= maxDeletion {
if _, ok := m.dirtyNew[key]; ok {
delete(m.dirtyNew, key)
m.deletionNew++
}
m.dirtyOld[key] = value
} else {
if _, ok := m.dirtyOld[key]; ok {
delete(m.dirtyOld, key)
m.deletionOld++
}
m.dirtyNew[key] = value
}
m.lock.Unlock()
}
func (m *SafeMap) Size() int {
m.lock.RLock()
size := len(m.dirtyOld) + len(m.dirtyNew)
m.lock.RUnlock()
return size
}

110
core/collection/safemap_test.go Normal file
View File

@@ -0,0 +1,110 @@
package collection
import (
"testing"
"zero/core/stringx"
"github.com/stretchr/testify/assert"
)
func TestSafeMap(t *testing.T) {
tests := []struct {
size int
exception int
}{
{
100000,
2000,
},
{
100000,
50,
},
}
for _, test := range tests {
t.Run(stringx.Rand(), func(t *testing.T) {
testSafeMapWithParameters(t, test.size, test.exception)
})
}
}
func TestSafeMap_CopyNew(t *testing.T) {
const (
size = 100000
exception1 = 5
exception2 = 500
)
m := NewSafeMap()
for i := 0; i < size; i++ {
m.Set(i, i)
}
for i := 0; i < size; i++ {
if i%exception1 == 0 {
m.Del(i)
}
}
for i := size; i < size<<1; i++ {
m.Set(i, i)
}
for i := size; i < size<<1; i++ {
if i%exception2 != 0 {
m.Del(i)
}
}
for i := 0; i < size; i++ {
val, ok := m.Get(i)
if i%exception1 != 0 {
assert.True(t, ok)
assert.Equal(t, i, val.(int))
} else {
assert.False(t, ok)
}
}
for i := size; i < size<<1; i++ {
val, ok := m.Get(i)
if i%exception2 == 0 {
assert.True(t, ok)
assert.Equal(t, i, val.(int))
} else {
assert.False(t, ok)
}
}
}
func testSafeMapWithParameters(t *testing.T, size, exception int) {
m := NewSafeMap()
for i := 0; i < size; i++ {
m.Set(i, i)
}
for i := 0; i < size; i++ {
if i%exception != 0 {
m.Del(i)
}
}
assert.Equal(t, size/exception, m.Size())
for i := size; i < size<<1; i++ {
m.Set(i, i)
}
for i := size; i < size<<1; i++ {
if i%exception != 0 {
m.Del(i)
}
}
for i := 0; i < size<<1; i++ {
val, ok := m.Get(i)
if i%exception == 0 {
assert.True(t, ok)
assert.Equal(t, i, val.(int))
} else {
assert.False(t, ok)
}
}
}

230
core/collection/set.go Normal file
View File

@@ -0,0 +1,230 @@
package collection
import (
"zero/core/lang"
"zero/core/logx"
)
const (
unmanaged = iota
untyped
intType
int64Type
uintType
uint64Type
stringType
)
type Set struct {
data map[interface{}]lang.PlaceholderType
tp int
}
func NewSet() *Set {
return &Set{
data: make(map[interface{}]lang.PlaceholderType),
tp: untyped,
}
}
func NewUnmanagedSet() *Set {
return &Set{
data: make(map[interface{}]lang.PlaceholderType),
tp: unmanaged,
}
}
func (s *Set) Add(i ...interface{}) {
for _, each := range i {
s.add(each)
}
}
func (s *Set) AddInt(ii ...int) {
for _, each := range ii {
s.add(each)
}
}
func (s *Set) AddInt64(ii ...int64) {
for _, each := range ii {
s.add(each)
}
}
func (s *Set) AddUint(ii ...uint) {
for _, each := range ii {
s.add(each)
}
}
func (s *Set) AddUint64(ii ...uint64) {
for _, each := range ii {
s.add(each)
}
}
func (s *Set) AddStr(ss ...string) {
for _, each := range ss {
s.add(each)
}
}
func (s *Set) Contains(i interface{}) bool {
if len(s.data) == 0 {
return false
}
s.validate(i)
_, ok := s.data[i]
return ok
}
func (s *Set) Keys() []interface{} {
var keys []interface{}
for key := range s.data {
keys = append(keys, key)
}
return keys
}
func (s *Set) KeysInt() []int {
var keys []int
for key := range s.data {
if intKey, ok := key.(int); !ok {
continue
} else {
keys = append(keys, intKey)
}
}
return keys
}
func (s *Set) KeysInt64() []int64 {
var keys []int64
for key := range s.data {
if intKey, ok := key.(int64); !ok {
continue
} else {
keys = append(keys, intKey)
}
}
return keys
}
func (s *Set) KeysUint() []uint {
var keys []uint
for key := range s.data {
if intKey, ok := key.(uint); !ok {
continue
} else {
keys = append(keys, intKey)
}
}
return keys
}
func (s *Set) KeysUint64() []uint64 {
var keys []uint64
for key := range s.data {
if intKey, ok := key.(uint64); !ok {
continue
} else {
keys = append(keys, intKey)
}
}
return keys
}
func (s *Set) KeysStr() []string {
var keys []string
for key := range s.data {
if strKey, ok := key.(string); !ok {
continue
} else {
keys = append(keys, strKey)
}
}
return keys
}
func (s *Set) Remove(i interface{}) {
s.validate(i)
delete(s.data, i)
}
func (s *Set) Count() int {
return len(s.data)
}
func (s *Set) add(i interface{}) {
switch s.tp {
case unmanaged:
// do nothing
case untyped:
s.setType(i)
default:
s.validate(i)
}
s.data[i] = lang.Placeholder
}
func (s *Set) setType(i interface{}) {
if s.tp != untyped {
return
}
switch i.(type) {
case int:
s.tp = intType
case int64:
s.tp = int64Type
case uint:
s.tp = uintType
case uint64:
s.tp = uint64Type
case string:
s.tp = stringType
}
}
func (s *Set) validate(i interface{}) {
if s.tp == unmanaged {
return
}
switch i.(type) {
case int:
if s.tp != intType {
logx.Errorf("Error: element is int, but set contains elements with type %d", s.tp)
}
case int64:
if s.tp != int64Type {
logx.Errorf("Error: element is int64, but set contains elements with type %d", s.tp)
}
case uint:
if s.tp != uintType {
logx.Errorf("Error: element is uint, but set contains elements with type %d", s.tp)
}
case uint64:
if s.tp != uint64Type {
logx.Errorf("Error: element is uint64, but set contains elements with type %d", s.tp)
}
case string:
if s.tp != stringType {
logx.Errorf("Error: element is string, but set contains elements with type %d", s.tp)
}
}
}

149
core/collection/set_test.go Normal file
View File

@@ -0,0 +1,149 @@
package collection
import (
"sort"
"testing"
"github.com/stretchr/testify/assert"
)
func BenchmarkRawSet(b *testing.B) {
m := make(map[interface{}]struct{})
for i := 0; i < b.N; i++ {
m[i] = struct{}{}
_ = m[i]
}
}
func BenchmarkUnmanagedSet(b *testing.B) {
s := NewUnmanagedSet()
for i := 0; i < b.N; i++ {
s.Add(i)
_ = s.Contains(i)
}
}
func BenchmarkSet(b *testing.B) {
s := NewSet()
for i := 0; i < b.N; i++ {
s.AddInt(i)
_ = s.Contains(i)
}
}
func TestAdd(t *testing.T) {
// given
set := NewUnmanagedSet()
values := []interface{}{1, 2, 3}
// when
set.Add(values...)
// then
assert.True(t, set.Contains(1) && set.Contains(2) && set.Contains(3))
assert.Equal(t, len(values), len(set.Keys()))
}
func TestAddInt(t *testing.T) {
// given
set := NewSet()
values := []int{1, 2, 3}
// when
set.AddInt(values...)
// then
assert.True(t, set.Contains(1) && set.Contains(2) && set.Contains(3))
keys := set.KeysInt()
sort.Ints(keys)
assert.EqualValues(t, values, keys)
}
func TestAddInt64(t *testing.T) {
// given
set := NewSet()
values := []int64{1, 2, 3}
// when
set.AddInt64(values...)
// then
assert.True(t, set.Contains(int64(1)) && set.Contains(int64(2)) && set.Contains(int64(3)))
assert.Equal(t, len(values), len(set.KeysInt64()))
}
func TestAddUint(t *testing.T) {
// given
set := NewSet()
values := []uint{1, 2, 3}
// when
set.AddUint(values...)
// then
assert.True(t, set.Contains(uint(1)) && set.Contains(uint(2)) && set.Contains(uint(3)))
assert.Equal(t, len(values), len(set.KeysUint()))
}
func TestAddUint64(t *testing.T) {
// given
set := NewSet()
values := []uint64{1, 2, 3}
// when
set.AddUint64(values...)
// then
assert.True(t, set.Contains(uint64(1)) && set.Contains(uint64(2)) && set.Contains(uint64(3)))
assert.Equal(t, len(values), len(set.KeysUint64()))
}
func TestAddStr(t *testing.T) {
// given
set := NewSet()
values := []string{"1", "2", "3"}
// when
set.AddStr(values...)
// then
assert.True(t, set.Contains("1") && set.Contains("2") && set.Contains("3"))
assert.Equal(t, len(values), len(set.KeysStr()))
}
func TestContainsWithoutElements(t *testing.T) {
// given
set := NewSet()
// then
assert.False(t, set.Contains(1))
}
func TestContainsUnmanagedWithoutElements(t *testing.T) {
// given
set := NewUnmanagedSet()
// then
assert.False(t, set.Contains(1))
}
func TestRemove(t *testing.T) {
// given
set := NewSet()
set.Add([]interface{}{1, 2, 3}...)
// when
set.Remove(2)
// then
assert.True(t, set.Contains(1) && !set.Contains(2) && set.Contains(3))
}
func TestCount(t *testing.T) {
// given
set := NewSet()
set.Add([]interface{}{1, 2, 3}...)
// then
assert.Equal(t, set.Count(), 3)
}

311
core/collection/timingwheel.go Normal file
View File

@@ -0,0 +1,311 @@
package collection
import (
"container/list"
"fmt"
"time"
"zero/core/lang"
"zero/core/threading"
"zero/core/timex"
)
const drainWorkers = 8
type (
Execute func(key, value interface{})
TimingWheel struct {
interval time.Duration
ticker timex.Ticker
slots []*list.List
timers *SafeMap
tickedPos int
numSlots int
execute Execute
setChannel chan timingEntry
moveChannel chan baseEntry
removeChannel chan interface{}
drainChannel chan func(key, value interface{})
stopChannel chan lang.PlaceholderType
}
timingEntry struct {
baseEntry
value interface{}
circle int
diff int
removed bool
}
baseEntry struct {
delay time.Duration
key interface{}
}
positionEntry struct {
pos int
item *timingEntry
}
timingTask struct {
key interface{}
value interface{}
}
)
func NewTimingWheel(interval time.Duration, numSlots int, execute Execute) (*TimingWheel, error) {
if interval <= 0 || numSlots <= 0 || execute == nil {
return nil, fmt.Errorf("interval: %v, slots: %d, execute: %p", interval, numSlots, execute)
}
return newTimingWheelWithClock(interval, numSlots, execute, timex.NewTicker(interval))
}
func newTimingWheelWithClock(interval time.Duration, numSlots int, execute Execute, ticker timex.Ticker) (
*TimingWheel, error) {
tw := &TimingWheel{
interval: interval,
ticker: ticker,
slots: make([]*list.List, numSlots),
timers: NewSafeMap(),
tickedPos: numSlots - 1, // at previous virtual circle
execute: execute,
numSlots: numSlots,
setChannel: make(chan timingEntry),
moveChannel: make(chan baseEntry),
removeChannel: make(chan interface{}),
drainChannel: make(chan func(key, value interface{})),
stopChannel: make(chan lang.PlaceholderType),
}
tw.initSlots()
go tw.run()
return tw, nil
}
func (tw *TimingWheel) Drain(fn func(key, value interface{})) {
tw.drainChannel <- fn
}
func (tw *TimingWheel) MoveTimer(key interface{}, delay time.Duration) {
if delay <= 0 || key == nil {
return
}
tw.moveChannel <- baseEntry{
delay: delay,
key: key,
}
}
func (tw *TimingWheel) RemoveTimer(key interface{}) {
if key == nil {
return
}
tw.removeChannel <- key
}
func (tw *TimingWheel) SetTimer(key, value interface{}, delay time.Duration) {
if delay <= 0 || key == nil {
return
}
tw.setChannel <- timingEntry{
baseEntry: baseEntry{
delay: delay,
key: key,
},
value: value,
}
}
func (tw *TimingWheel) Stop() {
close(tw.stopChannel)
}
func (tw *TimingWheel) drainAll(fn func(key, value interface{})) {
runner := threading.NewTaskRunner(drainWorkers)
for _, slot := range tw.slots {
for e := slot.Front(); e != nil; {
task := e.Value.(*timingEntry)
next := e.Next()
slot.Remove(e)
e = next
if !task.removed {
runner.Schedule(func() {
fn(task.key, task.value)
})
}
}
}
}
func (tw *TimingWheel) getPositionAndCircle(d time.Duration) (pos int, circle int) {
steps := int(d / tw.interval)
pos = (tw.tickedPos + steps) % tw.numSlots
circle = (steps - 1) / tw.numSlots
return
}
func (tw *TimingWheel) initSlots() {
for i := 0; i < tw.numSlots; i++ {
tw.slots[i] = list.New()
}
}
func (tw *TimingWheel) moveTask(task baseEntry) {
val, ok := tw.timers.Get(task.key)
if !ok {
return
}
timer := val.(*positionEntry)
if task.delay < tw.interval {
threading.GoSafe(func() {
tw.execute(timer.item.key, timer.item.value)
})
return
}
pos, circle := tw.getPositionAndCircle(task.delay)
if pos >= timer.pos {
timer.item.circle = circle
timer.item.diff = pos - timer.pos
} else if circle > 0 {
circle--
timer.item.circle = circle
timer.item.diff = tw.numSlots + pos - timer.pos
} else {
timer.item.removed = true
newItem := &timingEntry{
baseEntry: task,
value: timer.item.value,
}
tw.slots[pos].PushBack(newItem)
tw.setTimerPosition(pos, newItem)
}
}
func (tw *TimingWheel) onTick() {
tw.tickedPos = (tw.tickedPos + 1) % tw.numSlots
l := tw.slots[tw.tickedPos]
tw.scanAndRunTasks(l)
}
func (tw *TimingWheel) removeTask(key interface{}) {
val, ok := tw.timers.Get(key)
if !ok {
return
}
timer := val.(*positionEntry)
timer.item.removed = true
}
func (tw *TimingWheel) run() {
for {
select {
case <-tw.ticker.Chan():
tw.onTick()
case task := <-tw.setChannel:
tw.setTask(&task)
case key := <-tw.removeChannel:
tw.removeTask(key)
case task := <-tw.moveChannel:
tw.moveTask(task)
case fn := <-tw.drainChannel:
tw.drainAll(fn)
case <-tw.stopChannel:
tw.ticker.Stop()
return
}
}
}
func (tw *TimingWheel) runTasks(tasks []timingTask) {
if len(tasks) == 0 {
return
}
go func() {
for i := range tasks {
threading.RunSafe(func() {
tw.execute(tasks[i].key, tasks[i].value)
})
}
}()
}
func (tw *TimingWheel) scanAndRunTasks(l *list.List) {
var tasks []timingTask
for e := l.Front(); e != nil; {
task := e.Value.(*timingEntry)
if task.removed {
next := e.Next()
l.Remove(e)
tw.timers.Del(task.key)
e = next
continue
} else if task.circle > 0 {
task.circle--
e = e.Next()
continue
} else if task.diff > 0 {
next := e.Next()
l.Remove(e)
// (tw.tickedPos+task.diff)%tw.numSlots
// cannot be the same value of tw.tickedPos
pos := (tw.tickedPos + task.diff) % tw.numSlots
tw.slots[pos].PushBack(task)
tw.setTimerPosition(pos, task)
task.diff = 0
e = next
continue
}
tasks = append(tasks, timingTask{
key: task.key,
value: task.value,
})
next := e.Next()
l.Remove(e)
tw.timers.Del(task.key)
e = next
}
tw.runTasks(tasks)
}
func (tw *TimingWheel) setTask(task *timingEntry) {
if task.delay < tw.interval {
task.delay = tw.interval
}
if val, ok := tw.timers.Get(task.key); ok {
entry := val.(*positionEntry)
entry.item.value = task.value
tw.moveTask(task.baseEntry)
} else {
pos, circle := tw.getPositionAndCircle(task.delay)
task.circle = circle
tw.slots[pos].PushBack(task)
tw.setTimerPosition(pos, task)
}
}
func (tw *TimingWheel) setTimerPosition(pos int, task *timingEntry) {
if val, ok := tw.timers.Get(task.key); ok {
timer := val.(*positionEntry)
timer.pos = pos
} else {
tw.timers.Set(task.key, &positionEntry{
pos: pos,
item: task,
})
}
}

593
core/collection/timingwheel_test.go Normal file
View File

@@ -0,0 +1,593 @@
package collection
import (
"sort"
"sync"
"sync/atomic"
"testing"
"time"
"zero/core/lang"
"zero/core/stringx"
"zero/core/syncx"
"zero/core/timex"
"github.com/stretchr/testify/assert"
)
const (
testStep = time.Minute
waitTime = time.Second
)
func TestNewTimingWheel(t *testing.T) {
_, err := NewTimingWheel(0, 10, func(key, value interface{}) {})
assert.NotNil(t, err)
}
func TestTimingWheel_Drain(t *testing.T) {
ticker := timex.NewFakeTicker()
tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
}, ticker)
defer tw.Stop()
tw.SetTimer("first", 3, testStep*4)
tw.SetTimer("second", 5, testStep*7)
tw.SetTimer("third", 7, testStep*7)
var keys []string
var vals []int
var lock sync.Mutex
var wg sync.WaitGroup
wg.Add(3)
tw.Drain(func(key, value interface{}) {
lock.Lock()
defer lock.Unlock()
keys = append(keys, key.(string))
vals = append(vals, value.(int))
wg.Done()
})
wg.Wait()
sort.Strings(keys)
sort.Ints(vals)
assert.Equal(t, 3, len(keys))
assert.EqualValues(t, []string{"first", "second", "third"}, keys)
assert.EqualValues(t, []int{3, 5, 7}, vals)
var count int
tw.Drain(func(key, value interface{}) {
count++
})
time.Sleep(time.Millisecond * 100)
assert.Equal(t, 0, count)
}
func TestTimingWheel_SetTimerSoon(t *testing.T) {
run := syncx.NewAtomicBool()
ticker := timex.NewFakeTicker()
tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
assert.True(t, run.CompareAndSwap(false, true))
assert.Equal(t, "any", k)
assert.Equal(t, 3, v.(int))
ticker.Done()
}, ticker)
defer tw.Stop()
tw.SetTimer("any", 3, testStep>>1)
ticker.Tick()
assert.Nil(t, ticker.Wait(waitTime))
assert.True(t, run.True())
}
func TestTimingWheel_SetTimerTwice(t *testing.T) {
run := syncx.NewAtomicBool()
ticker := timex.NewFakeTicker()
tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
assert.True(t, run.CompareAndSwap(false, true))
assert.Equal(t, "any", k)
assert.Equal(t, 5, v.(int))
ticker.Done()
}, ticker)
defer tw.Stop()
tw.SetTimer("any", 3, testStep*4)
tw.SetTimer("any", 5, testStep*7)
for i := 0; i < 8; i++ {
ticker.Tick()
}
assert.Nil(t, ticker.Wait(waitTime))
assert.True(t, run.True())
}
func TestTimingWheel_SetTimerWrongDelay(t *testing.T) {
ticker := timex.NewFakeTicker()
tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {}, ticker)
defer tw.Stop()
assert.NotPanics(t, func() {
tw.SetTimer("any", 3, -testStep)
})
}
func TestTimingWheel_MoveTimer(t *testing.T) {
run := syncx.NewAtomicBool()
ticker := timex.NewFakeTicker()
tw, _ := newTimingWheelWithClock(testStep, 3, func(k, v interface{}) {
assert.True(t, run.CompareAndSwap(false, true))
assert.Equal(t, "any", k)
assert.Equal(t, 3, v.(int))
ticker.Done()
}, ticker)
defer tw.Stop()
tw.SetTimer("any", 3, testStep*4)
tw.MoveTimer("any", testStep*7)
tw.MoveTimer("any", -testStep)
tw.MoveTimer("none", testStep)
for i := 0; i < 5; i++ {
ticker.Tick()
}
assert.False(t, run.True())
for i := 0; i < 3; i++ {
ticker.Tick()
}
assert.Nil(t, ticker.Wait(waitTime))
assert.True(t, run.True())
}
func TestTimingWheel_MoveTimerSoon(t *testing.T) {
run := syncx.NewAtomicBool()
ticker := timex.NewFakeTicker()
tw, _ := newTimingWheelWithClock(testStep, 3, func(k, v interface{}) {
assert.True(t, run.CompareAndSwap(false, true))
assert.Equal(t, "any", k)
assert.Equal(t, 3, v.(int))
ticker.Done()
}, ticker)
defer tw.Stop()
tw.SetTimer("any", 3, testStep*4)
tw.MoveTimer("any", testStep>>1)
assert.Nil(t, ticker.Wait(waitTime))
assert.True(t, run.True())
}
func TestTimingWheel_MoveTimerEarlier(t *testing.T) {
run := syncx.NewAtomicBool()
ticker := timex.NewFakeTicker()
tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
assert.True(t, run.CompareAndSwap(false, true))
assert.Equal(t, "any", k)
assert.Equal(t, 3, v.(int))
ticker.Done()
}, ticker)
defer tw.Stop()
tw.SetTimer("any", 3, testStep*4)
tw.MoveTimer("any", testStep*2)
for i := 0; i < 3; i++ {
ticker.Tick()
}
assert.Nil(t, ticker.Wait(waitTime))
assert.True(t, run.True())
}
func TestTimingWheel_RemoveTimer(t *testing.T) {
ticker := timex.NewFakeTicker()
tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {}, ticker)
tw.SetTimer("any", 3, testStep)
assert.NotPanics(t, func() {
tw.RemoveTimer("any")
tw.RemoveTimer("none")
tw.RemoveTimer(nil)
})
for i := 0; i < 5; i++ {
ticker.Tick()
}
tw.Stop()
}
func TestTimingWheel_SetTimer(t *testing.T) {
tests := []struct {
slots int
setAt time.Duration
}{
{
slots: 5,
setAt: 5,
},
{
slots: 5,
setAt: 7,
},
{
slots: 5,
setAt: 10,
},
{
slots: 5,
setAt: 12,
},
{
slots: 5,
setAt: 7,
},
{
slots: 5,
setAt: 10,
},
{
slots: 5,
setAt: 12,
},
}
for _, test := range tests {
t.Run(stringx.RandId(), func(t *testing.T) {
var count int32
ticker := timex.NewFakeTicker()
tick := func() {
atomic.AddInt32(&count, 1)
ticker.Tick()
time.Sleep(time.Millisecond)
}
var actual int32
done := make(chan lang.PlaceholderType)
tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
assert.Equal(t, 1, key.(int))
assert.Equal(t, 2, value.(int))
actual = atomic.LoadInt32(&count)
close(done)
}, ticker)
assert.Nil(t, err)
defer tw.Stop()
tw.SetTimer(1, 2, testStep*test.setAt)
for {
select {
case <-done:
assert.Equal(t, int32(test.setAt), actual)
return
default:
tick()
}
}
})
}
}
func TestTimingWheel_SetAndMoveThenStart(t *testing.T) {
tests := []struct {
slots int
setAt time.Duration
moveAt time.Duration
}{
{
slots: 5,
setAt: 3,
moveAt: 5,
},
{
slots: 5,
setAt: 3,
moveAt: 7,
},
{
slots: 5,
setAt: 3,
moveAt: 10,
},
{
slots: 5,
setAt: 3,
moveAt: 12,
},
{
slots: 5,
setAt: 5,
moveAt: 7,
},
{
slots: 5,
setAt: 5,
moveAt: 10,
},
{
slots: 5,
setAt: 5,
moveAt: 12,
},
}
for _, test := range tests {
t.Run(stringx.RandId(), func(t *testing.T) {
var count int32
ticker := timex.NewFakeTicker()
tick := func() {
atomic.AddInt32(&count, 1)
ticker.Tick()
time.Sleep(time.Millisecond * 10)
}
var actual int32
done := make(chan lang.PlaceholderType)
tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
actual = atomic.LoadInt32(&count)
close(done)
}, ticker)
assert.Nil(t, err)
defer tw.Stop()
tw.SetTimer(1, 2, testStep*test.setAt)
tw.MoveTimer(1, testStep*test.moveAt)
for {
select {
case <-done:
assert.Equal(t, int32(test.moveAt), actual)
return
default:
tick()
}
}
})
}
}
func TestTimingWheel_SetAndMoveTwice(t *testing.T) {
tests := []struct {
slots int
setAt time.Duration
moveAt time.Duration
moveAgainAt time.Duration
}{
{
slots: 5,
setAt: 3,
moveAt: 5,
moveAgainAt: 10,
},
{
slots: 5,
setAt: 3,
moveAt: 7,
moveAgainAt: 12,
},
{
slots: 5,
setAt: 3,
moveAt: 10,
moveAgainAt: 15,
},
{
slots: 5,
setAt: 3,
moveAt: 12,
moveAgainAt: 17,
},
{
slots: 5,
setAt: 5,
moveAt: 7,
moveAgainAt: 12,
},
{
slots: 5,
setAt: 5,
moveAt: 10,
moveAgainAt: 17,
},
{
slots: 5,
setAt: 5,
moveAt: 12,
moveAgainAt: 17,
},
}
for _, test := range tests {
t.Run(stringx.RandId(), func(t *testing.T) {
var count int32
ticker := timex.NewFakeTicker()
tick := func() {
atomic.AddInt32(&count, 1)
ticker.Tick()
time.Sleep(time.Millisecond * 10)
}
var actual int32
done := make(chan lang.PlaceholderType)
tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
actual = atomic.LoadInt32(&count)
close(done)
}, ticker)
assert.Nil(t, err)
defer tw.Stop()
tw.SetTimer(1, 2, testStep*test.setAt)
tw.MoveTimer(1, testStep*test.moveAt)
tw.MoveTimer(1, testStep*test.moveAgainAt)
for {
select {
case <-done:
assert.Equal(t, int32(test.moveAgainAt), actual)
return
default:
tick()
}
}
})
}
}
func TestTimingWheel_ElapsedAndSet(t *testing.T) {
tests := []struct {
slots int
elapsed time.Duration
setAt time.Duration
}{
{
slots: 5,
elapsed: 3,
setAt: 5,
},
{
slots: 5,
elapsed: 3,
setAt: 7,
},
{
slots: 5,
elapsed: 3,
setAt: 10,
},
{
slots: 5,
elapsed: 3,
setAt: 12,
},
{
slots: 5,
elapsed: 5,
setAt: 7,
},
{
slots: 5,
elapsed: 5,
setAt: 10,
},
{
slots: 5,
elapsed: 5,
setAt: 12,
},
}
for _, test := range tests {
t.Run(stringx.RandId(), func(t *testing.T) {
var count int32
ticker := timex.NewFakeTicker()
tick := func() {
atomic.AddInt32(&count, 1)
ticker.Tick()
time.Sleep(time.Millisecond * 10)
}
var actual int32
done := make(chan lang.PlaceholderType)
tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
actual = atomic.LoadInt32(&count)
close(done)
}, ticker)
assert.Nil(t, err)
defer tw.Stop()
for i := 0; i < int(test.elapsed); i++ {
tick()
}
tw.SetTimer(1, 2, testStep*test.setAt)
for {
select {
case <-done:
assert.Equal(t, int32(test.elapsed+test.setAt), actual)
return
default:
tick()
}
}
})
}
}
func TestTimingWheel_ElapsedAndSetThenMove(t *testing.T) {
tests := []struct {
slots int
elapsed time.Duration
setAt time.Duration
moveAt time.Duration
}{
{
slots: 5,
elapsed: 3,
setAt: 5,
moveAt: 10,
},
{
slots: 5,
elapsed: 3,
setAt: 7,
moveAt: 12,
},
{
slots: 5,
elapsed: 3,
setAt: 10,
moveAt: 15,
},
{
slots: 5,
elapsed: 3,
setAt: 12,
moveAt: 16,
},
{
slots: 5,
elapsed: 5,
setAt: 7,
moveAt: 12,
},
{
slots: 5,
elapsed: 5,
setAt: 10,
moveAt: 15,
},
{
slots: 5,
elapsed: 5,
setAt: 12,
moveAt: 17,
},
}
for _, test := range tests {
t.Run(stringx.RandId(), func(t *testing.T) {
var count int32
ticker := timex.NewFakeTicker()
tick := func() {
atomic.AddInt32(&count, 1)
ticker.Tick()
time.Sleep(time.Millisecond * 10)
}
var actual int32
done := make(chan lang.PlaceholderType)
tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
actual = atomic.LoadInt32(&count)
close(done)
}, ticker)
assert.Nil(t, err)
defer tw.Stop()
for i := 0; i < int(test.elapsed); i++ {
tick()
}
tw.SetTimer(1, 2, testStep*test.setAt)
tw.MoveTimer(1, testStep*test.moveAt)
for {
select {
case <-done:
assert.Equal(t, int32(test.elapsed+test.moveAt), actual)
return
default:
tick()
}
}
})
}
}
func BenchmarkTimingWheel(b *testing.B) {
b.ReportAllocs()
tw, _ := NewTimingWheel(time.Second, 100, func(k, v interface{}) {})
for i := 0; i < b.N; i++ {
tw.SetTimer(i, i, time.Second)
tw.SetTimer(b.N+i, b.N+i, time.Second)
tw.MoveTimer(i, time.Second*time.Duration(i))
tw.RemoveTimer(i)
}
}