package pool import ( "code.squareroundforest.org/arpio/times" "time" ) type state[R any] struct { items []R stats Stats forcedCheckPending bool } type pool[R any] struct { state chan state[R] alloc func() (R, error) free func(r R) options Options } func makePool[R any](alloc func() (R, error), free func(r R), o Options) pool[R] { if o.Algo == nil { o.Algo = Adaptive() } if o.Testing.Clock == nil { o.Testing.Clock = times.Sys() } if sc, ok := o.Algo.(interface{ setClock(times.Clock) }); ok { sc.setClock(o.Testing.Clock) } s := make(chan state[R], 1) s <- state[R]{} return pool[R]{ state: s, alloc: alloc, free: free, options: o, } } func (p pool[R]) stats() Stats { s := <-p.state defer func() { p.state <- s }() return s.stats } func (p pool[R]) sendEvent(e EventType, s Stats) { if p.options.Events == nil { return } if p.options.EventMask&e == 0 { return } ev := Event{ Type: e, Stats: s, } select { case p.options.Events <- ev: default: } } func (p pool[R]) get(n int) ([]R, error) { if n <= 0 { return nil, nil } s := <-p.state defer func() { p.state <- s }() var event EventType event |= GetOperation s.stats.Get++ if len(s.items) < n && p.alloc == nil { return nil, ErrNoItems } r := make([]R, n) if len(s.items) <= n { copy(r, s.items) } if len(s.items) > n { // using the items from the end of the list: copy(r, s.items[len(s.items)-n:]) } if len(s.items) < n { var err error event |= AllocateOperation for i := len(s.items); i < n; i++ { s.stats.Alloc++ r[i], err = p.alloc() if err == nil { continue } s.items = append(s.items, r[len(s.items):i]...) s.stats.Idle = len(s.items) event |= AllocateError p.sendEvent(event, s.stats) return nil, err } } if len(s.items) <= n { var zero R for i := 0; i < len(s.items); i++ { s.items[i] = zero } s.items = nil } if len(s.items) > n { var zero R for i := len(s.items) - n; i < len(s.items); i++ { s.items[i] = zero } s.items = s.items[:len(s.items)-n] } s.stats.Idle = len(s.items) s.stats.Active += n p.sendEvent(event, s.stats) return r, nil } func (p pool[R]) put(r ...R) { if len(r) == 0 { return } s := <-p.state defer func() { p.state <- s }() var event EventType event |= PutOperation s.stats.Put++ s.stats.Idle += len(r) s.stats.Active -= len(r) // one may put in items that were allocated outside of the pool: if s.stats.Active < 0 { s.stats.Active = 0 } t, f := p.options.Algo.Target(s.stats) switch { case t >= len(s.items)+len(r): s.items = append(s.items, r...) case t > len(s.items) && t < len(s.items)+len(r): event |= FreeOperation s = p.freeItems(s, r[t-len(s.items):]) s.items = append(s.items, r[:t-len(s.items)]...) default: event |= FreeOperation s = p.freeItems(s, r) s = p.freeAboveTarget(s, t) } s.stats.Idle = len(s.items) p.sendEvent(event, s.stats) if f > 0 { s = p.forcedCheck(s, f) } } func (p pool[R]) drop(n int) { if n <= 0 { return } s := <-p.state defer func() { p.state <- s }() var event EventType event |= DropOperation s.stats.Drop += n s.stats.Active -= n if s.stats.Active < 0 { s.stats.Active = 0 } t, f := p.options.Algo.Target(s.stats) if t < len(s.items) { event |= FreeOperation s = p.freeAboveTarget(s, t) } s.stats.Idle = len(s.items) p.sendEvent(event, s.stats) if f > 0 { s = p.forcedCheck(s, f) } } func (p pool[R]) load(i []R) { if len(i) == 0 { return } s := <-p.state defer func() { p.state <- s }() s.items = append(s.items, i...) s.stats.Idle = len(s.items) s.stats.Load += len(i) p.options.Algo.Load(len(i)) event := LoadOperation p.sendEvent(event, s.stats) } func (p pool[R]) forcedCheck(s state[R], timeout time.Duration) state[R] { if s.forcedCheckPending { return s } s.forcedCheckPending = true go func(to time.Duration) { c := p.options.Testing.Clock.After(to) p.options.Testing.Bus.Signal("background-job-waiting") <-c p.checkAndFree() }(timeout) return s } func (p pool[R]) freeItems(s state[R], r []R) state[R] { s.stats.Free += len(r) if p.free == nil { return s } for _, ri := range r { p.free(ri) } return s } func (p pool[R]) freeAboveTarget(s state[R], target int) state[R] { if len(s.items) <= target { return s } f := s.items[:len(s.items)-target] var zero R for i := 0; i < len(s.items)-target; i++ { s.items[i] = zero } s.items = s.items[len(s.items)-target:] if len(s.items) == 0 { s.items = nil } return p.freeItems(s, f) } func (p pool[R]) checkAndFree() { s := <-p.state defer func() { p.state <- s }() s.forcedCheckPending = false t, f := p.options.Algo.Target(s.stats) prev := s.stats.Free s = p.freeAboveTarget(s, t) s.stats.Idle = len(s.items) if s.stats.Free > prev { p.sendEvent(FreeOperation, s.stats) } if f > 0 { s = p.forcedCheck(s, f) } p.options.Testing.Bus.Signal("free-idle-done") } func (p pool[R]) freePool() { s := <-p.state defer func() { p.state <- s }() s.stats.Idle = 0 prev := s.stats.Free s = p.freeItems(s, s.items) s.items = nil if s.stats.Free > prev { p.sendEvent(FreeOperation, s.stats) } }