bilibili-go-common/library/container/pool/slice.go

package pool

import (
	"context"
	"io"
	"sync"
	"time"
)

var _ Pool = &Slice{}

// Slice .
type Slice struct {
	// New is an application supplied function for creating and configuring a
	// item.
	//
	// The item returned from new must not be in a special state
	// (subscribed to pubsub channel, transaction started, ...).
	New  func(ctx context.Context) (io.Closer, error)
	stop func() // stop cancels the item opener.

	// mu protects fields defined below.
	mu           sync.Mutex
	freeItem     []*item
	itemRequests map[uint64]chan item
	nextRequest  uint64 // Next key to use in itemRequests.
	active       int    // number of opened and pending open items
	// Used to signal the need for new items
	// a goroutine running itemOpener() reads on this chan and
	// maybeOpenNewItems sends on the chan (one send per needed item)
	// It is closed during db.Close(). The close tells the itemOpener
	// goroutine to exit.
	openerCh  chan struct{}
	closed    bool
	cleanerCh chan struct{}

	// Config pool configuration
	conf *Config
}

// NewSlice creates a new pool.
func NewSlice(c *Config) *Slice {
	// check Config
	if c == nil || c.Active < c.Idle {
		panic("config nil or Idle Must <= Active")
	}
	ctx, cancel := context.WithCancel(context.Background())
	// new pool
	p := &Slice{
		conf:         c,
		stop:         cancel,
		itemRequests: make(map[uint64]chan item),
		openerCh:     make(chan struct{}, 1000000),
	}
	p.startCleanerLocked(time.Duration(c.IdleTimeout))

	go p.itemOpener(ctx)
	return p
}

// Reload reload config.
func (p *Slice) Reload(c *Config) error {
	p.mu.Lock()
	p.startCleanerLocked(time.Duration(c.IdleTimeout))
	p.setActive(c.Active)
	p.setIdle(c.Idle)
	p.conf = c
	p.mu.Unlock()
	return nil
}

// Get returns a newly-opened or cached *item.
func (p *Slice) Get(ctx context.Context) (io.Closer, error) {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return nil, ErrPoolClosed
	}
	idleTimeout := time.Duration(p.conf.IdleTimeout)
	// Prefer a free item, if possible.
	numFree := len(p.freeItem)
	for numFree > 0 {
		i := p.freeItem[0]
		copy(p.freeItem, p.freeItem[1:])
		p.freeItem = p.freeItem[:numFree-1]
		p.mu.Unlock()
		if i.expired(idleTimeout) {
			i.close()
			p.mu.Lock()
			p.release()
		} else {
			return i.c, nil
		}
		numFree = len(p.freeItem)
	}

	// Out of free items or we were asked not to use one. If we're not
	// allowed to open any more items, make a request and wait.
	if p.conf.Active > 0 && p.active >= p.conf.Active {
		// check WaitTimeout and return directly
		if p.conf.WaitTimeout == 0 && !p.conf.Wait {
			p.mu.Unlock()
			return nil, ErrPoolExhausted
		}
		// Make the item channel. It's buffered so that the
		// itemOpener doesn't block while waiting for the req to be read.
		req := make(chan item, 1)
		reqKey := p.nextRequestKeyLocked()
		p.itemRequests[reqKey] = req
		wt := p.conf.WaitTimeout
		p.mu.Unlock()

		// reset context timeout
		if wt > 0 {
			var cancel func()
			_, ctx, cancel = wt.Shrink(ctx)
			defer cancel()
		}
		// Timeout the item request with the context.
		select {
		case <-ctx.Done():
			// Remove the item request and ensure no value has been sent
			// on it after removing.
			p.mu.Lock()
			delete(p.itemRequests, reqKey)
			p.mu.Unlock()
			return nil, ctx.Err()
		case ret, ok := <-req:
			if !ok {
				return nil, ErrPoolClosed
			}
			if ret.expired(idleTimeout) {
				ret.close()
				p.mu.Lock()
				p.release()
			} else {
				return ret.c, nil
			}
		}
	}

	p.active++ // optimistically
	p.mu.Unlock()
	c, err := p.New(ctx)
	if err != nil {
		p.mu.Lock()
		p.release()
		p.mu.Unlock()
		return nil, err
	}
	return c, nil
}

// Put adds a item to the p's free pool.
// err is optionally the last error that occurred on this item.
func (p *Slice) Put(ctx context.Context, c io.Closer, forceClose bool) error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if forceClose {
		p.release()
		return c.Close()
	}
	added := p.putItemLocked(c)
	if !added {
		p.active--
		return c.Close()
	}
	return nil
}

// Satisfy a item or put the item in the idle pool and return true
// or return false.
// putItemLocked will satisfy a item if there is one, or it will
// return the *item to the freeItem list if err == nil and the idle
// item limit will not be exceeded.
// If err != nil, the value of i is ignored.
// If err == nil, then i must not equal nil.
// If a item was fulfilled or the *item was placed in the
// freeItem list, then true is returned, otherwise false is returned.
func (p *Slice) putItemLocked(c io.Closer) bool {
	if p.closed {
		return false
	}
	if p.conf.Active > 0 && p.active > p.conf.Active {
		return false
	}
	i := item{
		c:         c,
		createdAt: nowFunc(),
	}
	if l := len(p.itemRequests); l > 0 {
		var req chan item
		var reqKey uint64
		for reqKey, req = range p.itemRequests {
			break
		}
		delete(p.itemRequests, reqKey) // Remove from pending requests.
		req <- i
		return true
	} else if !p.closed && p.maxIdleItemsLocked() > len(p.freeItem) {
		p.freeItem = append(p.freeItem, &i)
		return true
	}
	return false
}

// Runs in a separate goroutine, opens new item when requested.
func (p *Slice) itemOpener(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-p.openerCh:
			p.openNewItem(ctx)
		}
	}
}

func (p *Slice) maybeOpenNewItems() {
	numRequests := len(p.itemRequests)
	if p.conf.Active > 0 {
		numCanOpen := p.conf.Active - p.active
		if numRequests > numCanOpen {
			numRequests = numCanOpen
		}
	}
	for numRequests > 0 {
		p.active++ // optimistically
		numRequests--
		if p.closed {
			return
		}
		p.openerCh <- struct{}{}
	}
}

// openNewItem one new item
func (p *Slice) openNewItem(ctx context.Context) {
	// maybeOpenNewConnctions has already executed p.active++ before it sent
	// on p.openerCh. This function must execute p.active-- if the
	// item fails or is closed before returning.
	c, err := p.New(ctx)
	p.mu.Lock()
	defer p.mu.Unlock()
	if err != nil {
		p.release()
		return
	}
	if !p.putItemLocked(c) {
		p.active--
		c.Close()
	}
}

// setIdle sets the maximum number of items in the idle
// item pool.
//
// If MaxOpenConns is greater than 0 but less than the new IdleConns
// then the new IdleConns will be reduced to match the MaxOpenConns limit
//
// If n <= 0, no idle items are retained.
func (p *Slice) setIdle(n int) {
	p.mu.Lock()
	if n > 0 {
		p.conf.Idle = n
	} else {
		// No idle items.
		p.conf.Idle = -1
	}
	// Make sure maxIdle doesn't exceed maxOpen
	if p.conf.Active > 0 && p.maxIdleItemsLocked() > p.conf.Active {
		p.conf.Idle = p.conf.Active
	}
	var closing []*item
	idleCount := len(p.freeItem)
	maxIdle := p.maxIdleItemsLocked()
	if idleCount > maxIdle {
		closing = p.freeItem[maxIdle:]
		p.freeItem = p.freeItem[:maxIdle]
	}
	p.mu.Unlock()
	for _, c := range closing {
		c.close()
	}
}

// setActive sets the maximum number of open items to the database.
//
// If IdleConns is greater than 0 and the new MaxOpenConns is less than
// IdleConns, then IdleConns will be reduced to match the new
// MaxOpenConns limit
//
// If n <= 0, then there is no limit on the number of open items.
// The default is 0 (unlimited).
func (p *Slice) setActive(n int) {
	p.mu.Lock()
	p.conf.Active = n
	if n < 0 {
		p.conf.Active = 0
	}
	syncIdle := p.conf.Active > 0 && p.maxIdleItemsLocked() > p.conf.Active
	p.mu.Unlock()
	if syncIdle {
		p.setIdle(n)
	}
}

// startCleanerLocked starts itemCleaner if needed.
func (p *Slice) startCleanerLocked(d time.Duration) {
	if d <= 0 {
		// if set 0, staleCleaner() will return directly
		return
	}
	if d < time.Duration(p.conf.IdleTimeout) && p.cleanerCh != nil {
		select {
		case p.cleanerCh <- struct{}{}:
		default:
		}
	}
	// run only one, clean stale items.
	if p.cleanerCh == nil {
		p.cleanerCh = make(chan struct{}, 1)
		go p.staleCleaner(time.Duration(p.conf.IdleTimeout))
	}
}

func (p *Slice) staleCleaner(d time.Duration) {
	const minInterval = 100 * time.Millisecond

	if d < minInterval {
		d = minInterval
	}
	t := time.NewTimer(d)

	for {
		select {
		case <-t.C:
		case <-p.cleanerCh: // maxLifetime was changed or db was closed.
		}
		p.mu.Lock()
		d = time.Duration(p.conf.IdleTimeout)
		if p.closed || d <= 0 {
			p.mu.Unlock()
			return
		}

		expiredSince := nowFunc().Add(-d)
		var closing []*item
		for i := 0; i < len(p.freeItem); i++ {
			c := p.freeItem[i]
			if c.createdAt.Before(expiredSince) {
				closing = append(closing, c)
				p.active--
				last := len(p.freeItem) - 1
				p.freeItem[i] = p.freeItem[last]
				p.freeItem[last] = nil
				p.freeItem = p.freeItem[:last]
				i--
			}
		}
		p.mu.Unlock()

		for _, c := range closing {
			c.close()
		}

		if d < minInterval {
			d = minInterval
		}
		t.Reset(d)
	}
}

// nextRequestKeyLocked returns the next item request key.
// It is assumed that nextRequest will not overflow.
func (p *Slice) nextRequestKeyLocked() uint64 {
	next := p.nextRequest
	p.nextRequest++
	return next
}

const defaultIdleItems = 2

func (p *Slice) maxIdleItemsLocked() int {
	n := p.conf.Idle
	switch {
	case n == 0:
		return defaultIdleItems
	case n < 0:
		return 0
	default:
		return n
	}
}

func (p *Slice) release() {
	p.active--
	p.maybeOpenNewItems()
}

// Close close pool.
func (p *Slice) Close() error {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return nil
	}
	if p.cleanerCh != nil {
		close(p.cleanerCh)
	}
	var err error
	for _, i := range p.freeItem {
		i.close()
	}
	p.freeItem = nil
	p.closed = true
	for _, req := range p.itemRequests {
		close(req)
	}
	p.mu.Unlock()
	p.stop()
	return err
}