pool/pool.go

package pool

import "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()
	}

	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() (R, error) {
	s := <-p.state
	defer func() {
		p.state <- s
	}()

	var (
		r     R
		event EventType
		err   error
	)

	event |= GetOperation
	s.stats.Get++
	switch {
	case len(s.items) == 0 && p.alloc == nil:
		s.stats.Alloc++
		event |= AllocateOperation
		event |= AllocateError
		err = ErrEmptyPool
	case len(s.items) == 0:
		s.stats.Alloc++
		event |= AllocateOperation
		r, err = p.alloc()
		if err != nil {
			event |= AllocateError
		} else {
			s.stats.Active++
		}
	default:
		r, s.items = s.items[len(s.items)-1], s.items[:len(s.items)-1]
		if len(s.items) == 0 {
			s.items = nil
		}

		s.stats.Active++
		s.stats.Idle = len(s.items)
	}

	p.sendEvent(event, s.stats)
	return r, err
}

func (p pool[R]) put(r R) {
	s := <-p.state
	defer func() {
		p.state <- s
	}()

	var event EventType
	event |= PutOperation
	s.stats.Put++
	s.stats.Idle++

	// one may put in items that were allocated outside of the pool:
	if s.stats.Active > 0 {
		s.stats.Active--
	}

	t, f := p.options.Algo.Target(s.stats)
	switch {
	case t > len(s.items):
		s.items = append(s.items, r)
	default:
		event |= FreeOperation
		s = p.freeItems(s, t, r)
	}

	// fix provisioned idle count:
	s.stats.Idle = len(s.items)
	p.sendEvent(event, s.stats)
	if f > 0 {
		s = p.forcedCheck(s, f)
	}
}

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) {
		<-time.After(to)
		p.freeIdle()
	}(timeout)

	return s
}

func (p pool[R]) freeItems(s state[R], target int, additional ...R) state[R] {
	if len(s.items)+len(additional) <= target {
		return s
	}

	s.stats.Free += len(additional)
	if p.free == nil && len(s.items) <= target {
		return s
	}

	var f []R
	if p.free != nil {
		f = additional
		if len(s.items) > target {
			f = append(f, s.items[:len(s.items)-target]...)
		}
	}

	if len(s.items) > target {
		s.stats.Free += 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
		}
	}

	if p.free == nil {
		return s
	}

	for _, fi := range f {
		p.free(fi)
	}

	return s
}

func (p pool[R]) freeIdle() {
	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.freeItems(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)
	}
}

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, 0)
	if s.stats.Free > prev {
		p.sendEvent(FreeOperation, s.stats)
	}
}