pool.go 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. package pool
  2. import (
  3. "context"
  4. "runtime"
  5. "time"
  6. )
  7. type Factory[T any] func(context.Context) (T, error)
  8. type entry[T any] struct {
  9. elm T
  10. time time.Time
  11. }
  12. type Option[T any] func(*pool[T])
  13. // WithEvict set the evict callback
  14. func WithEvict[T any](cb func(T)) Option[T] {
  15. return func(p *pool[T]) {
  16. p.evict = cb
  17. }
  18. }
  19. // WithAge defined element max age (millisecond)
  20. func WithAge[T any](maxAge int64) Option[T] {
  21. return func(p *pool[T]) {
  22. p.maxAge = maxAge
  23. }
  24. }
  25. // WithSize defined max size of Pool
  26. func WithSize[T any](maxSize int) Option[T] {
  27. return func(p *pool[T]) {
  28. p.ch = make(chan *entry[T], maxSize)
  29. }
  30. }
  31. // Pool is for GC, see New for detail
  32. type Pool[T any] struct {
  33. *pool[T]
  34. }
  35. type pool[T any] struct {
  36. ch chan *entry[T]
  37. factory Factory[T]
  38. evict func(T)
  39. maxAge int64
  40. }
  41. func (p *pool[T]) GetContext(ctx context.Context) (T, error) {
  42. now := time.Now()
  43. for {
  44. select {
  45. case item := <-p.ch:
  46. elm := item
  47. if p.maxAge != 0 && now.Sub(item.time).Milliseconds() > p.maxAge {
  48. if p.evict != nil {
  49. p.evict(elm.elm)
  50. }
  51. continue
  52. }
  53. return elm.elm, nil
  54. default:
  55. return p.factory(ctx)
  56. }
  57. }
  58. }
  59. func (p *pool[T]) Get() (T, error) {
  60. return p.GetContext(context.Background())
  61. }
  62. func (p *pool[T]) Put(item T) {
  63. e := &entry[T]{
  64. elm: item,
  65. time: time.Now(),
  66. }
  67. select {
  68. case p.ch <- e:
  69. return
  70. default:
  71. // pool is full
  72. if p.evict != nil {
  73. p.evict(item)
  74. }
  75. return
  76. }
  77. }
  78. func recycle[T any](p *Pool[T]) {
  79. for item := range p.pool.ch {
  80. if p.pool.evict != nil {
  81. p.pool.evict(item.elm)
  82. }
  83. }
  84. }
  85. func New[T any](factory Factory[T], options ...Option[T]) *Pool[T] {
  86. p := &pool[T]{
  87. ch: make(chan *entry[T], 10),
  88. factory: factory,
  89. }
  90. for _, option := range options {
  91. option(p)
  92. }
  93. P := &Pool[T]{p}
  94. runtime.SetFinalizer(P, recycle[T])
  95. return P
  96. }