client: add close method for client pool

client/common_test.go

@@ -25,6 +25,7 @@ func Test(t *testing.T) {
 	for _, seg := range segs {
 		Suite(&clientTestSuite{port: seg})
 		Suite(&connTestSuite{port: seg})
+		Suite(&poolTestSuite{port: seg})
 	}
 	TestingT(t)
 }

client/pool.go

@@ -5,6 +5,7 @@ import (
 	"math"
 	"math/rand"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/pingcap/errors"
@@ -42,6 +43,7 @@ type (
 		}
 
 		readyConnection chan Connection
+		closed          uint32
 	}
 
 	ConnectionStats struct {
@@ -143,6 +145,9 @@ func (pool *Pool) GetStats(stats *ConnectionStats) {
 // GetConn returns connection from the pool or create new
 func (pool *Pool) GetConn(ctx context.Context) (*Conn, error) {
 	for {
+		if atomic.LoadUint32(&pool.closed) == 1 {
+			return nil, errors.Errorf("failed get conn, pool closed")
+		}
 		connection, err := pool.getConnection(ctx)
 		if err != nil {
 			return nil, err
@@ -228,6 +233,9 @@ func (pool *Pool) newConnectionProducer() {
 	var err error
 
 	for {
+		if atomic.LoadUint32(&pool.closed) == 1 {
+			return
+		}
 		connection.conn = nil
 
 		pool.synchro.Lock()
@@ -298,6 +306,9 @@ func (pool *Pool) closeOldIdleConnections() {
 	ticker := time.NewTicker(5 * time.Second)
 
 	for range ticker.C {
+		if atomic.LoadUint32(&pool.closed) == 1 {
+			return
+		}
 		toPing = pool.getOldIdleConnections(toPing[:0])
 		if len(toPing) == 0 {
 			continue
@@ -475,3 +486,26 @@ func (pool *Pool) ping(conn *Conn) error {
 	}
 	return err
 }
+
+// Close only shutdown idle connections. we couldn't control the connection which not in the pool.
+// So before call Close, Call PutConn to put all connections that in use back to connection pool first.
+func (pool *Pool) Close() {
+	//already closed, return.
+	if !atomic.CompareAndSwapUint32(&pool.closed, 0, 1) {
+		return
+	}
+	//close idle connections
+	pool.synchro.Lock()
+	for _, connection := range pool.synchro.idleConnections {
+		pool.synchro.stats.TotalCount--
+		_ = connection.conn.Close()
+	}
+	pool.synchro.idleConnections = nil
+	pool.synchro.Unlock()
+	//close connection in ready
+	select {
+	case connection := <-pool.readyConnection:
+		pool.closeConn(connection.conn)
+	default:
+	}
+}

client/pool_test.go

@@ -0,0 +1,31 @@
+package client
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/go-mysql-org/go-mysql/test_util"
+	. "github.com/pingcap/check"
+	"github.com/siddontang/go-log/log"
+)
+
+type poolTestSuite struct {
+	port string
+}
+
+func (s poolTestSuite) TestPool_Close(c *C) {
+	addr := fmt.Sprintf("%s:%s", *test_util.MysqlHost, s.port)
+	pool := NewPool(log.Debugf, 5, 10, 5, addr, *testUser, *testPassword, "")
+	conn, err := pool.GetConn(context.Background())
+	c.Assert(err, IsNil)
+	err = conn.Ping()
+	c.Assert(err, IsNil)
+	pool.PutConn(conn)
+	pool.Close()
+	var poolStats ConnectionStats
+	pool.GetStats(&poolStats)
+	c.Assert(poolStats.TotalCount, Equals, 0)
+	c.Assert(pool.readyConnection, HasLen, 0)
+	_, err = pool.GetConn(context.Background())
+	c.Assert(err, NotNil)
+}