Fix connection initialization.

internal/pool/bench_test.go

@@ -20,7 +20,7 @@ func benchmarkPoolGetPut(b *testing.B, poolSize int) {
 
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
-			conn, _, err := pool.Get()
+			conn, err := pool.Get()
 			if err != nil {
 				b.Fatalf("no error expected on pool.Get but received: %s", err.Error())
 			}
@@ -56,7 +56,7 @@ func benchmarkPoolGetReplace(b *testing.B, poolSize int) {
 
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
-			conn, _, err := pool.Get()
+			conn, err := pool.Get()
 			if err != nil {
 				b.Fatalf("no error expected on pool.Get but received: %s", err.Error())
 			}

internal/pool/conn.go

@@ -18,7 +18,9 @@ type Conn struct {
 	Rd      *bufio.Reader
 	Buf     []byte
 
-	UsedAt       time.Time
+	Inited bool
+	UsedAt time.Time
+
 	ReadTimeout  time.Duration
 	WriteTimeout time.Duration
 }
@@ -40,8 +42,12 @@ func (cn *Conn) Index() int {
 	return int(atomic.LoadInt32(&cn.idx))
 }
 
-func (cn *Conn) SetIndex(idx int) {
-	atomic.StoreInt32(&cn.idx, int32(idx))
+func (cn *Conn) SetIndex(newIdx int) int {
+	oldIdx := cn.Index()
+	if !atomic.CompareAndSwapInt32(&cn.idx, int32(oldIdx), int32(newIdx)) {
+		return -1
+	}
+	return oldIdx
 }
 
 func (cn *Conn) IsStale(timeout time.Duration) bool {
@@ -72,11 +78,6 @@ func (cn *Conn) RemoteAddr() net.Addr {
 	return cn.NetConn.RemoteAddr()
 }
 
-func (cn *Conn) Close() int {
-	idx := cn.Index()
-	if !atomic.CompareAndSwapInt32(&cn.idx, int32(idx), -1) {
-		return -1
-	}
-	_ = cn.NetConn.Close()
-	return idx
+func (cn *Conn) Close() error {
+	return cn.NetConn.Close()
 }

internal/pool/conn_list.go

@@ -43,7 +43,7 @@ func (l *connList) Add(cn *Conn) {
 	l.mu.Lock()
 	for i, c := range l.cns {
 		if c == nil {
-			cn.SetIndex(i)
+			cn.idx = int32(i)
 			l.cns[i] = cn
 			l.mu.Unlock()
 			return
@@ -76,6 +76,7 @@ func (l *connList) Close() error {
 		if c == nil {
 			continue
 		}
+		c.idx = -1
 		c.Close()
 	}
 	l.cns = nil

internal/pool/pool.go

@@ -32,7 +32,7 @@ type PoolStats struct {
 
 type Pooler interface {
 	First() *Conn
-	Get() (*Conn, bool, error)
+	Get() (*Conn, error)
 	Put(*Conn) error
 	Replace(*Conn, error) error
 	Len() int
@@ -146,7 +146,7 @@ func (p *ConnPool) dial() (net.Conn, error) {
 	return cn, nil
 }
 
-func (p *ConnPool) newConn() (*Conn, error) {
+func (p *ConnPool) NewConn() (*Conn, error) {
 	netConn, err := p.dial()
 	if err != nil {
 		return nil, err
@@ -155,42 +155,38 @@ func (p *ConnPool) newConn() (*Conn, error) {
 }
 
 // Get returns existed connection from the pool or creates a new one.
-func (p *ConnPool) Get() (cn *Conn, isNew bool, err error) {
+func (p *ConnPool) Get() (*Conn, error) {
 	if p.Closed() {
-		err = ErrClosed
-		return
+		return nil, ErrClosed
 	}
 
 	atomic.AddUint32(&p.stats.Requests, 1)
 
 	// Fetch first non-idle connection, if available.
-	if cn = p.First(); cn != nil {
+	if cn := p.First(); cn != nil {
 		atomic.AddUint32(&p.stats.Hits, 1)
-		return
+		return cn, nil
 	}
 
 	// Try to create a new one.
 	if p.conns.Reserve() {
-		isNew = true
-
-		cn, err = p.newConn()
+		cn, err := p.NewConn()
 		if err != nil {
 			p.conns.CancelReservation()
-			return
+			return nil, err
 		}
 		p.conns.Add(cn)
-		return
+		return cn, nil
 	}
 
 	// Otherwise, wait for the available connection.
 	atomic.AddUint32(&p.stats.Waits, 1)
-	if cn = p.wait(); cn != nil {
-		return
+	if cn := p.wait(); cn != nil {
+		return cn, nil
 	}
 
 	atomic.AddUint32(&p.stats.Timeouts, 1)
-	err = ErrPoolTimeout
-	return
+	return nil, ErrPoolTimeout
 }
 
 func (p *ConnPool) Put(cn *Conn) error {
@@ -205,7 +201,9 @@ func (p *ConnPool) Put(cn *Conn) error {
 }
 
 func (p *ConnPool) replace(cn *Conn) (*Conn, error) {
-	idx := cn.Close()
+	_ = cn.Close()
+
+	idx := cn.SetIndex(-1)
 	if idx == -1 {
 		return nil, errConnClosed
 	}
@@ -236,7 +234,9 @@ func (p *ConnPool) Replace(cn *Conn, reason error) error {
 }
 
 func (p *ConnPool) Remove(cn *Conn, reason error) error {
-	idx := cn.Close()
+	_ = cn.Close()
+
+	idx := cn.SetIndex(-1)
 	if idx == -1 {
 		return errConnClosed
 	}

internal/pool/pool_single.go

@@ -16,8 +16,8 @@ func (p *SingleConnPool) First() *Conn {
 	return p.cn
 }
 
-func (p *SingleConnPool) Get() (*Conn, bool, error) {
-	return p.cn, false, nil
+func (p *SingleConnPool) Get() (*Conn, error) {
+	return p.cn, nil
 }
 
 func (p *SingleConnPool) Put(cn *Conn) error {

internal/pool/pool_sticky.go

@@ -30,25 +30,23 @@ func (p *StickyConnPool) First() *Conn {
 	return cn
 }
 
-func (p *StickyConnPool) Get() (cn *Conn, isNew bool, err error) {
+func (p *StickyConnPool) Get() (*Conn, error) {
 	defer p.mx.Unlock()
 	p.mx.Lock()
 
 	if p.closed {
-		err = ErrClosed
-		return
+		return nil, ErrClosed
 	}
 	if p.cn != nil {
-		cn = p.cn
-		return
+		return p.cn, nil
 	}
 
-	cn, isNew, err = p.pool.Get()
+	cn, err := p.pool.Get()
 	if err != nil {
-		return
+		return nil, err
 	}
 	p.cn = cn
-	return
+	return cn, nil
 }
 
 func (p *StickyConnPool) put() (err error) {

internal/pool/pool_test.go

@@ -69,9 +69,8 @@ var _ = Describe("conns reapser", func() {
 			cn := connPool.First()
 			Expect(cn).To(BeNil())
 
-			cn, isNew, err := connPool.Get()
+			cn, err := connPool.Get()
 			Expect(err).NotTo(HaveOccurred())
-			Expect(isNew).To(BeTrue())
 			Expect(cn).NotTo(BeNil())
 
 			Expect(connPool.Len()).To(Equal(4))