database/sql: add SetConnMaxIdleTime

Allow removing a connection from the connection pool after
it has been idle for a period of time, without regard to the
total lifespan of the connection.

Fixes #25232

Change-Id: Icff157b906769a2d2d45c67525e04a72feb8d832
Reviewed-on: https://go-review.googlesource.com/c/go/+/145758
Run-TryBot: Daniel Theophanes <kardianos@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

2 files changed, 149 insertions, 31 deletions

diff --git a/src/database/sql/sql.go b/src/database/sql/sql.go
index 0f5bbc01c9..550b58753f 100644
--- a/src/database/sql/sql.go
+++ b/src/database/sql/sql.go
@@ -425,12 +425,14 @@ type DB 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
+	maxIdleCount      int                    // zero means defaultMaxIdleConns; negative means 0
 	maxOpen           int                    // <= 0 means unlimited
 	maxLifetime       time.Duration          // maximum amount of time a connection may be reused
+	maxIdleTime       time.Duration          // maximum amount of time a connection may be idle before being closed
 	cleanerCh         chan struct{}
 	waitCount         int64 // Total number of connections waited for.
-	maxIdleClosed     int64 // Total number of connections closed due to idle.
+	maxIdleClosed     int64 // Total number of connections closed due to idle count.
+	maxIdleTimeClosed int64 // Total number of connections closed due to idle time.
 	maxLifetimeClosed int64 // Total number of connections closed due to max free limit.
 
 	stop func() // stop cancels the connection opener and the session resetter.
@@ -465,8 +467,9 @@ type driverConn struct {
 
 	// guarded by db.mu
 	inUse      bool
-	onPut      []func() // code (with db.mu held) run when conn is next returned
-	dbmuClosed bool     // same as closed, but guarded by db.mu, for removeClosedStmtLocked
+	returnedAt time.Time // Time the connection was created or returned.
+	onPut      []func()  // code (with db.mu held) run when conn is next returned
+	dbmuClosed bool      // same as closed, but guarded by db.mu, for removeClosedStmtLocked
 }
 
 func (dc *driverConn) releaseConn(err error) {
@@ -839,7 +842,7 @@ func (db *DB) Close() error {
 const defaultMaxIdleConns = 2
 
 func (db *DB) maxIdleConnsLocked() int {
-	n := db.maxIdle
+	n := db.maxIdleCount
 	switch {
 	case n == 0:
 		// TODO(bradfitz): ask driver, if supported, for its default preference
@@ -851,6 +854,14 @@ func (db *DB) maxIdleConnsLocked() int {
 	}
 }
 
+func (db *DB) shortestIdleTimeLocked() time.Duration {
+	min := db.maxIdleTime
+	if min > db.maxLifetime {
+		min = db.maxLifetime
+	}
+	return min
+}
+
 // SetMaxIdleConns sets the maximum number of connections in the idle
 // connection pool.
 //
@@ -864,14 +875,14 @@ func (db *DB) maxIdleConnsLocked() int {
 func (db *DB) SetMaxIdleConns(n int) {
 	db.mu.Lock()
 	if n > 0 {
-		db.maxIdle = n
+		db.maxIdleCount = n
 	} else {
 		// No idle connections.
-		db.maxIdle = -1
+		db.maxIdleCount = -1
 	}
 	// Make sure maxIdle doesn't exceed maxOpen
 	if db.maxOpen > 0 && db.maxIdleConnsLocked() > db.maxOpen {
-		db.maxIdle = db.maxOpen
+		db.maxIdleCount = db.maxOpen
 	}
 	var closing []*driverConn
 	idleCount := len(db.freeConn)
@@ -912,13 +923,13 @@ func (db *DB) SetMaxOpenConns(n int) {
 //
 // Expired connections may be closed lazily before reuse.
 //
-// If d <= 0, connections are reused forever.
+// If d <= 0, connections are not closed due to a connection's age.
 func (db *DB) SetConnMaxLifetime(d time.Duration) {
 	if d < 0 {
 		d = 0
 	}
 	db.mu.Lock()
-	// wake cleaner up when lifetime is shortened.
+	// Wake cleaner up when lifetime is shortened.
 	if d > 0 && d < db.maxLifetime && db.cleanerCh != nil {
 		select {
 		case db.cleanerCh <- struct{}{}:
@@ -930,11 +941,34 @@ func (db *DB) SetConnMaxLifetime(d time.Duration) {
 	db.mu.Unlock()
 }
 
+// SetConnMaxIdleTime sets the maximum amount of time a connection may be idle.
+//
+// Expired connections may be closed lazily before reuse.
+//
+// If d <= 0, connections are not closed due to a connection's idle time.
+func (db *DB) SetConnMaxIdleTime(d time.Duration) {
+	if d < 0 {
+		d = 0
+	}
+	db.mu.Lock()
+	defer db.mu.Unlock()
+
+	// Wake cleaner up when idle time is shortened.
+	if d > 0 && d < db.maxIdleTime && db.cleanerCh != nil {
+		select {
+		case db.cleanerCh <- struct{}{}:
+		default:
+		}
+	}
+	db.maxIdleTime = d
+	db.startCleanerLocked()
+}
+
 // startCleanerLocked starts connectionCleaner if needed.
 func (db *DB) startCleanerLocked() {
-	if db.maxLifetime > 0 && db.numOpen > 0 && db.cleanerCh == nil {
+	if (db.maxLifetime > 0 || db.maxIdleTime > 0) && db.numOpen > 0 && db.cleanerCh == nil {
 		db.cleanerCh = make(chan struct{}, 1)
-		go db.connectionCleaner(db.maxLifetime)
+		go db.connectionCleaner(db.shortestIdleTimeLocked())
 	}
 }
 
@@ -953,15 +987,30 @@ func (db *DB) connectionCleaner(d time.Duration) {
 		}
 
 		db.mu.Lock()
-		d = db.maxLifetime
+
+		d = db.shortestIdleTimeLocked()
 		if db.closed || db.numOpen == 0 || d <= 0 {
 			db.cleanerCh = nil
 			db.mu.Unlock()
 			return
 		}
 
-		expiredSince := nowFunc().Add(-d)
-		var closing []*driverConn
+		closing := db.connectionCleanerRunLocked()
+		db.mu.Unlock()
+		for _, c := range closing {
+			c.Close()
+		}
+
+		if d < minInterval {
+			d = minInterval
+		}
+		t.Reset(d)
+	}
+}
+
+func (db *DB) connectionCleanerRunLocked() (closing []*driverConn) {
+	if db.maxLifetime > 0 {
+		expiredSince := nowFunc().Add(-db.maxLifetime)
 		for i := 0; i < len(db.freeConn); i++ {
 			c := db.freeConn[i]
 			if c.createdAt.Before(expiredSince) {
@@ -974,17 +1023,26 @@ func (db *DB) connectionCleaner(d time.Duration) {
 			}
 		}
 		db.maxLifetimeClosed += int64(len(closing))
-		db.mu.Unlock()
-
-		for _, c := range closing {
-			c.Close()
-		}
+	}
 
-		if d < minInterval {
-			d = minInterval
+	if db.maxIdleTime > 0 {
+		expiredSince := nowFunc().Add(-db.maxIdleTime)
+		var expiredCount int64
+		for i := 0; i < len(db.freeConn); i++ {
+			c := db.freeConn[i]
+			if db.maxIdleTime > 0 && c.returnedAt.Before(expiredSince) {
+				closing = append(closing, c)
+				expiredCount++
+				last := len(db.freeConn) - 1
+				db.freeConn[i] = db.freeConn[last]
+				db.freeConn[last] = nil
+				db.freeConn = db.freeConn[:last]
+				i--
+			}
 		}
-		t.Reset(d)
+		db.maxIdleTimeClosed += expiredCount
 	}
+	return
 }
 
 // DBStats contains database statistics.
@@ -1000,6 +1058,7 @@ type DBStats struct {
 	WaitCount         int64         // The total number of connections waited for.
 	WaitDuration      time.Duration // The total time blocked waiting for a new connection.
 	MaxIdleClosed     int64         // The total number of connections closed due to SetMaxIdleConns.
+	MaxIdleTimeClosed int64         // The total number of connections closed due to SetConnMaxIdleTime.
 	MaxLifetimeClosed int64         // The total number of connections closed due to SetConnMaxLifetime.
 }
 
@@ -1020,6 +1079,7 @@ func (db *DB) Stats() DBStats {
 		WaitCount:         db.waitCount,
 		WaitDuration:      time.Duration(wait),
 		MaxIdleClosed:     db.maxIdleClosed,
+		MaxIdleTimeClosed: db.maxIdleTimeClosed,
 		MaxLifetimeClosed: db.maxLifetimeClosed,
 	}
 	return stats
@@ -1097,9 +1157,10 @@ func (db *DB) openNewConnection(ctx context.Context) {
 		return
 	}
 	dc := &driverConn{
-		db:        db,
-		createdAt: nowFunc(),
-		ci:        ci,
+		db:         db,
+		createdAt:  nowFunc(),
+		returnedAt: nowFunc(),
+		ci:         ci,
 	}
 	if db.putConnDBLocked(dc, err) {
 		db.addDepLocked(dc, dc)
@@ -1177,7 +1238,7 @@ func (db *DB) conn(ctx context.Context, strategy connReuseStrategy) (*driverConn
 		db.waitCount++
 		db.mu.Unlock()
 
-		waitStart := time.Now()
+		waitStart := nowFunc()
 
 		// Timeout the connection request with the context.
 		select {
@@ -1235,10 +1296,11 @@ func (db *DB) conn(ctx context.Context, strategy connReuseStrategy) (*driverConn
 	}
 	db.mu.Lock()
 	dc := &driverConn{
-		db:        db,
-		createdAt: nowFunc(),
-		ci:        ci,
-		inUse:     true,
+		db:         db,
+		createdAt:  nowFunc(),
+		returnedAt: nowFunc(),
+		ci:         ci,
+		inUse:      true,
 	}
 	db.addDepLocked(dc, dc)
 	db.mu.Unlock()
@@ -1286,6 +1348,7 @@ func (db *DB) putConn(dc *driverConn, err error, resetSession bool) {
 		db.lastPut[dc] = stack()
 	}
 	dc.inUse = false
+	dc.returnedAt = nowFunc()
 
 	for _, fn := range dc.onPut {
 		fn()
diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 6f59260cda..a1437c46c9 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -3593,6 +3593,61 @@ func TestStatsMaxIdleClosedTen(t *testing.T) {
 	}
 }
 
+func TestMaxIdleTime(t *testing.T) {
+	list := []struct {
+		wantMaxIdleTime time.Duration
+		wantIdleClosed  int64
+		timeOffset      time.Duration
+	}{
+		{time.Nanosecond, 1, 10 * time.Millisecond},
+		{time.Hour, 0, 10 * time.Millisecond},
+	}
+	baseTime := time.Unix(0, 0)
+	defer func() {
+		nowFunc = time.Now
+	}()
+	for _, item := range list {
+		nowFunc = func() time.Time {
+			return baseTime
+		}
+		t.Run(fmt.Sprintf("%v", item.wantMaxIdleTime), func(t *testing.T) {
+			db := newTestDB(t, "people")
+			defer closeDB(t, db)
+
+			db.SetMaxOpenConns(1)
+			db.SetMaxIdleConns(1)
+			db.SetConnMaxIdleTime(item.wantMaxIdleTime)
+			db.SetConnMaxLifetime(0)
+
+			preMaxIdleClosed := db.Stats().MaxIdleTimeClosed
+
+			if err := db.Ping(); err != nil {
+				t.Fatal(err)
+			}
+
+			nowFunc = func() time.Time {
+				return baseTime.Add(item.timeOffset)
+			}
+
+			db.mu.Lock()
+			closing := db.connectionCleanerRunLocked()
+			db.mu.Unlock()
+			for _, c := range closing {
+				c.Close()
+			}
+			if g, w := int64(len(closing)), item.wantIdleClosed; g != w {
+				t.Errorf("got: %d; want %d closed conns", g, w)
+			}
+
+			st := db.Stats()
+			maxIdleClosed := st.MaxIdleTimeClosed - preMaxIdleClosed
+			if g, w := maxIdleClosed, item.wantIdleClosed; g != w {
+				t.Errorf(" got: %d; want %d max idle closed conns", g, w)
+			}
+		})
+	}
+}
+
 type nvcDriver struct {
 	fakeDriver
 	skipNamedValueCheck bool