github.com/jackc/puddle@v1.3.0/pool.go (about) 1 package puddle 2 3 import ( 4 "context" 5 "errors" 6 "sync" 7 "time" 8 ) 9 10 const ( 11 resourceStatusConstructing = 0 12 resourceStatusIdle = iota 13 resourceStatusAcquired = iota 14 resourceStatusHijacked = iota 15 ) 16 17 // ErrClosedPool occurs on an attempt to acquire a connection from a closed pool 18 // or a pool that is closed while the acquire is waiting. 19 var ErrClosedPool = errors.New("closed pool") 20 21 // ErrNotAvailable occurs on an attempt to acquire a resource from a pool 22 // that is at maximum capacity and has no available resources. 23 var ErrNotAvailable = errors.New("resource not available") 24 25 // Constructor is a function called by the pool to construct a resource. 26 type Constructor func(ctx context.Context) (res interface{}, err error) 27 28 // Destructor is a function called by the pool to destroy a resource. 29 type Destructor func(res interface{}) 30 31 // Resource is the resource handle returned by acquiring from the pool. 32 type Resource struct { 33 value interface{} 34 pool *Pool 35 creationTime time.Time 36 lastUsedNano int64 37 status byte 38 } 39 40 // Value returns the resource value. 41 func (res *Resource) Value() interface{} { 42 if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) { 43 panic("tried to access resource that is not acquired or hijacked") 44 } 45 return res.value 46 } 47 48 // Release returns the resource to the pool. res must not be subsequently used. 49 func (res *Resource) Release() { 50 if res.status != resourceStatusAcquired { 51 panic("tried to release resource that is not acquired") 52 } 53 res.pool.releaseAcquiredResource(res, nanotime()) 54 } 55 56 // ReleaseUnused returns the resource to the pool without updating when it was last used used. i.e. LastUsedNanotime 57 // will not change. res must not be subsequently used. 58 func (res *Resource) ReleaseUnused() { 59 if res.status != resourceStatusAcquired { 60 panic("tried to release resource that is not acquired") 61 } 62 res.pool.releaseAcquiredResource(res, res.lastUsedNano) 63 } 64 65 // Destroy returns the resource to the pool for destruction. res must not be 66 // subsequently used. 67 func (res *Resource) Destroy() { 68 if res.status != resourceStatusAcquired { 69 panic("tried to destroy resource that is not acquired") 70 } 71 go res.pool.destroyAcquiredResource(res) 72 } 73 74 // Hijack assumes ownership of the resource from the pool. Caller is responsible 75 // for cleanup of resource value. 76 func (res *Resource) Hijack() { 77 if res.status != resourceStatusAcquired { 78 panic("tried to hijack resource that is not acquired") 79 } 80 res.pool.hijackAcquiredResource(res) 81 } 82 83 // CreationTime returns when the resource was created by the pool. 84 func (res *Resource) CreationTime() time.Time { 85 if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) { 86 panic("tried to access resource that is not acquired or hijacked") 87 } 88 return res.creationTime 89 } 90 91 // LastUsedNanotime returns when Release was last called on the resource measured in nanoseconds from an arbitrary time 92 // (a monotonic time). Returns creation time if Release has never been called. This is only useful to compare with 93 // other calls to LastUsedNanotime. In almost all cases, IdleDuration should be used instead. 94 func (res *Resource) LastUsedNanotime() int64 { 95 if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) { 96 panic("tried to access resource that is not acquired or hijacked") 97 } 98 99 return res.lastUsedNano 100 } 101 102 // IdleDuration returns the duration since Release was last called on the resource. This is equivalent to subtracting 103 // LastUsedNanotime to the current nanotime. 104 func (res *Resource) IdleDuration() time.Duration { 105 if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) { 106 panic("tried to access resource that is not acquired or hijacked") 107 } 108 109 return time.Duration(nanotime() - res.lastUsedNano) 110 } 111 112 // Pool is a concurrency-safe resource pool. 113 type Pool struct { 114 cond *sync.Cond 115 destructWG *sync.WaitGroup 116 117 allResources []*Resource 118 idleResources []*Resource 119 120 constructor Constructor 121 destructor Destructor 122 maxSize int32 123 124 acquireCount int64 125 acquireDuration time.Duration 126 emptyAcquireCount int64 127 canceledAcquireCount int64 128 129 closed bool 130 } 131 132 // NewPool creates a new pool. Panics if maxSize is less than 1. 133 func NewPool(constructor Constructor, destructor Destructor, maxSize int32) *Pool { 134 if maxSize < 1 { 135 panic("maxSize is less than 1") 136 } 137 138 return &Pool{ 139 cond: sync.NewCond(new(sync.Mutex)), 140 destructWG: &sync.WaitGroup{}, 141 maxSize: maxSize, 142 constructor: constructor, 143 destructor: destructor, 144 } 145 } 146 147 // Close destroys all resources in the pool and rejects future Acquire calls. 148 // Blocks until all resources are returned to pool and destroyed. 149 func (p *Pool) Close() { 150 p.cond.L.Lock() 151 if p.closed { 152 p.cond.L.Unlock() 153 return 154 } 155 p.closed = true 156 157 for _, res := range p.idleResources { 158 p.allResources = removeResource(p.allResources, res) 159 go p.destructResourceValue(res.value) 160 } 161 p.idleResources = nil 162 p.cond.L.Unlock() 163 164 // Wake up all go routines waiting for a resource to be returned so they can terminate. 165 p.cond.Broadcast() 166 167 p.destructWG.Wait() 168 } 169 170 // Stat is a snapshot of Pool statistics. 171 type Stat struct { 172 constructingResources int32 173 acquiredResources int32 174 idleResources int32 175 maxResources int32 176 acquireCount int64 177 acquireDuration time.Duration 178 emptyAcquireCount int64 179 canceledAcquireCount int64 180 } 181 182 // TotalResource returns the total number of resources currently in the pool. 183 // The value is the sum of ConstructingResources, AcquiredResources, and 184 // IdleResources. 185 func (s *Stat) TotalResources() int32 { 186 return s.constructingResources + s.acquiredResources + s.idleResources 187 } 188 189 // ConstructingResources returns the number of resources with construction in progress in 190 // the pool. 191 func (s *Stat) ConstructingResources() int32 { 192 return s.constructingResources 193 } 194 195 // AcquiredResources returns the number of currently acquired resources in the pool. 196 func (s *Stat) AcquiredResources() int32 { 197 return s.acquiredResources 198 } 199 200 // IdleResources returns the number of currently idle resources in the pool. 201 func (s *Stat) IdleResources() int32 { 202 return s.idleResources 203 } 204 205 // MaxResources returns the maximum size of the pool. 206 func (s *Stat) MaxResources() int32 { 207 return s.maxResources 208 } 209 210 // AcquireCount returns the cumulative count of successful acquires from the pool. 211 func (s *Stat) AcquireCount() int64 { 212 return s.acquireCount 213 } 214 215 // AcquireDuration returns the total duration of all successful acquires from 216 // the pool. 217 func (s *Stat) AcquireDuration() time.Duration { 218 return s.acquireDuration 219 } 220 221 // EmptyAcquireCount returns the cumulative count of successful acquires from the pool 222 // that waited for a resource to be released or constructed because the pool was 223 // empty. 224 func (s *Stat) EmptyAcquireCount() int64 { 225 return s.emptyAcquireCount 226 } 227 228 // CanceledAcquireCount returns the cumulative count of acquires from the pool 229 // that were canceled by a context. 230 func (s *Stat) CanceledAcquireCount() int64 { 231 return s.canceledAcquireCount 232 } 233 234 // Stat returns the current pool statistics. 235 func (p *Pool) Stat() *Stat { 236 p.cond.L.Lock() 237 s := &Stat{ 238 maxResources: p.maxSize, 239 acquireCount: p.acquireCount, 240 emptyAcquireCount: p.emptyAcquireCount, 241 canceledAcquireCount: p.canceledAcquireCount, 242 acquireDuration: p.acquireDuration, 243 } 244 245 for _, res := range p.allResources { 246 switch res.status { 247 case resourceStatusConstructing: 248 s.constructingResources += 1 249 case resourceStatusIdle: 250 s.idleResources += 1 251 case resourceStatusAcquired: 252 s.acquiredResources += 1 253 } 254 } 255 256 p.cond.L.Unlock() 257 return s 258 } 259 260 // Acquire gets a resource from the pool. If no resources are available and the pool 261 // is not at maximum capacity it will create a new resource. If the pool is at 262 // maximum capacity it will block until a resource is available. ctx can be used 263 // to cancel the Acquire. 264 func (p *Pool) Acquire(ctx context.Context) (*Resource, error) { 265 startNano := nanotime() 266 if doneChan := ctx.Done(); doneChan != nil { 267 select { 268 case <-ctx.Done(): 269 p.cond.L.Lock() 270 p.canceledAcquireCount += 1 271 p.cond.L.Unlock() 272 return nil, ctx.Err() 273 default: 274 } 275 } 276 277 p.cond.L.Lock() 278 279 emptyAcquire := false 280 281 for { 282 if p.closed { 283 p.cond.L.Unlock() 284 return nil, ErrClosedPool 285 } 286 287 // If a resource is available now 288 if len(p.idleResources) > 0 { 289 res := p.idleResources[len(p.idleResources)-1] 290 p.idleResources[len(p.idleResources)-1] = nil // Avoid memory leak 291 p.idleResources = p.idleResources[:len(p.idleResources)-1] 292 res.status = resourceStatusAcquired 293 if emptyAcquire { 294 p.emptyAcquireCount += 1 295 } 296 p.acquireCount += 1 297 p.acquireDuration += time.Duration(nanotime() - startNano) 298 p.cond.L.Unlock() 299 return res, nil 300 } 301 302 emptyAcquire = true 303 304 // If there is room to create a resource do so 305 if len(p.allResources) < int(p.maxSize) { 306 res := &Resource{pool: p, creationTime: time.Now(), lastUsedNano: nanotime(), status: resourceStatusConstructing} 307 p.allResources = append(p.allResources, res) 308 p.destructWG.Add(1) 309 p.cond.L.Unlock() 310 311 // we create the resource in the background because the constructor might 312 // outlive the context and we want to continue constructing it as long as 313 // necessary but the acquire should be cancelled when the context is cancelled 314 // see: https://github.com/jackc/pgx/issues/1287 and https://github.com/jackc/pgx/issues/1259 315 constructErrCh := make(chan error) 316 go func() { 317 value, err := p.constructResourceValue(ctx) 318 p.cond.L.Lock() 319 if err != nil { 320 p.allResources = removeResource(p.allResources, res) 321 p.destructWG.Done() 322 323 // we can't use default here in case we get here before the caller is 324 // in the select 325 select { 326 case constructErrCh <- err: 327 case <-ctx.Done(): 328 p.canceledAcquireCount += 1 329 } 330 p.cond.L.Unlock() 331 p.cond.Signal() 332 return 333 } 334 res.value = value 335 336 // assume that we will acquire it 337 res.status = resourceStatusAcquired 338 // we can't use default here in case we get here before the caller is 339 // in the select 340 select { 341 case constructErrCh <- nil: 342 p.emptyAcquireCount += 1 343 p.acquireCount += 1 344 p.acquireDuration += time.Duration(nanotime() - startNano) 345 p.cond.L.Unlock() 346 // we don't call Signal here we didn't change any of the resource pools 347 case <-ctx.Done(): 348 p.canceledAcquireCount += 1 349 p.cond.L.Unlock() 350 // we don't call Signal here we didn't change any of the resopurce pools 351 // since we couldn't send the constructed resource to the acquire 352 // function that means the caller has stopped waiting and we should 353 // just put this resource back in the pool 354 p.releaseAcquiredResource(res, res.lastUsedNano) 355 } 356 }() 357 358 select { 359 case <-ctx.Done(): 360 return nil, ctx.Err() 361 case err := <-constructErrCh: 362 if err != nil { 363 return nil, err 364 } 365 // we don't call signal here because we didn't change the resource pools 366 // at all so waking anything else up won't help 367 return res, nil 368 } 369 } 370 371 if ctx.Done() == nil { 372 p.cond.Wait() 373 } else { 374 // Convert p.cond.Wait into a channel 375 waitChan := make(chan struct{}, 1) 376 go func() { 377 p.cond.Wait() 378 waitChan <- struct{}{} 379 }() 380 381 select { 382 case <-ctx.Done(): 383 // Allow goroutine waiting for signal to exit. Re-signal since we couldn't 384 // do anything with it. Another goroutine might be waiting. 385 go func() { 386 <-waitChan 387 p.cond.L.Unlock() 388 p.cond.Signal() 389 }() 390 391 p.cond.L.Lock() 392 p.canceledAcquireCount += 1 393 p.cond.L.Unlock() 394 return nil, ctx.Err() 395 case <-waitChan: 396 } 397 } 398 } 399 } 400 401 // TryAcquire gets a resource from the pool if one is immediately available. If not, it returns ErrNotAvailable. If no 402 // resources are available but the pool has room to grow, a resource will be created in the background. ctx is only 403 // used to cancel the background creation. 404 func (p *Pool) TryAcquire(ctx context.Context) (*Resource, error) { 405 p.cond.L.Lock() 406 defer p.cond.L.Unlock() 407 408 if p.closed { 409 return nil, ErrClosedPool 410 } 411 412 // If a resource is available now 413 if len(p.idleResources) > 0 { 414 res := p.idleResources[len(p.idleResources)-1] 415 p.idleResources[len(p.idleResources)-1] = nil // Avoid memory leak 416 p.idleResources = p.idleResources[:len(p.idleResources)-1] 417 p.acquireCount += 1 418 res.status = resourceStatusAcquired 419 return res, nil 420 } 421 422 if len(p.allResources) < int(p.maxSize) { 423 res := &Resource{pool: p, creationTime: time.Now(), lastUsedNano: nanotime(), status: resourceStatusConstructing} 424 p.allResources = append(p.allResources, res) 425 p.destructWG.Add(1) 426 427 go func() { 428 value, err := p.constructResourceValue(ctx) 429 defer p.cond.Signal() 430 p.cond.L.Lock() 431 defer p.cond.L.Unlock() 432 433 if err != nil { 434 p.allResources = removeResource(p.allResources, res) 435 p.destructWG.Done() 436 return 437 } 438 439 res.value = value 440 res.status = resourceStatusIdle 441 p.idleResources = append(p.idleResources, res) 442 }() 443 } 444 445 return nil, ErrNotAvailable 446 } 447 448 // AcquireAllIdle atomically acquires all currently idle resources. Its intended 449 // use is for health check and keep-alive functionality. It does not update pool 450 // statistics. 451 func (p *Pool) AcquireAllIdle() []*Resource { 452 p.cond.L.Lock() 453 if p.closed { 454 p.cond.L.Unlock() 455 return nil 456 } 457 458 for _, res := range p.idleResources { 459 res.status = resourceStatusAcquired 460 } 461 resources := p.idleResources // Swap out current slice 462 p.idleResources = nil 463 464 p.cond.L.Unlock() 465 return resources 466 } 467 468 // CreateResource constructs a new resource without acquiring it. 469 // It goes straight in the IdlePool. It does not check against maxSize. 470 // It can be useful to maintain warm resources under little load. 471 func (p *Pool) CreateResource(ctx context.Context) error { 472 p.cond.L.Lock() 473 if p.closed { 474 p.cond.L.Unlock() 475 return ErrClosedPool 476 } 477 p.cond.L.Unlock() 478 479 value, err := p.constructResourceValue(ctx) 480 if err != nil { 481 return err 482 } 483 484 res := &Resource{ 485 pool: p, 486 creationTime: time.Now(), 487 status: resourceStatusIdle, 488 value: value, 489 lastUsedNano: nanotime(), 490 } 491 p.destructWG.Add(1) 492 493 p.cond.L.Lock() 494 // If closed while constructing resource then destroy it and return an error 495 if p.closed { 496 go p.destructResourceValue(res.value) 497 p.cond.L.Unlock() 498 return ErrClosedPool 499 } 500 p.allResources = append(p.allResources, res) 501 p.idleResources = append(p.idleResources, res) 502 p.cond.L.Unlock() 503 504 return nil 505 } 506 507 // releaseAcquiredResource returns res to the the pool. 508 func (p *Pool) releaseAcquiredResource(res *Resource, lastUsedNano int64) { 509 p.cond.L.Lock() 510 511 if !p.closed { 512 res.lastUsedNano = lastUsedNano 513 res.status = resourceStatusIdle 514 p.idleResources = append(p.idleResources, res) 515 } else { 516 p.allResources = removeResource(p.allResources, res) 517 go p.destructResourceValue(res.value) 518 } 519 520 p.cond.L.Unlock() 521 p.cond.Signal() 522 } 523 524 // Remove removes res from the pool and closes it. If res is not part of the 525 // pool Remove will panic. 526 func (p *Pool) destroyAcquiredResource(res *Resource) { 527 p.destructResourceValue(res.value) 528 p.cond.L.Lock() 529 p.allResources = removeResource(p.allResources, res) 530 p.cond.L.Unlock() 531 p.cond.Signal() 532 } 533 534 func (p *Pool) hijackAcquiredResource(res *Resource) { 535 p.cond.L.Lock() 536 537 p.allResources = removeResource(p.allResources, res) 538 res.status = resourceStatusHijacked 539 p.destructWG.Done() // not responsible for destructing hijacked resources 540 541 p.cond.L.Unlock() 542 p.cond.Signal() 543 } 544 545 func removeResource(slice []*Resource, res *Resource) []*Resource { 546 for i := range slice { 547 if slice[i] == res { 548 slice[i] = slice[len(slice)-1] 549 slice[len(slice)-1] = nil // Avoid memory leak 550 return slice[:len(slice)-1] 551 } 552 } 553 554 panic("BUG: removeResource could not find res in slice") 555 } 556 557 func (p *Pool) constructResourceValue(ctx context.Context) (interface{}, error) { 558 return p.constructor(ctx) 559 } 560 561 func (p *Pool) destructResourceValue(value interface{}) { 562 p.destructor(value) 563 p.destructWG.Done() 564 }