common/utils/sync2/token_pool.go

package sync2

import "sync"

type TokenPool[T any] struct {
	tokens []T
	cond   *sync.Cond
	closed bool
}

func NewTokenPool[T any]() *TokenPool[T] {
	return &TokenPool[T]{
		tokens: make([]T, 0),
		cond:   sync.NewCond(new(sync.Mutex)),
	}
}

func NewTokenPoolWithTokens[T any](tokens ...T) *TokenPool[T] {
	b := NewTokenPool[T]()
	b.tokens = append(b.tokens, tokens...)
	return b
}

func (b *TokenPool[T]) Take() (T, bool) {
	b.cond.L.Lock()
	defer b.cond.L.Unlock()
	var ret T

	if b.closed {
		return ret, false
	}

	if len(b.tokens) == 0 {
		b.cond.Wait()

		if len(b.tokens) == 0 {
			return ret, false
		}
	}

	ret = b.tokens[0]
	b.tokens = b.tokens[1:]

	return ret, true
}

func (b *TokenPool[T]) Put(t T) {
	b.cond.L.Lock()
	b.tokens = append(b.tokens, t)
	b.cond.L.Unlock()
	b.cond.Signal()
}

func (b *TokenPool[T]) Close() {
	b.cond.L.Lock()
	b.closed = true
	b.cond.L.Unlock()

	b.cond.Broadcast()
}