From 3a2fe62f44a8a8513a087f75798425db7f9cc7bd Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Mon, 18 Mar 2013 15:33:04 -0700
Subject: [PATCH] database/sql: add DB.SetMaxIdleConns

Update #4805

R=golang-dev, r
CC=golang-dev
https://golang.org/cl/7634045
---
 src/pkg/database/sql/sql.go      | 47 ++++++++++++++++++++++++++++----
 src/pkg/database/sql/sql_test.go | 29 ++++++++++++++++++++
 2 files changed, 70 insertions(+), 6 deletions(-)

diff --git a/src/pkg/database/sql/sql.go b/src/pkg/database/sql/sql.go
index 556580eaecf6ea..8c6ffbc7008833 100644
--- a/src/pkg/database/sql/sql.go
+++ b/src/pkg/database/sql/sql.go
@@ -197,6 +197,7 @@ type DB struct {
 	dep       map[finalCloser]depSet
 	onConnPut map[*driverConn][]func() // code (with mu held) run when conn is next returned
 	lastPut   map[*driverConn]string   // stacktrace of last conn's put; debug only
+	maxIdle   int                      // zero means defaultMaxIdleConns; negative means 0
 }
 
 // driverConn wraps a driver.Conn with a mutex, to
@@ -332,11 +333,45 @@ func (db *DB) Close() error {
 	return err
 }
 
-func (db *DB) maxIdleConns() int {
-	const defaultMaxIdleConns = 2
-	// TODO(bradfitz): ask driver, if supported, for its default preference
-	// TODO(bradfitz): let users override?
-	return defaultMaxIdleConns
+const defaultMaxIdleConns = 2
+
+func (db *DB) maxIdleConnsLocked() int {
+	n := db.maxIdle
+	switch {
+	case n == 0:
+		// TODO(bradfitz): ask driver, if supported, for its default preference
+		return defaultMaxIdleConns
+	case n < 0:
+		return 0
+	default:
+		return n
+	}
+}
+
+// SetMaxIdleConns sets the maximum number of connections in the idle
+// connection pool.
+//
+// If n <= 0, no idle connections are retained.
+func (db *DB) SetMaxIdleConns(n int) {
+	db.mu.Lock()
+	defer db.mu.Unlock()
+	if n > 0 {
+		db.maxIdle = n
+	} else {
+		// No idle connections.
+		db.maxIdle = -1
+	}
+	for len(db.freeConn) > 0 && len(db.freeConn) > n {
+		nfree := len(db.freeConn)
+		dc := db.freeConn[nfree-1]
+		db.freeConn[nfree-1] = nil
+		db.freeConn = db.freeConn[:nfree-1]
+		go func() {
+			dc.Lock()
+			dc.ci.Close()
+			dc.Unlock()
+		}()
+	}
 }
 
 // conn returns a newly-opened or cached *driverConn
@@ -441,7 +476,7 @@ func (db *DB) putConn(dc *driverConn, err error) {
 	if putConnHook != nil {
 		putConnHook(db, dc)
 	}
-	if n := len(db.freeConn); !db.closed && n < db.maxIdleConns() {
+	if n := len(db.freeConn); !db.closed && n < db.maxIdleConnsLocked() {
 		db.freeConn = append(db.freeConn, dc)
 		db.mu.Unlock()
 		return
diff --git a/src/pkg/database/sql/sql_test.go b/src/pkg/database/sql/sql_test.go
index 5d3df721ed122f..2a9592e104c30a 100644
--- a/src/pkg/database/sql/sql_test.go
+++ b/src/pkg/database/sql/sql_test.go
@@ -761,3 +761,32 @@ func TestSimultaneousQueries(t *testing.T) {
 	}
 	defer r2.Close()
 }
+
+func TestMaxIdleConns(t *testing.T) {
+	db := newTestDB(t, "people")
+	defer closeDB(t, db)
+
+	tx, err := db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+	tx.Commit()
+	if got := len(db.freeConn); got != 1 {
+		t.Errorf("freeConns = %d; want 1", got)
+	}
+
+	db.SetMaxIdleConns(0)
+
+	if got := len(db.freeConn); got != 0 {
+		t.Errorf("freeConns after set to zero = %d; want 0", got)
+	}
+
+	tx, err = db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+	tx.Commit()
+	if got := len(db.freeConn); got != 0 {
+		t.Errorf("freeConns = %d; want 0", got)
+	}
+}