micro/util/pool/default.go

package pool

import (
	"context"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/unistack-org/micro/v3/network/transport"
)

type pool struct {
	size int
	ttl  time.Duration
	tr   transport.Transport

	sync.Mutex
	conns map[string][]*poolConn
}

type poolConn struct {
	transport.Client
	id      string
	created time.Time
}

func newPool(options Options) *pool {
	return &pool{
		size:  options.Size,
		tr:    options.Transport,
		ttl:   options.TTL,
		conns: make(map[string][]*poolConn),
	}
}

func (p *pool) Close() error {
	p.Lock()
	for k, c := range p.conns {
		for _, conn := range c {
			conn.Client.Close()
		}
		delete(p.conns, k)
	}
	p.Unlock()
	return nil
}

// NoOp the Close since we manage it
func (p *poolConn) Close() error {
	return nil
}

func (p *poolConn) Id() string {
	return p.id
}

func (p *poolConn) Created() time.Time {
	return p.created
}

func (p *pool) Get(ctx context.Context, addr string, opts ...transport.DialOption) (Conn, error) {
	p.Lock()
	conns := p.conns[addr]

	// while we have conns check age and then return one
	// otherwise we'll create a new conn
	for len(conns) > 0 {
		conn := conns[len(conns)-1]
		conns = conns[:len(conns)-1]
		p.conns[addr] = conns

		// if conn is old kill it and move on
		if d := time.Since(conn.Created()); d > p.ttl {
			conn.Client.Close()
			continue
		}

		// we got a good conn, lets unlock and return it
		p.Unlock()

		return conn, nil
	}

	p.Unlock()

	// create new conn
	c, err := p.tr.Dial(ctx, addr, opts...)
	if err != nil {
		return nil, err
	}
	return &poolConn{
		Client:  c,
		id:      uuid.New().String(),
		created: time.Now(),
	}, nil
}

func (p *pool) Release(conn Conn, err error) error {
	// don't store the conn if it has errored
	if err != nil {
		return conn.(*poolConn).Client.Close()
	}

	// otherwise put it back for reuse
	p.Lock()
	conns := p.conns[conn.Remote()]
	if len(conns) >= p.size {
		p.Unlock()
		return conn.(*poolConn).Client.Close()
	}
	p.conns[conn.Remote()] = append(conns, conn.(*poolConn))
	p.Unlock()

	return nil
}