/
default.go
115 lines (94 loc) · 1.94 KB
/
default.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
115
package pool
import (
"sync"
"time"
"c-z.dev/go-micro/transport"
"github.com/google/uuid"
)
type pool struct {
size int
ttl time.Duration
tr transport.Transport
sync.Mutex
conns map[string][]*poolConn
}
type poolConn struct {
transport.Client
id string
created time.Time
}
func newPool(options Options) *pool {
return &pool{
size: options.Size,
tr: options.Transport,
ttl: options.TTL,
conns: make(map[string][]*poolConn),
}
}
func (p *pool) Close() error {
p.Lock()
for k, c := range p.conns {
for _, conn := range c {
conn.Client.Close()
}
delete(p.conns, k)
}
p.Unlock()
return nil
}
// NoOp the Close since we manage it
func (p *poolConn) Close() error {
return nil
}
func (p *poolConn) Id() string {
return p.id
}
func (p *poolConn) Created() time.Time {
return p.created
}
func (p *pool) Get(addr string, opts ...transport.DialOption) (Conn, error) {
p.Lock()
conns := p.conns[addr]
// while we have conns check age and then return one
// otherwise we'll create a new conn
for len(conns) > 0 {
conn := conns[len(conns)-1]
conns = conns[:len(conns)-1]
p.conns[addr] = conns
// if conn is old kill it and move on
if d := time.Since(conn.Created()); d > p.ttl {
conn.Client.Close()
continue
}
// we got a good conn, lets unlock and return it
p.Unlock()
return conn, nil
}
p.Unlock()
// create new conn
c, err := p.tr.Dial(addr, opts...)
if err != nil {
return nil, err
}
return &poolConn{
Client: c,
id: uuid.New().String(),
created: time.Now(),
}, nil
}
func (p *pool) Release(conn Conn, err error) error {
// don't store the conn if it has errored
if err != nil {
return conn.(*poolConn).Client.Close()
}
// otherwise put it back for reuse
p.Lock()
conns := p.conns[conn.Remote()]
if len(conns) >= p.size {
p.Unlock()
return conn.(*poolConn).Client.Close()
}
p.conns[conn.Remote()] = append(conns, conn.(*poolConn))
p.Unlock()
return nil
}