/
pool.go
139 lines (124 loc) · 2.5 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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
package spcdb
import (
"fmt"
"sync"
"time"
)
var (
MaxConnsInPool int = 20
TimeoutPing time.Duration = 2 // in minutes
DefaultDriverName string = "postgres"
)
type DBConfiguer interface {
DriverName() string
IsPing() bool
String() string
}
type poolType struct {
conns []*DB
busy []bool
driver string
dsn string
ping bool
m sync.RWMutex
}
type ptrType struct {
index int
nameConn string
}
var pools map[string]*poolType
var poolPtr map[*DB]ptrType
var mPtr sync.RWMutex
func init() {
pools = make(map[string]*poolType, 10)
poolPtr = make(map[*DB]ptrType, 40)
go func() {
for {
<-time.After(TimeoutPing * time.Minute)
for _, pool := range pools {
if pool.ping {
pingPool(pool)
}
}
}
}()
}
func pingPool(pool *poolType) {
freePools := make([]bool, MaxConnsInPool)
pool.m.RLock()
copy(freePools, pool.busy)
pool.m.RUnlock()
for index, busy := range freePools {
if !busy {
if db := pool.conns[index]; db != nil {
db.Ping()
}
}
}
}
func (db *DB) ConnectionName() string {
mPtr.RLock()
defer mPtr.RUnlock()
if ptr, found := poolPtr[db]; found {
return ptr.nameConn
}
return ""
}
func (db *DB) ReturnToPool() bool {
return ReturnToPool(db)
}
func NewPoolConnection(connectionName string, cfg DBConfiguer) {
drvName := cfg.DriverName()
if drvName == "" {
drvName = DefaultDriverName
}
pools[connectionName] = &poolType{
conns: make([]*DB, MaxConnsInPool),
busy: make([]bool, MaxConnsInPool),
driver: drvName,
dsn: cfg.String(),
ping: cfg.IsPing(),
}
}
func GetFromPool(connectionName string) (*DB, error) {
pool, found := pools[connectionName]
if !found {
return nil, fmt.Errorf("spcdb: No DB connection by name '%s'", connectionName)
}
pool.m.Lock()
defer pool.m.Unlock()
for index, busy := range pool.busy {
if !busy {
if db := pool.conns[index]; db != nil {
pool.busy[index] = true
return db, nil
}
db, err := Open(pool.driver, pool.dsn)
if err == nil {
pool.conns[index] = db
mPtr.Lock()
poolPtr[db] = ptrType{index, connectionName}
mPtr.Unlock()
pool.busy[index] = true
}
return db, err
}
}
return nil, fmt.Errorf("spcdb: No idle DB connections; '%s'", connectionName)
}
func ReturnToPool(db *DB) bool {
mPtr.RLock()
ptr, found := poolPtr[db]
mPtr.RUnlock()
if !found {
return false
}
pool, found := pools[ptr.nameConn]
if !found {
return false
}
pool.m.Lock()
pool.busy[ptr.index] = false
pool.m.Unlock()
return true
}