pool/pool.go

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)
	}
}