pool.go

package utils

import (
	"sync"
	"time"
)

type PoolElem struct {
	Value interface{}
}

type Pool struct {
	freeList    []*PoolElem
	allocedList []*PoolElem
	mu          *sync.RWMutex
	cond        *sync.Cond
}

func NewPool() *Pool {
	return &Pool{
		freeList:    make([]*PoolElem, 0),
		allocedList: make([]*PoolElem, 0),
		mu:          &sync.RWMutex{},
		cond:        sync.NewCond(&sync.Mutex{}),
	}
}

func (p *Pool) freeListEmpty() bool {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return len(p.freeList) == 0
}

func (p *Pool) getImpl() *PoolElem {
	p.mu.Lock()
	defer p.mu.Unlock()
	obj := p.freeList[0]
	p.freeList = p.freeList[1:]
	p.allocedList = append(p.allocedList, obj)
	return obj
}

func find(elements []*PoolElem, elem *PoolElem) int {
	for i, el := range elements {
		if el == elem {
			return i
		}
	}
	return -1
}

func (p *Pool) putImpl(elem *PoolElem) {
	p.mu.Lock()
	defer p.mu.Unlock()
	index := find(p.allocedList, elem)
	if index >= 0 {
		p.allocedList = append(p.allocedList[0:index], p.allocedList[index+1:]...)
	}
	p.freeList = append(p.freeList, elem)
}

func (p *Pool) TryGet() *PoolElem {
	if p.freeListEmpty() {
		return nil
	}
	return p.getImpl()
}

func (p *Pool) Get() *PoolElem {
	p.cond.L.Lock()
	defer p.cond.L.Unlock()
	for p.freeListEmpty() {
		p.cond.Wait()
	}
	return p.getImpl()
}

func (p *Pool) Put(value interface{}) {
	switch v := value.(type) {
	case *PoolElem:
		p.putImpl(v)
	default:
		p.putImpl(&PoolElem{Value: value})
	}
	p.cond.Signal()
}

func (p *Pool) ForEach(fn func(value interface{})) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, elem := range p.freeList {
		fn(elem.Value)
	}
	for _, elem := range p.allocedList {
		fn(elem.Value)
	}
}

func (p *Pool) allocedListEmpty() bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	return len(p.allocedList) == 0
}

func (p *Pool) Wait() {
	p.cond.L.Lock()
	for !p.allocedListEmpty() {
		p.cond.Wait()
	}
	p.cond.L.Unlock()
}

func (p *Pool) WaitTimeout(timeoutTime int) chan bool {
	c := make(chan bool)
	expireTime := time.Now().Unix() + int64(timeoutTime)
	go func() {
		p.cond.L.Lock()
		defer p.cond.L.Unlock()
		for !p.allocedListEmpty() {
			p.cond.Wait()
			if time.Now().Unix() > expireTime {
				c <- true
				return
			}
		}
		c <- false
	}()
	return c
}