/
manager.go
109 lines (90 loc) · 1.9 KB
/
manager.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
package xgrpc
import (
"context"
"errors"
"sync"
"google.golang.org/grpc"
)
var (
ErrManagerClosed = errors.New("manager is closed")
)
type DialFunc func(ctx context.Context, addr string) (*grpc.ClientConn, error)
type ConnFunc func(addr string, dialFunc DialFunc) Conn
type managerOptions struct {
dial DialFunc
connFunc ConnFunc
}
var defaultOptions = managerOptions{
dial: func(ctx context.Context, addr string) (conn *grpc.ClientConn, e error) {
return grpc.DialContext(ctx, addr, grpc.WithInsecure())
},
connFunc: NewDefaultConn,
}
type ManagerOption func(options *managerOptions)
func WithDialFunc(dialFunc DialFunc) ManagerOption {
return func(options *managerOptions) {
options.dial = dialFunc
}
}
func WithConnFunc(connFunc ConnFunc) ManagerOption {
return func(options *managerOptions) {
options.connFunc = connFunc
}
}
// Manage a Conn for an address
type Manager struct {
managerOptions
rw sync.RWMutex
connMap map[string]Conn
closed bool
}
func NewManager(opts ...ManagerOption) *Manager {
options := defaultOptions
for _, opt := range opts {
opt(&options)
}
return &Manager{
managerOptions: options,
connMap: make(map[string]Conn),
}
}
func (m *Manager) GetConn(ctx context.Context, addr string) (conn Conn, err error) {
// fast-path
m.rw.RLock()
if m.closed {
m.rw.RUnlock()
return nil, ErrManagerClosed
}
var ok bool
conn, ok = m.connMap[addr]
m.rw.RUnlock()
if ok && !conn.IsClosed() {
return
}
// slow-path
m.rw.Lock()
defer m.rw.Unlock()
if m.closed {
return nil, ErrManagerClosed
}
// double-check
conn, ok = m.connMap[addr]
if !ok || conn.IsClosed() {
// make Conn
conn = m.connFunc(addr, m.dial)
m.connMap[addr] = conn
}
return
}
func (m *Manager) Close() error {
m.rw.Lock()
defer m.rw.Unlock()
if !m.closed {
m.closed = true
for _, c := range m.connMap {
_ = c.Close()
}
m.connMap = nil
}
return nil
}