package pool import ( "code.squareroundforest.org/arpio/times" "errors" "time" ) type key struct { null bool value string } type state[R any] struct { keys []key items map[key][]R stats Stats forcedCheckPending bool deletedKeys int } type pool[R any] struct { state chan state[R] alloc func(string) (R, error) free func(r R) options Options } var nullKey = key{null: true} func makePool[R any](alloc func(string) (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]{items: make(map[key][]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 appendKeys(all []key, k key, n int) []key { for i := 0; i < n; i++ { all = append(all, k) } return all } func dropKeys(all []key, k key, n int) []key { first := len(all) last := first for n > 0 { if first == 0 && last == 0 { panic(errors.New("damaged pool: missing keys")) } if first > 0 && all[first-1] == k && last-first < n { first-- continue } if first == last { first-- last-- continue } var zero key copy(all[first:], all[last:]) for i := len(all) + first - last; i < len(all); i++ { all[i] = zero } all = all[:len(all)+first-last] n += first - last first-- last = first } return all } func resetItemsMap[R any](s state[R]) state[R] { if s.deletedKeys < 1<<12 { return s } if s.deletedKeys <= len(s.items)*2 { return s } next := make(map[key][]R) for k := range s.items { next[k] = s.items[k] } s.items = next s.deletedKeys = 0 return s } func (p pool[R]) getKey(k key, 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[k]) < n && p.alloc == nil { return nil, ErrNoItems } r := make([]R, n) if len(s.items[k]) <= n { copy(r, s.items[k]) } if len(s.items[k]) > n { // using the items from the end of the list: copy(r, s.items[k][len(s.items[k])-n:]) } if len(s.items[k]) < n { var err error event |= AllocateOperation for i := len(s.items[k]); i < n; i++ { s.stats.Alloc++ if r[i], err = p.alloc(k.value); err == nil { continue } s.items[k] = append(s.items[k], r[len(s.items[k]):i]...) s.keys = appendKeys(s.keys, k, i-len(s.items[k])) s.stats.Idle = len(s.keys) event |= AllocateError p.sendEvent(event, s.stats) return nil, err } } // drop enough keys from the end: d := n if len(s.items[k]) < n { d = len(s.items[k]) } s.keys = dropKeys(s.keys, k, d) if len(s.items[k]) <= n { var zero R for i := 0; i < len(s.items[k]); i++ { s.items[k][i] = zero } delete(s.items, k) s.deletedKeys++ } s = resetItemsMap(s) if len(s.items[k]) > n { var zero R for i := len(s.items[k]) - n; i < len(s.items[k]); i++ { s.items[k][i] = zero } s.items[k] = s.items[k][:len(s.items[k])-n] } s.stats.Idle = len(s.keys) s.stats.Active += n p.sendEvent(event, s.stats) return r, nil } func (p pool[R]) get(n int) ([]R, error) { return p.getKey(nullKey, n) } func (p pool[R]) putKey(k key, 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.keys)+len(r): s.items[k] = append(s.items[k], r...) s.keys = appendKeys(s.keys, k, len(r)) case t > len(s.keys) && t < len(s.keys)+len(r): event |= FreeOperation s = p.freeItems(s, r[t-len(s.keys):]) s.items[k] = append(s.items[k], r[:t-len(s.keys)]...) s.keys = appendKeys(s.keys, k, t-len(s.keys)) default: event |= FreeOperation s = p.freeItems(s, r) s = p.freeAboveTarget(s, t) } s.stats.Idle = len(s.keys) p.sendEvent(event, s.stats) if f > 0 { s = p.forcedCheck(s, f) } } func (p pool[R]) put(r ...R) { p.putKey(nullKey, r...) } 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.keys) { event |= FreeOperation s = p.freeAboveTarget(s, t) } s.stats.Idle = len(s.keys) p.sendEvent(event, s.stats) if f > 0 { s = p.forcedCheck(s, f) } } func (p pool[R]) loadKey(k key, r []R) { if len(r) == 0 { return } s := <-p.state defer func() { p.state <- s }() s.items[k] = append(s.items[k], r...) s.keys = appendKeys(s.keys, k, len(r)) s.stats.Idle = len(s.keys) s.stats.Load += len(r) p.options.Algo.Load(len(r)) event := LoadOperation p.sendEvent(event, s.stats) } func (p pool[R]) load(i []R) { p.loadKey(nullKey, i) } 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.keys) <= target { return s } var ( zeroKey key zeroItem R ) d := len(s.keys) - target r := make([]R, 0, d) for _, k := range s.keys[:d] { if len(s.items[k]) == 0 { panic(errors.New("damaged pool: missing item")) } r = append(r, s.items[k][0]) s.items[k][0] = zeroItem s.items[k] = s.items[k][1:] if len(s.items[k]) == 0 { delete(s.items, k) s.deletedKeys++ } } s = resetItemsMap(s) for i := 0; i < d; i++ { s.keys[i] = zeroKey } s.keys = s.keys[d:] return p.freeItems(s, r) } 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.keys) 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 for k := range s.items { s = p.freeItems(s, s.items[k]) } s.keys = nil s.items = make(map[key][]R) s.deletedKeys = 0 if s.stats.Free > prev { p.sendEvent(FreeOperation, s.stats) } }