pool/pool.go

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