├── LICENSE └── pool ├── README.md ├── bench_test.go ├── conn.go ├── go.mod ├── go.sum ├── main_test.go ├── pool.go └── pool_test.go /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 bilibili 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /pool/README.md: -------------------------------------------------------------------------------- 1 | # pool 2 | 3 | 通用连接池，源自 https://github.com/go-redis/redis 4 | 5 | 主要添加了 ctx 支持。 6 | -------------------------------------------------------------------------------- /pool/bench_test.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "testing" 7 | "time" 8 | ) 9 | 10 | type poolGetPutBenchmark struct { 11 | poolSize int 12 | } 13 | 14 | func (bm poolGetPutBenchmark) String() string { 15 | return fmt.Sprintf("pool=%d", bm.poolSize) 16 | } 17 | 18 | func BenchmarkPoolGetPut(b *testing.B) { 19 | benchmarks := []poolGetPutBenchmark{ 20 | {1}, 21 | {2}, 22 | {8}, 23 | {32}, 24 | {64}, 25 | {128}, 26 | } 27 | for _, bm := range benchmarks { 28 | b.Run(bm.String(), func(b *testing.B) { 29 | connPool := New(Options{ 30 | Dialer: dummyDialer, 31 | PoolSize: bm.poolSize, 32 | PoolTimeout: time.Second, 33 | IdleTimeout: time.Hour, 34 | IdleCheckFrequency: time.Hour, 35 | }) 36 | 37 | b.ResetTimer() 38 | 39 | b.RunParallel(func(pb *testing.PB) { 40 | for pb.Next() { 41 | cn, err := connPool.Get(context.Background()) 42 | if err != nil { 43 | b.Fatal(err) 44 | } 45 | connPool.Put(cn) 46 | } 47 | }) 48 | }) 49 | } 50 | } 51 | 52 | type poolGetRemoveBenchmark struct { 53 | poolSize int 54 | } 55 | 56 | func (bm poolGetRemoveBenchmark) String() string { 57 | return fmt.Sprintf("pool=%d", bm.poolSize) 58 | } 59 | 60 | func BenchmarkPoolGetRemove(b *testing.B) { 61 | benchmarks := []poolGetRemoveBenchmark{ 62 | {1}, 63 | {2}, 64 | {8}, 65 | {32}, 66 | {64}, 67 | {128}, 68 | } 69 | for _, bm := range benchmarks { 70 | b.Run(bm.String(), func(b *testing.B) { 71 | connPool := New(Options{ 72 | Dialer: dummyDialer, 73 | PoolSize: bm.poolSize, 74 | PoolTimeout: time.Second, 75 | IdleTimeout: time.Hour, 76 | IdleCheckFrequency: time.Hour, 77 | }) 78 | 79 | b.ResetTimer() 80 | 81 | b.RunParallel(func(pb *testing.PB) { 82 | for pb.Next() { 83 | cn, err := connPool.Get(context.Background()) 84 | if err != nil { 85 | b.Fatal(err) 86 | } 87 | connPool.Remove(cn) 88 | } 89 | }) 90 | }) 91 | } 92 | } 93 | -------------------------------------------------------------------------------- /pool/conn.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "sync/atomic" 5 | "time" 6 | ) 7 | 8 | var noDeadline = time.Time{} 9 | 10 | type Conn struct { 11 | C Closer 12 | 13 | Inited bool 14 | pooled bool 15 | createdAt time.Time 16 | usedAt atomic.Value 17 | } 18 | 19 | type Closer interface { 20 | Close() error 21 | } 22 | 23 | func newConn(netConn Closer) *Conn { 24 | cn := &Conn{ 25 | C: netConn, 26 | createdAt: time.Now(), 27 | } 28 | cn.setUsedAt(time.Now()) 29 | return cn 30 | } 31 | 32 | func (cn *Conn) UsedAt() time.Time { 33 | return cn.usedAt.Load().(time.Time) 34 | } 35 | 36 | func (cn *Conn) setUsedAt(tm time.Time) { 37 | cn.usedAt.Store(tm) 38 | } 39 | 40 | func (cn *Conn) Close() error { 41 | return cn.C.Close() 42 | } 43 | -------------------------------------------------------------------------------- /pool/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/go-kiss/net/pool 2 | 3 | go 1.12 4 | 5 | require ( 6 | github.com/onsi/ginkgo v1.8.0 7 | github.com/onsi/gomega v1.5.0 8 | ) 9 | -------------------------------------------------------------------------------- /pool/go.sum: -------------------------------------------------------------------------------- 1 | github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I= 2 | github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= 3 | github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM= 4 | github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= 5 | github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI= 6 | github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= 7 | github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= 8 | github.com/onsi/ginkgo v1.8.0 h1:VkHVNpR4iVnU8XQR6DBm8BqYjN7CRzw+xKUbVVbbW9w= 9 | github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= 10 | github.com/onsi/gomega v1.5.0 h1:izbySO9zDPmjJ8rDjLvkA2zJHIo+HkYXHnf7eN7SSyo= 11 | github.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= 12 | golang.org/x/net v0.0.0-20180906233101-161cd47e91fd h1:nTDtHvHSdCn1m6ITfMRqtOd/9+7a3s8RBNOZ3eYZzJA= 13 | golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= 14 | golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f h1:wMNYb4v58l5UBM7MYRLPG6ZhfOqbKu7X5eyFl8ZhKvA= 15 | golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 16 | golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e h1:o3PsSEY8E4eXWkXrIP9YJALUkVZqzHJT5DOasTyn8Vs= 17 | golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 18 | golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg= 19 | golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= 20 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 21 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 22 | gopkg.in/fsnotify.v1 v1.4.7 h1:xOHLXZwVvI9hhs+cLKq5+I5onOuwQLhQwiu63xxlHs4= 23 | gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= 24 | gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= 25 | gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= 26 | gopkg.in/yaml.v2 v2.2.1 h1:mUhvW9EsL+naU5Q3cakzfE91YhliOondGd6ZrsDBHQE= 27 | gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= 28 | -------------------------------------------------------------------------------- /pool/main_test.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "net" 6 | "sync" 7 | "testing" 8 | 9 | . "github.com/onsi/ginkgo" 10 | . "github.com/onsi/gomega" 11 | ) 12 | 13 | func TestGinkgoSuite(t *testing.T) { 14 | RegisterFailHandler(Fail) 15 | RunSpecs(t, "pool") 16 | } 17 | 18 | func perform(n int, cbs ...func(int)) { 19 | var wg sync.WaitGroup 20 | for _, cb := range cbs { 21 | for i := 0; i < n; i++ { 22 | wg.Add(1) 23 | go func(cb func(int), i int) { 24 | defer GinkgoRecover() 25 | defer wg.Done() 26 | 27 | cb(i) 28 | }(cb, i) 29 | } 30 | } 31 | wg.Wait() 32 | } 33 | 34 | func dummyDialer(ctx context.Context) (Closer, error) { 35 | return &net.TCPConn{}, nil 36 | } 37 | -------------------------------------------------------------------------------- /pool/pool.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "log" 7 | "sync" 8 | "sync/atomic" 9 | "time" 10 | ) 11 | 12 | var ( 13 | ErrClosed = errors.New("net/pool: connection is closed") 14 | ErrPoolTimeout = errors.New("net/pool: connection pool timeout") 15 | ) 16 | 17 | var timers = sync.Pool{ 18 | New: func() interface{} { 19 | t := time.NewTimer(time.Hour) 20 | t.Stop() 21 | return t 22 | }, 23 | } 24 | 25 | // Stats contains pool state information and accumulated stats. 26 | type Stats struct { 27 | Hits uint32 // number of times free connection was found in the pool 28 | Misses uint32 // number of times free connection was NOT found in the pool 29 | Timeouts uint32 // number of times a wait timeout occurred 30 | 31 | TotalConns uint32 // number of total connections in the pool 32 | IdleConns uint32 // number of idle connections in the pool 33 | StaleConns uint32 // number of stale connections removed from the pool 34 | } 35 | 36 | type Pooler interface { 37 | NewConn(context.Context) (*Conn, error) 38 | CloseConn(*Conn) error 39 | 40 | Get(context.Context) (*Conn, error) 41 | Put(*Conn) 42 | Remove(*Conn) 43 | 44 | Len() int 45 | IdleLen() int 46 | Stats() *Stats 47 | 48 | Close() error 49 | } 50 | 51 | type Options struct { 52 | Dialer func(context.Context) (Closer, error) 53 | OnClose func(*Conn) error 54 | 55 | PoolSize int 56 | MinIdleConns int 57 | MaxConnAge time.Duration 58 | PoolTimeout time.Duration 59 | IdleTimeout time.Duration 60 | IdleCheckFrequency time.Duration 61 | } 62 | 63 | type ConnPool struct { 64 | opt Options 65 | 66 | dialErrorsNum uint32 // atomic 67 | 68 | lastDialErrorMu sync.RWMutex 69 | lastDialError error 70 | 71 | queue chan struct{} 72 | 73 | connsMu sync.Mutex 74 | conns []*Conn 75 | idleConns []*Conn 76 | poolSize int 77 | idleConnsLen int 78 | 79 | stats Stats 80 | 81 | _closed uint32 // atomic 82 | } 83 | 84 | var _ Pooler = (*ConnPool)(nil) 85 | 86 | func New(opt Options) Pooler { 87 | p := &ConnPool{ 88 | opt: opt, 89 | 90 | queue: make(chan struct{}, opt.PoolSize), 91 | conns: make([]*Conn, 0, opt.PoolSize), 92 | 93 | idleConns: make([]*Conn, 0, opt.PoolSize), 94 | } 95 | 96 | for i := 0; i < opt.MinIdleConns; i++ { 97 | p.checkMinIdleConns() 98 | } 99 | 100 | if opt.IdleTimeout > 0 && opt.IdleCheckFrequency > 0 { 101 | go p.reaper(opt.IdleCheckFrequency) 102 | } 103 | 104 | return p 105 | } 106 | 107 | func (p *ConnPool) checkMinIdleConns() { 108 | if p.opt.MinIdleConns == 0 { 109 | return 110 | } 111 | if p.poolSize < p.opt.PoolSize && p.idleConnsLen < p.opt.MinIdleConns { 112 | p.poolSize++ 113 | p.idleConnsLen++ 114 | go p.addIdleConn() 115 | } 116 | } 117 | 118 | func (p *ConnPool) addIdleConn() { 119 | cn, err := p.newConn(context.Background(), true) 120 | if err != nil { 121 | return 122 | } 123 | 124 | p.connsMu.Lock() 125 | p.conns = append(p.conns, cn) 126 | p.idleConns = append(p.idleConns, cn) 127 | p.connsMu.Unlock() 128 | } 129 | 130 | func (p *ConnPool) NewConn(ctx context.Context) (*Conn, error) { 131 | return p._NewConn(ctx, false) 132 | } 133 | 134 | func (p *ConnPool) _NewConn(ctx context.Context, pooled bool) (*Conn, error) { 135 | cn, err := p.newConn(ctx, pooled) 136 | if err != nil { 137 | return nil, err 138 | } 139 | 140 | p.connsMu.Lock() 141 | p.conns = append(p.conns, cn) 142 | if pooled { 143 | if p.poolSize < p.opt.PoolSize { 144 | p.poolSize++ 145 | } else { 146 | cn.pooled = false 147 | } 148 | } 149 | p.connsMu.Unlock() 150 | return cn, nil 151 | } 152 | 153 | func (p *ConnPool) newConn(ctx context.Context, pooled bool) (*Conn, error) { 154 | if p.closed() { 155 | return nil, ErrClosed 156 | } 157 | 158 | if atomic.LoadUint32(&p.dialErrorsNum) >= uint32(p.opt.PoolSize) { 159 | return nil, p.getLastDialError() 160 | } 161 | 162 | netConn, err := p.opt.Dialer(ctx) 163 | if err != nil { 164 | p.setLastDialError(err) 165 | if atomic.AddUint32(&p.dialErrorsNum, 1) == uint32(p.opt.PoolSize) { 166 | go p.tryDial() 167 | } 168 | return nil, err 169 | } 170 | 171 | cn := newConn(netConn) 172 | cn.pooled = pooled 173 | return cn, nil 174 | } 175 | 176 | func (p *ConnPool) tryDial() { 177 | for { 178 | if p.closed() { 179 | return 180 | } 181 | 182 | conn, err := p.opt.Dialer(context.Background()) 183 | if err != nil { 184 | p.setLastDialError(err) 185 | time.Sleep(time.Second) 186 | continue 187 | } 188 | 189 | atomic.StoreUint32(&p.dialErrorsNum, 0) 190 | _ = conn.Close() 191 | return 192 | } 193 | } 194 | 195 | func (p *ConnPool) setLastDialError(err error) { 196 | p.lastDialErrorMu.Lock() 197 | p.lastDialError = err 198 | p.lastDialErrorMu.Unlock() 199 | } 200 | 201 | func (p *ConnPool) getLastDialError() error { 202 | p.lastDialErrorMu.RLock() 203 | err := p.lastDialError 204 | p.lastDialErrorMu.RUnlock() 205 | return err 206 | } 207 | 208 | // Get returns existed connection from the pool or creates a new one. 209 | func (p *ConnPool) Get(ctx context.Context) (*Conn, error) { 210 | if p.closed() { 211 | return nil, ErrClosed 212 | } 213 | 214 | err := p.waitTurn(ctx) 215 | if err != nil { 216 | return nil, err 217 | } 218 | 219 | for { 220 | p.connsMu.Lock() 221 | cn := p.popIdle() 222 | p.connsMu.Unlock() 223 | 224 | if cn == nil { 225 | break 226 | } 227 | 228 | if p.isStaleConn(cn) { 229 | _ = p.CloseConn(cn) 230 | continue 231 | } 232 | 233 | atomic.AddUint32(&p.stats.Hits, 1) 234 | return cn, nil 235 | } 236 | 237 | atomic.AddUint32(&p.stats.Misses, 1) 238 | 239 | newcn, err := p._NewConn(ctx, true) 240 | if err != nil { 241 | p.freeTurn() 242 | return nil, err 243 | } 244 | 245 | return newcn, nil 246 | } 247 | 248 | func (p *ConnPool) getTurn() { 249 | p.queue <- struct{}{} 250 | } 251 | 252 | func (p *ConnPool) waitTurn(ctx context.Context) error { 253 | select { 254 | case p.queue <- struct{}{}: 255 | return nil 256 | default: 257 | timer := timers.Get().(*time.Timer) 258 | timer.Reset(p.opt.PoolTimeout) 259 | 260 | select { 261 | case p.queue <- struct{}{}: 262 | if !timer.Stop() { 263 | <-timer.C 264 | } 265 | timers.Put(timer) 266 | return nil 267 | case <-timer.C: 268 | timers.Put(timer) 269 | atomic.AddUint32(&p.stats.Timeouts, 1) 270 | return ErrPoolTimeout 271 | case <-ctx.Done(): 272 | timers.Put(timer) 273 | atomic.AddUint32(&p.stats.Timeouts, 1) 274 | return ErrPoolTimeout 275 | } 276 | } 277 | } 278 | 279 | func (p *ConnPool) freeTurn() { 280 | <-p.queue 281 | } 282 | 283 | func (p *ConnPool) popIdle() *Conn { 284 | if len(p.idleConns) == 0 { 285 | return nil 286 | } 287 | 288 | idx := len(p.idleConns) - 1 289 | cn := p.idleConns[idx] 290 | p.idleConns = p.idleConns[:idx] 291 | p.idleConnsLen-- 292 | p.checkMinIdleConns() 293 | return cn 294 | } 295 | 296 | func (p *ConnPool) Put(cn *Conn) { 297 | if !cn.pooled { 298 | p.Remove(cn) 299 | return 300 | } 301 | 302 | p.connsMu.Lock() 303 | p.idleConns = append(p.idleConns, cn) 304 | p.idleConnsLen++ 305 | p.connsMu.Unlock() 306 | p.freeTurn() 307 | } 308 | 309 | func (p *ConnPool) Remove(cn *Conn) { 310 | p.removeConn(cn) 311 | p.freeTurn() 312 | _ = p.closeConn(cn) 313 | } 314 | 315 | func (p *ConnPool) CloseConn(cn *Conn) error { 316 | p.removeConn(cn) 317 | return p.closeConn(cn) 318 | } 319 | 320 | func (p *ConnPool) removeConn(cn *Conn) { 321 | p.connsMu.Lock() 322 | for i, c := range p.conns { 323 | if c == cn { 324 | p.conns = append(p.conns[:i], p.conns[i+1:]...) 325 | if cn.pooled { 326 | p.poolSize-- 327 | p.checkMinIdleConns() 328 | } 329 | break 330 | } 331 | } 332 | p.connsMu.Unlock() 333 | } 334 | 335 | func (p *ConnPool) closeConn(cn *Conn) error { 336 | if p.opt.OnClose != nil { 337 | _ = p.opt.OnClose(cn) 338 | } 339 | return cn.Close() 340 | } 341 | 342 | // Len returns total number of connections. 343 | func (p *ConnPool) Len() int { 344 | p.connsMu.Lock() 345 | n := len(p.conns) 346 | p.connsMu.Unlock() 347 | return n 348 | } 349 | 350 | // IdleLen returns number of idle connections. 351 | func (p *ConnPool) IdleLen() int { 352 | p.connsMu.Lock() 353 | n := p.idleConnsLen 354 | p.connsMu.Unlock() 355 | return n 356 | } 357 | 358 | func (p *ConnPool) Stats() *Stats { 359 | idleLen := p.IdleLen() 360 | return &Stats{ 361 | Hits: atomic.LoadUint32(&p.stats.Hits), 362 | Misses: atomic.LoadUint32(&p.stats.Misses), 363 | Timeouts: atomic.LoadUint32(&p.stats.Timeouts), 364 | 365 | TotalConns: uint32(p.Len()), 366 | IdleConns: uint32(idleLen), 367 | StaleConns: atomic.LoadUint32(&p.stats.StaleConns), 368 | } 369 | } 370 | 371 | func (p *ConnPool) closed() bool { 372 | return atomic.LoadUint32(&p._closed) == 1 373 | } 374 | 375 | func (p *ConnPool) Filter(fn func(*Conn) bool) error { 376 | var firstErr error 377 | p.connsMu.Lock() 378 | for _, cn := range p.conns { 379 | if fn(cn) { 380 | if err := p.closeConn(cn); err != nil && firstErr == nil { 381 | firstErr = err 382 | } 383 | } 384 | } 385 | p.connsMu.Unlock() 386 | return firstErr 387 | } 388 | 389 | func (p *ConnPool) Close() error { 390 | if !atomic.CompareAndSwapUint32(&p._closed, 0, 1) { 391 | return ErrClosed 392 | } 393 | 394 | var firstErr error 395 | p.connsMu.Lock() 396 | for _, cn := range p.conns { 397 | if err := p.closeConn(cn); err != nil && firstErr == nil { 398 | firstErr = err 399 | } 400 | } 401 | p.conns = nil 402 | p.poolSize = 0 403 | p.idleConns = nil 404 | p.idleConnsLen = 0 405 | p.connsMu.Unlock() 406 | 407 | return firstErr 408 | } 409 | 410 | func (p *ConnPool) reapStaleConn() *Conn { 411 | if len(p.idleConns) == 0 { 412 | return nil 413 | } 414 | 415 | cn := p.idleConns[0] 416 | if !p.isStaleConn(cn) { 417 | return nil 418 | } 419 | 420 | p.idleConns = append(p.idleConns[:0], p.idleConns[1:]...) 421 | p.idleConnsLen-- 422 | 423 | return cn 424 | } 425 | 426 | func (p *ConnPool) reapStaleConns() (int, error) { 427 | var n int 428 | for { 429 | p.getTurn() 430 | 431 | p.connsMu.Lock() 432 | cn := p.reapStaleConn() 433 | p.connsMu.Unlock() 434 | 435 | if cn != nil { 436 | p.removeConn(cn) 437 | } 438 | 439 | p.freeTurn() 440 | 441 | if cn != nil { 442 | p.closeConn(cn) 443 | n++ 444 | } else { 445 | break 446 | } 447 | } 448 | return n, nil 449 | } 450 | 451 | func (p *ConnPool) reaper(frequency time.Duration) { 452 | ticker := time.NewTicker(frequency) 453 | defer ticker.Stop() 454 | 455 | for range ticker.C { 456 | if p.closed() { 457 | break 458 | } 459 | n, err := p.reapStaleConns() 460 | if err != nil { 461 | log.Printf("ReapStaleConns failed: %s", err) 462 | continue 463 | } 464 | atomic.AddUint32(&p.stats.StaleConns, uint32(n)) 465 | } 466 | } 467 | 468 | func (p *ConnPool) isStaleConn(cn *Conn) bool { 469 | if p.opt.IdleTimeout == 0 && p.opt.MaxConnAge == 0 { 470 | return false 471 | } 472 | 473 | now := time.Now() 474 | if p.opt.IdleTimeout > 0 && now.Sub(cn.UsedAt()) >= p.opt.IdleTimeout { 475 | return true 476 | } 477 | if p.opt.MaxConnAge > 0 && now.Sub(cn.createdAt) >= p.opt.MaxConnAge { 478 | return true 479 | } 480 | 481 | return false 482 | } 483 | -------------------------------------------------------------------------------- /pool/pool_test.go: -------------------------------------------------------------------------------- 1 | package pool 2 | 3 | import ( 4 | "context" 5 | "sync" 6 | "testing" 7 | "time" 8 | 9 | . "github.com/onsi/ginkgo" 10 | . "github.com/onsi/gomega" 11 | ) 12 | 13 | var _ = Describe("ConnPool", func() { 14 | var connPool Pooler 15 | 16 | BeforeEach(func() { 17 | connPool = New(Options{ 18 | Dialer: dummyDialer, 19 | PoolSize: 10, 20 | PoolTimeout: time.Hour, 21 | IdleTimeout: time.Millisecond, 22 | IdleCheckFrequency: time.Millisecond, 23 | }) 24 | }) 25 | 26 | AfterEach(func() { 27 | connPool.Close() 28 | }) 29 | 30 | It("should unblock client when conn is removed", func() { 31 | // Reserve one connection. 32 | cn, err := connPool.Get(context.Background()) 33 | Expect(err).NotTo(HaveOccurred()) 34 | 35 | // Reserve all other connections. 36 | var cns []*Conn 37 | for i := 0; i < 9; i++ { 38 | cn, err := connPool.Get(context.Background()) 39 | Expect(err).NotTo(HaveOccurred()) 40 | cns = append(cns, cn) 41 | } 42 | 43 | started := make(chan bool, 1) 44 | done := make(chan bool, 1) 45 | go func() { 46 | defer GinkgoRecover() 47 | 48 | started <- true 49 | _, err := connPool.Get(context.Background()) 50 | Expect(err).NotTo(HaveOccurred()) 51 | done <- true 52 | 53 | connPool.Put(cn) 54 | }() 55 | <-started 56 | 57 | // Check that Get is blocked. 58 | select { 59 | case <-done: 60 | Fail("Get is not blocked") 61 | case <-time.After(time.Millisecond): 62 | // ok 63 | } 64 | 65 | connPool.Remove(cn) 66 | 67 | // Check that Get is unblocked. 68 | select { 69 | case <-done: 70 | // ok 71 | case <-time.After(time.Second): 72 | Fail("Get is not unblocked") 73 | } 74 | 75 | for _, cn := range cns { 76 | connPool.Put(cn) 77 | } 78 | }) 79 | }) 80 | 81 | var _ = Describe("MinIdleConns", func() { 82 | const poolSize = 100 83 | var minIdleConns int 84 | var connPool Pooler 85 | 86 | newConnPool := func() Pooler { 87 | connPool := New(Options{ 88 | Dialer: dummyDialer, 89 | PoolSize: poolSize, 90 | MinIdleConns: minIdleConns, 91 | PoolTimeout: 100 * time.Millisecond, 92 | IdleTimeout: -1, 93 | IdleCheckFrequency: -1, 94 | }) 95 | Eventually(func() int { 96 | return connPool.Len() 97 | }).Should(Equal(minIdleConns)) 98 | return connPool 99 | } 100 | 101 | assert := func() { 102 | It("has idle connections when created", func() { 103 | Expect(connPool.Len()).To(Equal(minIdleConns)) 104 | Expect(connPool.IdleLen()).To(Equal(minIdleConns)) 105 | }) 106 | 107 | Context("after Get", func() { 108 | var cn *Conn 109 | 110 | BeforeEach(func() { 111 | var err error 112 | cn, err = connPool.Get(context.Background()) 113 | Expect(err).NotTo(HaveOccurred()) 114 | 115 | Eventually(func() int { 116 | return connPool.Len() 117 | }).Should(Equal(minIdleConns + 1)) 118 | }) 119 | 120 | It("has idle connections", func() { 121 | Expect(connPool.Len()).To(Equal(minIdleConns + 1)) 122 | Expect(connPool.IdleLen()).To(Equal(minIdleConns)) 123 | }) 124 | 125 | Context("after Remove", func() { 126 | BeforeEach(func() { 127 | connPool.Remove(cn) 128 | }) 129 | 130 | It("has idle connections", func() { 131 | Expect(connPool.Len()).To(Equal(minIdleConns)) 132 | Expect(connPool.IdleLen()).To(Equal(minIdleConns)) 133 | }) 134 | }) 135 | }) 136 | 137 | Describe("Get does not exceed pool size", func() { 138 | var mu sync.RWMutex 139 | var cns []*Conn 140 | 141 | BeforeEach(func() { 142 | cns = make([]*Conn, 0) 143 | 144 | perform(poolSize, func(_ int) { 145 | defer GinkgoRecover() 146 | 147 | cn, err := connPool.Get(context.Background()) 148 | Expect(err).NotTo(HaveOccurred()) 149 | mu.Lock() 150 | cns = append(cns, cn) 151 | mu.Unlock() 152 | }) 153 | 154 | Eventually(func() int { 155 | return connPool.Len() 156 | }).Should(BeNumerically(">=", poolSize)) 157 | }) 158 | 159 | It("Get is blocked", func() { 160 | done := make(chan struct{}) 161 | go func() { 162 | connPool.Get(context.Background()) 163 | close(done) 164 | }() 165 | 166 | select { 167 | case <-done: 168 | Fail("Get is not blocked") 169 | case <-time.After(time.Millisecond): 170 | // ok 171 | } 172 | 173 | select { 174 | case <-done: 175 | // ok 176 | case <-time.After(time.Second): 177 | Fail("Get is not unblocked") 178 | } 179 | }) 180 | 181 | Context("after Put", func() { 182 | BeforeEach(func() { 183 | perform(len(cns), func(i int) { 184 | mu.RLock() 185 | connPool.Put(cns[i]) 186 | mu.RUnlock() 187 | }) 188 | 189 | Eventually(func() int { 190 | return connPool.Len() 191 | }).Should(Equal(poolSize)) 192 | }) 193 | 194 | It("pool.Len is back to normal", func() { 195 | Expect(connPool.Len()).To(Equal(poolSize)) 196 | Expect(connPool.IdleLen()).To(Equal(poolSize)) 197 | }) 198 | }) 199 | 200 | Context("after Remove", func() { 201 | BeforeEach(func() { 202 | perform(len(cns), func(i int) { 203 | mu.RLock() 204 | connPool.Remove(cns[i]) 205 | mu.RUnlock() 206 | }) 207 | 208 | Eventually(func() int { 209 | return connPool.Len() 210 | }).Should(Equal(minIdleConns)) 211 | }) 212 | 213 | It("has idle connections", func() { 214 | Expect(connPool.Len()).To(Equal(minIdleConns)) 215 | Expect(connPool.IdleLen()).To(Equal(minIdleConns)) 216 | }) 217 | }) 218 | }) 219 | } 220 | 221 | Context("minIdleConns = 1", func() { 222 | BeforeEach(func() { 223 | minIdleConns = 1 224 | connPool = newConnPool() 225 | }) 226 | 227 | AfterEach(func() { 228 | connPool.Close() 229 | }) 230 | 231 | assert() 232 | }) 233 | 234 | Context("minIdleConns = 32", func() { 235 | BeforeEach(func() { 236 | minIdleConns = 32 237 | connPool = newConnPool() 238 | }) 239 | 240 | AfterEach(func() { 241 | connPool.Close() 242 | }) 243 | 244 | assert() 245 | }) 246 | }) 247 | 248 | var _ = Describe("conns reaper", func() { 249 | const idleTimeout = time.Minute 250 | const maxAge = time.Hour 251 | 252 | var connPool Pooler 253 | var conns, staleConns, closedConns []*Conn 254 | 255 | assert := func(typ string) { 256 | BeforeEach(func() { 257 | closedConns = nil 258 | connPool = New(Options{ 259 | Dialer: dummyDialer, 260 | PoolSize: 10, 261 | IdleTimeout: idleTimeout, 262 | MaxConnAge: maxAge, 263 | PoolTimeout: time.Second, 264 | IdleCheckFrequency: time.Hour, 265 | OnClose: func(cn *Conn) error { 266 | closedConns = append(closedConns, cn) 267 | return nil 268 | }, 269 | }) 270 | 271 | conns = nil 272 | 273 | // add stale connections 274 | staleConns = nil 275 | for i := 0; i < 3; i++ { 276 | cn, err := connPool.Get(context.Background()) 277 | Expect(err).NotTo(HaveOccurred()) 278 | switch typ { 279 | case "idle": 280 | cn.setUsedAt(time.Now().Add(-2 * idleTimeout)) 281 | case "aged": 282 | cn.createdAt = time.Now().Add(-2 * maxAge) 283 | } 284 | conns = append(conns, cn) 285 | staleConns = append(staleConns, cn) 286 | } 287 | 288 | // add fresh connections 289 | for i := 0; i < 3; i++ { 290 | cn, err := connPool.Get(context.Background()) 291 | Expect(err).NotTo(HaveOccurred()) 292 | conns = append(conns, cn) 293 | } 294 | 295 | for _, cn := range conns { 296 | connPool.Put(cn) 297 | } 298 | 299 | Expect(connPool.Len()).To(Equal(6)) 300 | Expect(connPool.IdleLen()).To(Equal(6)) 301 | 302 | n, err := connPool.(*ConnPool).reapStaleConns() 303 | Expect(err).NotTo(HaveOccurred()) 304 | Expect(n).To(Equal(3)) 305 | }) 306 | 307 | AfterEach(func() { 308 | _ = connPool.Close() 309 | Expect(connPool.Len()).To(Equal(0)) 310 | Expect(connPool.IdleLen()).To(Equal(0)) 311 | Expect(len(closedConns)).To(Equal(len(conns))) 312 | Expect(closedConns).To(ConsistOf(conns)) 313 | }) 314 | 315 | It("reaps stale connections", func() { 316 | Expect(connPool.Len()).To(Equal(3)) 317 | Expect(connPool.IdleLen()).To(Equal(3)) 318 | }) 319 | 320 | It("does not reap fresh connections", func() { 321 | n, err := connPool.(*ConnPool).reapStaleConns() 322 | Expect(err).NotTo(HaveOccurred()) 323 | Expect(n).To(Equal(0)) 324 | }) 325 | 326 | It("stale connections are closed", func() { 327 | Expect(len(closedConns)).To(Equal(len(staleConns))) 328 | Expect(closedConns).To(ConsistOf(staleConns)) 329 | }) 330 | 331 | It("pool is functional", func() { 332 | for j := 0; j < 3; j++ { 333 | var freeCns []*Conn 334 | for i := 0; i < 3; i++ { 335 | cn, err := connPool.Get(context.Background()) 336 | Expect(err).NotTo(HaveOccurred()) 337 | Expect(cn).NotTo(BeNil()) 338 | freeCns = append(freeCns, cn) 339 | } 340 | 341 | Expect(connPool.Len()).To(Equal(3)) 342 | Expect(connPool.IdleLen()).To(Equal(0)) 343 | 344 | cn, err := connPool.Get(context.Background()) 345 | Expect(err).NotTo(HaveOccurred()) 346 | Expect(cn).NotTo(BeNil()) 347 | conns = append(conns, cn) 348 | 349 | Expect(connPool.Len()).To(Equal(4)) 350 | Expect(connPool.IdleLen()).To(Equal(0)) 351 | 352 | connPool.Remove(cn) 353 | 354 | Expect(connPool.Len()).To(Equal(3)) 355 | Expect(connPool.IdleLen()).To(Equal(0)) 356 | 357 | for _, cn := range freeCns { 358 | connPool.Put(cn) 359 | } 360 | 361 | Expect(connPool.Len()).To(Equal(3)) 362 | Expect(connPool.IdleLen()).To(Equal(3)) 363 | } 364 | }) 365 | } 366 | 367 | assert("idle") 368 | assert("aged") 369 | }) 370 | 371 | var _ = Describe("race", func() { 372 | var connPool Pooler 373 | var C, N int 374 | 375 | BeforeEach(func() { 376 | C, N = 10, 1000 377 | if testing.Short() { 378 | C = 4 379 | N = 100 380 | } 381 | }) 382 | 383 | AfterEach(func() { 384 | connPool.Close() 385 | }) 386 | 387 | It("does not happen on Get, Put, and Remove", func() { 388 | connPool = New(Options{ 389 | Dialer: dummyDialer, 390 | PoolSize: 10, 391 | PoolTimeout: time.Minute, 392 | IdleTimeout: time.Millisecond, 393 | IdleCheckFrequency: time.Millisecond, 394 | }) 395 | 396 | perform(C, func(id int) { 397 | for i := 0; i < N; i++ { 398 | cn, err := connPool.Get(context.Background()) 399 | Expect(err).NotTo(HaveOccurred()) 400 | if err == nil { 401 | connPool.Put(cn) 402 | } 403 | } 404 | }, func(id int) { 405 | for i := 0; i < N; i++ { 406 | cn, err := connPool.Get(context.Background()) 407 | Expect(err).NotTo(HaveOccurred()) 408 | if err == nil { 409 | connPool.Remove(cn) 410 | } 411 | } 412 | }) 413 | }) 414 | }) 415 | --------------------------------------------------------------------------------