-
Notifications
You must be signed in to change notification settings - Fork 0
/
pool.go
114 lines (96 loc) · 1.8 KB
/
pool.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
package pool
import (
"context"
"runtime"
"time"
)
type Factory[T any] func(context.Context) (T, error)
type entry[T any] struct {
elm T
time time.Time
}
type Option[T any] func(*pool[T])
// WithEvict set the evict callback
func WithEvict[T any](cb func(T)) Option[T] {
return func(p *pool[T]) {
p.evict = cb
}
}
// WithAge defined element max age (millisecond)
func WithAge[T any](maxAge int64) Option[T] {
return func(p *pool[T]) {
p.maxAge = maxAge
}
}
// WithSize defined max size of Pool
func WithSize[T any](maxSize int) Option[T] {
return func(p *pool[T]) {
p.ch = make(chan *entry[T], maxSize)
}
}
// Pool is for GC, see New for detail
type Pool[T any] struct {
*pool[T]
}
type pool[T any] struct {
ch chan *entry[T]
factory Factory[T]
evict func(T)
maxAge int64
}
func (p *pool[T]) GetContext(ctx context.Context) (T, error) {
now := time.Now()
for {
select {
case item := <-p.ch:
elm := item
if p.maxAge != 0 && now.Sub(item.time).Milliseconds() > p.maxAge {
if p.evict != nil {
p.evict(elm.elm)
}
continue
}
return elm.elm, nil
default:
return p.factory(ctx)
}
}
}
func (p *pool[T]) Get() (T, error) {
return p.GetContext(context.Background())
}
func (p *pool[T]) Put(item T) {
e := &entry[T]{
elm: item,
time: time.Now(),
}
select {
case p.ch <- e:
return
default:
// pool is full
if p.evict != nil {
p.evict(item)
}
return
}
}
func recycle[T any](p *Pool[T]) {
for item := range p.pool.ch {
if p.pool.evict != nil {
p.pool.evict(item.elm)
}
}
}
func New[T any](factory Factory[T], options ...Option[T]) *Pool[T] {
p := &pool[T]{
ch: make(chan *entry[T], 10),
factory: factory,
}
for _, option := range options {
option(p)
}
P := &Pool[T]{p}
runtime.SetFinalizer(P, recycle[T])
return P
}