database/sql: Add DB.SetConnMaxLifetime

Long lived connections may make some DB operation difficult.
(e.g. retiring load balanced DB server.)
So SetConnMaxLifetime closes long lived connections.

It can be used to limit maximum idle time, too.
Closing idle connections reduces active connections while application is idle
and avoids connections are closed by server side (cause errBadConn while querying).

fixes #9851

Change-Id: I2e8e824219c1bee7f4b885d38ed96d11b7202b56
Reviewed-on: https://go-review.googlesource.com/6580
Run-TryBot: Brad Fitzpatrick <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>
Reviewed-by: Brad Fitzpatrick <[email protected]>

Author: methane

src/database/sql/sql.go

@@ -21,13 +21,17 @@ import (
 	"sort"
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
 var (
 	driversMu sync.RWMutex
 	drivers   = make(map[string]driver.Driver)
 )
 
+// nowFunc returns the current time; it's overridden in tests.
+var nowFunc = time.Now
+
 // Register makes a database driver available by the provided name.
 // If Register is called twice with the same name or if driver is nil,
 // it panics.
@@ -235,12 +239,14 @@ type DB struct {
 	// maybeOpenNewConnections sends on the chan (one send per needed connection)
 	// It is closed during db.Close(). The close tells the connectionOpener
 	// goroutine to exit.
-	openerCh chan struct{}
-	closed   bool
-	dep      map[finalCloser]depSet
-	lastPut  map[*driverConn]string // stacktrace of last conn's put; debug only
-	maxIdle  int                    // zero means defaultMaxIdleConns; negative means 0
-	maxOpen  int                    // <= 0 means unlimited
+	openerCh    chan struct{}
+	closed      bool
+	dep         map[finalCloser]depSet
+	lastPut     map[*driverConn]string // stacktrace of last conn's put; debug only
+	maxIdle     int                    // zero means defaultMaxIdleConns; negative means 0
+	maxOpen     int                    // <= 0 means unlimited
+	maxLifetime time.Duration          // maximum amount of time a connection may be reused
+	cleanerCh   chan struct{}
 }
 
 // connReuseStrategy determines how (*DB).conn returns database connections.
@@ -260,7 +266,8 @@ const (
 // interfaces returned via that Conn, such as calls on Tx, Stmt,
 // Result, Rows)
 type driverConn struct {
-	db *DB
+	db        *DB
+	createdAt time.Time
 
 	sync.Mutex  // guards following
 	ci          driver.Conn
@@ -284,6 +291,13 @@ func (dc *driverConn) removeOpenStmt(si driver.Stmt) {
 	delete(dc.openStmt, si)
 }
 
+func (dc *driverConn) expired(timeout time.Duration) bool {
+	if timeout <= 0 {
+		return false
+	}
+	return dc.createdAt.Add(timeout).Before(nowFunc())
+}
+
 func (dc *driverConn) prepareLocked(query string) (driver.Stmt, error) {
 	si, err := dc.ci.Prepare(query)
 	if err == nil {
@@ -506,6 +520,9 @@ func (db *DB) Close() error {
 		return nil
 	}
 	close(db.openerCh)
+	if db.cleanerCh != nil {
+		close(db.cleanerCh)
+	}
 	var err error
 	fns := make([]func() error, 0, len(db.freeConn))
 	for _, dc := range db.freeConn {
@@ -594,6 +611,84 @@ func (db *DB) SetMaxOpenConns(n int) {
 	}
 }
 
+// SetConnMaxLifetime sets the maximum amount of time a connection may be reused.
+//
+// Expired connections may be closed lazily before reuse.
+//
+// If d <= 0, connections are reused forever.
+func (db *DB) SetConnMaxLifetime(d time.Duration) {
+	if d < 0 {
+		d = 0
+	}
+	db.mu.Lock()
+	// wake cleaner up when lifetime is shortened.
+	if d > 0 && d < db.maxLifetime && db.cleanerCh != nil {
+		select {
+		case db.cleanerCh <- struct{}{}:
+		default:
+		}
+	}
+	db.maxLifetime = d
+	db.startCleanerLocked()
+	db.mu.Unlock()
+}
+
+// startCleanerLocked starts connectionCleaner if needed.
+func (db *DB) startCleanerLocked() {
+	if db.maxLifetime > 0 && db.numOpen > 0 && db.cleanerCh == nil {
+		db.cleanerCh = make(chan struct{}, 1)
+		go db.connectionCleaner(db.maxLifetime)
+	}
+}
+
+func (db *DB) connectionCleaner(d time.Duration) {
+	const minInterval = time.Second
+
+	if d < minInterval {
+		d = minInterval
+	}
+	t := time.NewTimer(d)
+
+	for {
+		select {
+		case <-t.C:
+		case <-db.cleanerCh: // maxLifetime was changed or db was closed.
+		}
+
+		db.mu.Lock()
+		d = db.maxLifetime
+		if db.closed || db.numOpen == 0 || d <= 0 {
+			db.cleanerCh = nil
+			db.mu.Unlock()
+			return
+		}
+
+		expiredSince := nowFunc().Add(-d)
+		var closing []*driverConn
+		for i := 0; i < len(db.freeConn); i++ {
+			c := db.freeConn[i]
+			if c.createdAt.Before(expiredSince) {
+				closing = append(closing, c)
+				last := len(db.freeConn) - 1
+				db.freeConn[i] = db.freeConn[last]
+				db.freeConn[last] = nil
+				db.freeConn = db.freeConn[:last]
+				i--
+			}
+		}
+		db.mu.Unlock()
+
+		for _, c := range closing {
+			c.Close()
+		}
+
+		if d < minInterval {
+			d = minInterval
+		}
+		t.Reset(d)
+	}
+}
+
 // DBStats contains database statistics.
 type DBStats struct {
 	// OpenConnections is the number of open connections to the database.
@@ -657,8 +752,9 @@ func (db *DB) openNewConnection() {
 		return
 	}
 	dc := &driverConn{
-		db: db,
-		ci: ci,
+		db:        db,
+		createdAt: nowFunc(),
+		ci:        ci,
 	}
 	if db.putConnDBLocked(dc, err) {
 		db.addDepLocked(dc, dc)
@@ -685,6 +781,7 @@ func (db *DB) conn(strategy connReuseStrategy) (*driverConn, error) {
 		db.mu.Unlock()
 		return nil, errDBClosed
 	}
+	lifetime := db.maxLifetime
 
 	// Prefer a free connection, if possible.
 	numFree := len(db.freeConn)
@@ -694,6 +791,10 @@ func (db *DB) conn(strategy connReuseStrategy) (*driverConn, error) {
 		db.freeConn = db.freeConn[:numFree-1]
 		conn.inUse = true
 		db.mu.Unlock()
+		if conn.expired(lifetime) {
+			conn.Close()
+			return nil, driver.ErrBadConn
+		}
 		return conn, nil
 	}
 
@@ -709,6 +810,10 @@ func (db *DB) conn(strategy connReuseStrategy) (*driverConn, error) {
 		if !ok {
 			return nil, errDBClosed
 		}
+		if ret.err == nil && ret.conn.expired(lifetime) {
+			ret.conn.Close()
+			return nil, driver.ErrBadConn
+		}
 		return ret.conn, ret.err
 	}
 
@@ -724,8 +829,9 @@ func (db *DB) conn(strategy connReuseStrategy) (*driverConn, error) {
 	}
 	db.mu.Lock()
 	dc := &driverConn{
-		db: db,
-		ci: ci,
+		db:        db,
+		createdAt: nowFunc(),
+		ci:        ci,
 	}
 	db.addDepLocked(dc, dc)
 	dc.inUse = true
@@ -835,6 +941,7 @@ func (db *DB) putConnDBLocked(dc *driverConn, err error) bool {
 		return true
 	} else if err == nil && !db.closed && db.maxIdleConnsLocked() > len(db.freeConn) {
 		db.freeConn = append(db.freeConn, dc)
+		db.startCleanerLocked()
 		return true
 	}
 	return false

src/database/sql/sql_test.go

@@ -142,6 +142,20 @@ func (db *DB) numFreeConns() int {
 	return len(db.freeConn)
 }
 
+// clearAllConns closes all connections in db.
+func (db *DB) clearAllConns(t *testing.T) {
+	db.SetMaxIdleConns(0)
+
+	if g, w := db.numFreeConns(), 0; g != w {
+		t.Errorf("free conns = %d; want %d", g, w)
+	}
+
+	if n := db.numDepsPollUntil(0, time.Second); n > 0 {
+		t.Errorf("number of dependencies = %d; expected 0", n)
+		db.dumpDeps(t)
+	}
+}
+
 func (db *DB) dumpDeps(t *testing.T) {
 	for fc := range db.dep {
 		db.dumpDep(t, 0, fc, map[finalCloser]bool{})
@@ -991,16 +1005,7 @@ func TestMaxOpenConns(t *testing.T) {
 
 	// Force the number of open connections to 0 so we can get an accurate
 	// count for the test
-	db.SetMaxIdleConns(0)
-
-	if g, w := db.numFreeConns(), 0; g != w {
-		t.Errorf("free conns = %d; want %d", g, w)
-	}
-
-	if n := db.numDepsPollUntil(0, time.Second); n > 0 {
-		t.Errorf("number of dependencies = %d; expected 0", n)
-		db.dumpDeps(t)
-	}
+	db.clearAllConns(t)
 
 	driver.mu.Lock()
 	opens0 := driver.openCount
@@ -1096,16 +1101,7 @@ func TestMaxOpenConns(t *testing.T) {
 		db.dumpDeps(t)
 	}
 
-	db.SetMaxIdleConns(0)
-
-	if g, w := db.numFreeConns(), 0; g != w {
-		t.Errorf("free conns = %d; want %d", g, w)
-	}
-
-	if n := db.numDepsPollUntil(0, time.Second); n > 0 {
-		t.Errorf("number of dependencies = %d; expected 0", n)
-		db.dumpDeps(t)
-	}
+	db.clearAllConns(t)
 }
 
 // Issue 9453: tests that SetMaxOpenConns can be lowered at runtime
@@ -1263,6 +1259,90 @@ func TestStats(t *testing.T) {
 	}
 }
 
+func TestConnMaxLifetime(t *testing.T) {
+	t0 := time.Unix(1000000, 0)
+	offset := time.Duration(0)
+
+	nowFunc = func() time.Time { return t0.Add(offset) }
+	defer func() { nowFunc = time.Now }()
+
+	db := newTestDB(t, "magicquery")
+	defer closeDB(t, db)
+
+	driver := db.driver.(*fakeDriver)
+
+	// Force the number of open connections to 0 so we can get an accurate
+	// count for the test
+	db.clearAllConns(t)
+
+	driver.mu.Lock()
+	opens0 := driver.openCount
+	closes0 := driver.closeCount
+	driver.mu.Unlock()
+
+	db.SetMaxIdleConns(10)
+	db.SetMaxOpenConns(10)
+
+	tx, err := db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	offset = time.Second
+	tx2, err := db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tx.Commit()
+	tx2.Commit()
+
+	driver.mu.Lock()
+	opens := driver.openCount - opens0
+	closes := driver.closeCount - closes0
+	driver.mu.Unlock()
+
+	if opens != 2 {
+		t.Errorf("opens = %d; want 2", opens)
+	}
+	if closes != 0 {
+		t.Errorf("closes = %d; want 0", closes)
+	}
+	if g, w := db.numFreeConns(), 2; g != w {
+		t.Errorf("free conns = %d; want %d", g, w)
+	}
+
+	// Expire first conn
+	offset = time.Second * 11
+	db.SetConnMaxLifetime(time.Second * 10)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tx, err = db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+	tx2, err = db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+	tx.Commit()
+	tx2.Commit()
+
+	driver.mu.Lock()
+	opens = driver.openCount - opens0
+	closes = driver.closeCount - closes0
+	driver.mu.Unlock()
+
+	if opens != 3 {
+		t.Errorf("opens = %d; want 3", opens)
+	}
+	if closes != 1 {
+		t.Errorf("closes = %d; want 1", closes)
+	}
+}
+
 // golang.org/issue/5323
 func TestStmtCloseDeps(t *testing.T) {
 	if testing.Short() {
@@ -1356,16 +1436,7 @@ func TestStmtCloseDeps(t *testing.T) {
 		db.dumpDeps(t)
 	}
 
-	db.SetMaxIdleConns(0)
-
-	if g, w := db.numFreeConns(), 0; g != w {
-		t.Errorf("free conns = %d; want %d", g, w)
-	}
-
-	if n := db.numDepsPollUntil(0, time.Second); n > 0 {
-		t.Errorf("number of dependencies = %d; expected 0", n)
-		db.dumpDeps(t)
-	}
+	db.clearAllConns(t)
 }
 
 // golang.org/issue/5046