add connNum for db

shenli · Aug 20, 2014 · 2e11f54 · 2e11f54
1 parent 935b4b2
commit 2e11f54
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 16 deletions.
diff --git a/client/db.go b/client/db.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	. "github.com/siddontang/mixer/mysql"
 	"sync"
+	"sync/atomic"
 )
 
 type DB struct {
@@ -16,7 +17,9 @@ type DB struct {
 	db           string
 	maxIdleConns int
 
-	conns *list.List
+	idleConns *list.List
+
+	connNum int32
 }
 
 func Open(addr string, user string, password string, dbName string) (*DB, error) {
@@ -27,7 +30,8 @@ func Open(addr string, user string, password string, dbName string) (*DB, error)
 	db.password = password
 	db.db = dbName
 
-	db.conns = list.New()
+	db.idleConns = list.New()
+	db.connNum = 0
 
 	return db, nil
 }
@@ -45,10 +49,10 @@ func (db *DB) Close() error {
 	db.Lock()
 
 	for {
-		if db.conns.Len() > 0 {
-			v := db.conns.Back()
+		if db.idleConns.Len() > 0 {
+			v := db.idleConns.Back()
 			co := v.Value.(*Conn)
-			db.conns.Remove(v)
+			db.idleConns.Remove(v)
 
 			co.Close()
 
@@ -73,10 +77,18 @@ func (db *DB) Ping() error {
 	return err
 }
 
-func (db *DB) SetIdleConns(num int) {
+func (db *DB) SetMaxIdleConnNum(num int) {
 	db.maxIdleConns = num
 }
 
+func (db *DB) GetIdleConnNum() int {
+	return db.idleConns.Len()
+}
+
+func (db *DB) GetConnNum() int {
+	return int(db.connNum)
+}
+
 func (db *DB) newConn() (*Conn, error) {
 	co := new(Conn)
 
@@ -115,10 +127,10 @@ func (db *DB) tryReuse(co *Conn) error {
 
 func (db *DB) PopConn() (co *Conn, err error) {
 	db.Lock()
-	if db.conns.Len() > 0 {
-		v := db.conns.Front()
+	if db.idleConns.Len() > 0 {
+		v := db.idleConns.Front()
 		co = v.Value.(*Conn)
-		db.conns.Remove(v)
+		db.idleConns.Remove(v)
 	}
 	db.Unlock()
 
@@ -132,7 +144,11 @@ func (db *DB) PopConn() (co *Conn, err error) {
 		co.Close()
 	}
 
-	return db.newConn()
+	co, err = db.newConn()
+	if err == nil {
+		atomic.AddInt32(&db.connNum, 1)
+	}
+	return
 }
 
 func (db *DB) PushConn(co *Conn, err error) {
@@ -144,13 +160,13 @@ func (db *DB) PushConn(co *Conn, err error) {
 		if db.maxIdleConns > 0 {
 			db.Lock()
 
-			if db.conns.Len() >= db.maxIdleConns {
-				v := db.conns.Front()
+			if db.idleConns.Len() >= db.maxIdleConns {
+				v := db.idleConns.Front()
 				closeConn = v.Value.(*Conn)
-				db.conns.Remove(v)
+				db.idleConns.Remove(v)
 			}
 
-			db.conns.PushBack(co)
+			db.idleConns.PushBack(co)
 
 			db.Unlock()
 
@@ -161,6 +177,8 @@ func (db *DB) PushConn(co *Conn, err error) {
 	}
 
 	if closeConn != nil {
+		atomic.AddInt32(&db.connNum, -1)
+
 		closeConn.Close()
 	}
 }

diff --git a/proxy/node.go b/proxy/node.go
@@ -138,7 +138,7 @@ func (n *Node) openDB(addr string) (*client.DB, error) {
 		return nil, err
 	}
 
-	db.SetIdleConns(n.cfg.IdleConns)
+	db.SetMaxIdleConnNum(n.cfg.IdleConns)
 	return db, nil
 }
 

diff --git a/proxy/server_test.go b/proxy/server_test.go
@@ -109,7 +109,7 @@ func newTestDB(t *testing.T) *client.DB {
 			t.Fatal(err)
 		}
 
-		testDB.SetIdleConns(4)
+		testDB.SetMaxIdleConnNum(4)
 	}
 
 	testDBOnce.Do(f)