chore: import upstream snapshot with attribution

internal/util/pool/default.go

@@ -0,0 +1,159 @@
+package pool
+
+import (
+	"errors"
+	"sync"
+	"time"
+
+	"github.com/google/uuid"
+
+	"go-micro.dev/v6/transport"
+)
+
+type pool struct {
+	tr transport.Transport
+
+	closeTimeout time.Duration
+	conns        map[string][]*poolConn
+	mu           sync.Mutex
+	size         int
+	ttl          time.Duration
+}
+
+type poolConn struct {
+	transport.Client
+
+	closeTimeout time.Duration
+	created      time.Time
+	id           string
+}
+
+func newPool(options Options) *pool {
+	return &pool{
+		size:         options.Size,
+		tr:           options.Transport,
+		ttl:          options.TTL,
+		closeTimeout: options.CloseTimeout,
+		conns:        make(map[string][]*poolConn),
+	}
+}
+
+func (p *pool) Close() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	var err error
+
+	for k, c := range p.conns {
+		for _, conn := range c {
+			if nerr := conn.close(); nerr != nil {
+				err = nerr
+			}
+		}
+
+		delete(p.conns, k)
+	}
+
+	return err
+}
+
+// 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(addr string, opts ...transport.DialOption) (Conn, error) {
+	p.mu.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 {
+			if err := conn.close(); err != nil {
+				p.mu.Unlock()
+				c, errConn := p.newConn(addr, opts)
+				if errConn != nil {
+					return nil, errConn
+				}
+				return c, err
+			}
+
+			continue
+		}
+
+		// We got a good conn, lets unlock and return it
+		p.mu.Unlock()
+
+		return conn, nil
+	}
+
+	p.mu.Unlock()
+
+	return p.newConn(addr, opts)
+}
+
+func (p *pool) newConn(addr string, opts []transport.DialOption) (Conn, error) {
+	// create new conn
+	c, err := p.tr.Dial(addr, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	return &poolConn{
+		Client:       c,
+		id:           uuid.New().String(),
+		closeTimeout: p.closeTimeout,
+		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).close()
+	}
+
+	// otherwise put it back for reuse
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	conns := p.conns[conn.Remote()]
+	if len(conns) >= p.size {
+		return conn.(*poolConn).close()
+	}
+
+	p.conns[conn.Remote()] = append(conns, conn.(*poolConn))
+
+	return nil
+}
+
+func (p *poolConn) close() error {
+	ch := make(chan error)
+	go func() {
+		defer close(ch)
+		ch <- p.Client.Close()
+	}()
+	t := time.NewTimer(p.closeTimeout)
+	var err error
+	select {
+	case <-t.C:
+		err = errors.New("unable to close in time")
+	case err = <-ch:
+		t.Stop()
+	}
+	return err
+}

internal/util/pool/default_test.go

@@ -0,0 +1,88 @@
+package pool
+
+import (
+	"testing"
+	"time"
+
+	"go-micro.dev/v6/transport"
+)
+
+func testPool(t *testing.T, size int, ttl time.Duration) {
+	// mock transport
+	tr := transport.NewMemoryTransport()
+
+	options := Options{
+		TTL:       ttl,
+		Size:      size,
+		Transport: tr,
+	}
+	// zero pool
+	p := newPool(options)
+
+	// listen
+	l, err := tr.Listen(":0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer l.Close()
+
+	// accept loop
+	go func() {
+		for {
+			if err := l.Accept(func(s transport.Socket) {
+				for {
+					var msg transport.Message
+					if err := s.Recv(&msg); err != nil {
+						return
+					}
+					if err := s.Send(&msg); err != nil {
+						return
+					}
+				}
+			}); err != nil {
+				return
+			}
+		}
+	}()
+
+	for i := 0; i < 10; i++ {
+		// get a conn
+		c, err := p.Get(l.Addr())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		msg := &transport.Message{
+			Body: []byte(`hello world`),
+		}
+
+		if err := c.Send(msg); err != nil {
+			t.Fatal(err)
+		}
+
+		var rcv transport.Message
+
+		if err := c.Recv(&rcv); err != nil {
+			t.Fatal(err)
+		}
+
+		if string(rcv.Body) != string(msg.Body) {
+			t.Fatalf("got %v, expected %v", rcv.Body, msg.Body)
+		}
+
+		// release the conn
+		p.Release(c, nil)
+
+		p.mu.Lock()
+		if i := len(p.conns[l.Addr()]); i > size {
+			p.mu.Unlock()
+			t.Fatalf("pool size %d is greater than expected %d", i, size)
+		}
+		p.mu.Unlock()
+	}
+}
+
+func TestClientPool(t *testing.T) {
+	testPool(t, 0, time.Minute)
+	testPool(t, 2, time.Minute)
+}

internal/util/pool/options.go

@@ -0,0 +1,40 @@
+package pool
+
+import (
+	"time"
+
+	"go-micro.dev/v6/transport"
+)
+
+type Options struct {
+	Transport    transport.Transport
+	TTL          time.Duration
+	CloseTimeout time.Duration
+	Size         int
+}
+
+type Option func(*Options)
+
+func Size(i int) Option {
+	return func(o *Options) {
+		o.Size = i
+	}
+}
+
+func Transport(t transport.Transport) Option {
+	return func(o *Options) {
+		o.Transport = t
+	}
+}
+
+func TTL(t time.Duration) Option {
+	return func(o *Options) {
+		o.TTL = t
+	}
+}
+
+func CloseTimeout(t time.Duration) Option {
+	return func(o *Options) {
+		o.CloseTimeout = t
+	}
+}

internal/util/pool/pool.go

@@ -0,0 +1,38 @@
+// Package pool is a connection pool
+package pool
+
+import (
+	"time"
+
+	"go-micro.dev/v6/transport"
+)
+
+// Pool is an interface for connection pooling.
+type Pool interface {
+	// Close the pool
+	Close() error
+	// Get a connection
+	Get(addr string, opts ...transport.DialOption) (Conn, error)
+	// Release the connection
+	Release(c Conn, status error) error
+}
+
+// Conn interface represents a pool connection.
+type Conn interface {
+	// unique id of connection
+	Id() string
+	// time it was created
+	Created() time.Time
+	// embedded connection
+	transport.Client
+}
+
+// NewPool will return a new pool object.
+func NewPool(opts ...Option) Pool {
+	var options Options
+	for _, o := range opts {
+		o(&options)
+	}
+
+	return newPool(options)
+}