Mercurial > gemma

comparison auth/connection.go @ 26:96a429c5f227

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Fundamental connection pool based on tokens.
author Sascha L. Teichmann <sascha.teichmann@intevation.de>
date Mon, 25 Jun 2018 13:19:55 +0200
parents
children c3e2cd7fa46f
comparison
equal deleted inserted replaced
25:a728610b3812 26:96a429c5f227
1 package auth
2
3 import (
4 "database/sql"
5 "errors"
6 "time"
7 )
8
9 var ErrNoSuchToken = errors.New("No such token")
10
11 var ConnPool = NewConnectionPool()
12
13 const (
14 maxOpen = 16
15 maxDBIdle = time.Minute * 5
16 maxTokenIdle = time.Minute * 20
17 )
18
19 type Connection struct {
20 user string
21 password string
22
23 access time.Time
24 db *sql.DB
25 }
26
27 func (c *Connection) set(user, password string) {
28 c.user = user
29 c.password = password
30 c.access = time.Now()
31 }
32
33 type ConnectionPool struct {
34 conns map[string]*Connection
35 cmds chan func(*ConnectionPool)
36 }
37
38 func NewConnectionPool() *ConnectionPool {
39 cp := &ConnectionPool{
40 conns: make(map[string]*Connection),
41 cmds: make(chan func(*ConnectionPool)),
42 }
43 go cp.run()
44 return cp
45 }
46
47 func (cp *ConnectionPool) run() {
48 for {
49 select {
50 case cmd := <-cp.cmds:
51 cmd(cp)
52 case <-time.After(time.Minute):
53 cp.cleanDB()
54 case <-time.After(time.Minute * 5):
55 cp.cleanToken()
56 }
57 }
58 }
59
60 func (cp *ConnectionPool) cleanDB() {
61 valid := time.Now().Add(-maxDBIdle)
62 for _, con := range cp.conns {
63 if con.access.Before(valid) {
64 con.db.Close()
65 con.db = nil
66 }
67 }
68 }
69
70 func (cp *ConnectionPool) cleanToken() {
71 valid := time.Now().Add(-maxDBIdle)
72 for token, con := range cp.conns {
73 if con.access.Before(valid) {
74 if con.db != nil {
75 // TODO: Be more graceful here?
76 con.db.Close()
77 con.db = nil
78 }
79 delete(cp.conns, token)
80 }
81 }
82 }
83
84 func (cp *ConnectionPool) Add(token, user, password string) *Connection {
85 res := make(chan *Connection)
86
87 cp.cmds <- func(cp *ConnectionPool) {
88
89 con := cp.conns[token]
90 if con == nil {
91 con = &Connection{}
92 cp.conns[token] = con
93 }
94 con.set(user, password)
95 res <- con
96 }
97
98 con := <-res
99 return con
100 }
101
102 func trim(cp *ConnectionPool) {
103
104 least := time.Now()
105 var count int
106 var oldest *Connection
107
108 for _, con := range cp.conns {
109 if con.db != nil {
110 if con.access.Before(least) {
111 least = con.access
112 oldest = con
113 }
114 count++
115 }
116 }
117 if count > maxOpen {
118 oldest.db.Close()
119 oldest.db = nil
120 }
121 return
122 }
123
124 func (cp *ConnectionPool) Do(token string, fn func(*sql.DB) error) error {
125
126 type result struct {
127 con *Connection
128 err error
129 }
130
131 res := make(chan result)
132
133 cp.cmds <- func(cp *ConnectionPool) {
134 con := cp.conns[token]
135 if con == nil {
136 res <- result{err: ErrNoSuchToken}
137 return
138 }
139 con.access = time.Now()
140 if con.db != nil {
141 res <- result{con: con}
142 return
143 }
144
145 db, err := opendb(con.user, con.password)
146 if err != nil {
147 res <- result{err: err}
148 return
149 }
150 con.db = db
151 }
152
153 r := <-res
154
155 if r.err != nil {
156 return r.err
157 }
158
159 defer func() { cp.cmds <- trim }()
160
161 return fn(r.con.db)
162 }