简单的 goroutine pool

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
package gorpool

import (
"sync"
"time"
)

type gor struct {
ch chan func()
idx int
}

type Pool struct {
mu sync.Mutex
idle []*gor
Reserve uint
}

func reserve(kb uint) {
var stack [1024]byte
_ = stack[1023]
if kb <= 1 {
return
}
reserve(kb - 1)
}

func (p *Pool) gorproc(g *gor) {
// only a new gor reserves stack
if r := p.Reserve; r >= 1024 {
reserve(r / 1024)
}
tm := time.NewTimer(time.Minute)
down := false
for !down {
tm.Reset(time.Minute)
select {
case fn := <-g.ch:
fn()
p.mu.Lock()
g.idx = len(p.idle)
p.idle = append(p.idle, g)
p.mu.Unlock()
case <-tm.C:
p.mu.Lock()
if g.idx >= 0 {
t := p.idle[len(p.idle)-1]
t.idx = g.idx
p.idle[t.idx] = t
p.idle = p.idle[:len(p.idle)-1]
down = true
}
p.mu.Unlock()
}
}
}

func (p *Pool) Go(fn func()) {
if fn == nil {
panic("nil func")
}
var g *gor
p.mu.Lock()
if len(p.idle) > 0 {
g = p.idle[len(p.idle)-1]
p.idle = p.idle[:len(p.idle)-1]
g.idx = -1
}
p.mu.Unlock()
if g == nil {
g = &gor{
ch: make(chan func(), 1),
idx: -1,
}
go p.gorproc(g)
}
g.ch <- fn
}