database/sql: fix driver Conn refcounting with prepared statements

The refcounting of driver Conns was completedly busted and
would leak (be held open forever) with any reasonable
load. This was a significant regression from Go 1.0.

The core of this patch is removing one line:

     s.db.addDep(dc, s)

A database conn (dc) is a resource that be re-created any time
(but cached for speed) should not be held open forever with a
dependency refcount just because the Stmt (s) is alive (which
typically last for long periods of time, like forever).

The meat of the patch is new tests. In fixing the real issue,
a lot of tests then failed due to the fakedb_test.go's paranoia
about closing a fakeConn while it has open fakeStmts on it. I
could've ignored that, but that's been a problem in the past for
other bugs.

Instead, I now track per-Conn open statements and close them
when the the conn closes.  The proper way to do this would've
been making *driverStmt a finalCloser and using the dep mechanism,
but it was much more invasive. Added a TODO instead.

I'd like to give a way for drivers to opt-out of caring about
driver.Stmt closes before a driver.Conn close, but that's a TODO
for the future, and that TODO is added in this CL.

I know this is very late for Go 1.1, but database/sql is
currently nearly useless without this.

I'd like to believe all these database/sql bugs in the past
release cycle are the result of increased usage, number of
drivers, and good feedback from increasingly-capable Go
developers, and not the result of me sucking.  It's also hard
with all the real drivers being out-of-tree, so I'm having to
add more and more hooks to fakedb_test.go to simulate things
which real drivers end up doing.

Fixes #5323

R=golang-dev, snaury, gwenn.kahz, google, r
CC=golang-dev
https://golang.org/cl/8836045

1 files changed, 98 insertions, 27 deletions

diff --git a/src/pkg/database/sql/sql.go b/src/pkg/database/sql/sql.go
index 72289407c9..0646fb796f 100644
--- a/src/pkg/database/sql/sql.go
+++ b/src/pkg/database/sql/sql.go
@@ -207,13 +207,46 @@ type DB struct {
 type driverConn struct {
 	db *DB
 
-	sync.Mutex // guards following
-	ci         driver.Conn
-	closed     bool
+	sync.Mutex  // guards following
+	ci          driver.Conn
+	closed      bool
+	finalClosed bool // ci.Close has been called
+	openStmt    map[driver.Stmt]bool
 
 	// guarded by db.mu
-	inUse bool
-	onPut []func() // code (with db.mu held) run when conn is next returned
+	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 connIfFree
+}
+
+func (dc *driverConn) removeOpenStmt(si driver.Stmt) {
+	dc.Lock()
+	defer dc.Unlock()
+	delete(dc.openStmt, si)
+}
+
+func (dc *driverConn) prepareLocked(query string) (driver.Stmt, error) {
+	si, err := dc.ci.Prepare(query)
+	if err == nil {
+		// Track each driverConn's open statements, so we can close them
+		// before closing the conn.
+		//
+		// TODO(bradfitz): let drivers opt out of caring about
+		// stmt closes if the conn is about to close anyway? For now
+		// do the safe thing, in case stmts need to be closed.
+		//
+		// TODO(bradfitz): after Go 1.1, closing driver.Stmts
+		// should be moved to driverStmt, using unique
+		// *driverStmts everywhere (including from
+		// *Stmt.connStmt, instead of returning a
+		// driver.Stmt), using driverStmt as a pointer
+		// everywhere, and making it a finalCloser.
+		if dc.openStmt == nil {
+			dc.openStmt = make(map[driver.Stmt]bool)
+		}
+		dc.openStmt[si] = true
+	}
+	return si, err
 }
 
 // the dc.db's Mutex is held.
@@ -236,13 +269,27 @@ func (dc *driverConn) Close() error {
 	}
 	dc.closed = true
 	dc.Unlock() // not defer; removeDep finalClose calls may need to lock
-	return dc.db.removeDep(dc, dc)
+
+	// And now updates that require holding dc.mu.Lock.
+	dc.db.mu.Lock()
+	dc.dbmuClosed = true
+	fn := dc.db.removeDepLocked(dc, dc)
+	dc.db.mu.Unlock()
+	return fn()
 }
 
 func (dc *driverConn) finalClose() error {
 	dc.Lock()
+
+	for si := range dc.openStmt {
+		si.Close()
+	}
+	dc.openStmt = nil
+
 	err := dc.ci.Close()
 	dc.ci = nil
+	dc.finalClosed = true
+
 	dc.Unlock()
 	return err
 }
@@ -264,7 +311,8 @@ func (ds *driverStmt) Close() error {
 // depSet is a finalCloser's outstanding dependencies
 type depSet map[interface{}]bool // set of true bools
 
-// The finalCloser interface is used by (*DB).addDep and (*DB).get
+// The finalCloser interface is used by (*DB).addDep and related
+// dependency reference counting.
 type finalCloser interface {
 	// finalClose is called when the reference count of an object
 	// goes to zero. (*DB).mu is not held while calling it.
@@ -448,16 +496,26 @@ func (db *DB) conn() (*driverConn, error) {
 	return dc, nil
 }
 
-// connIfFree returns (wanted, true) if wanted is still a valid conn and
+var (
+	errConnClosed = errors.New("database/sql: internal sentinel error: conn is closed")
+	errConnBusy   = errors.New("database/sql: internal sentinel error: conn is busy")
+)
+
+// connIfFree returns (wanted, nil) if wanted is still a valid conn and
 // isn't in use.
 //
-// If wanted is valid but in use, connIfFree returns (wanted, false).
-// If wanted is invalid, connIfFre returns (nil, false).
-func (db *DB) connIfFree(wanted *driverConn) (conn *driverConn, ok bool) {
+// The error is errConnClosed if the connection if the requested connection
+// is invalid because it's been closed.
+//
+// The error is errConnBusy if the connection is in use.
+func (db *DB) connIfFree(wanted *driverConn) (*driverConn, error) {
 	db.mu.Lock()
 	defer db.mu.Unlock()
 	if wanted.inUse {
-		return conn, false
+		return nil, errConnBusy
+	}
+	if wanted.dbmuClosed {
+		return nil, errConnClosed
 	}
 	for i, conn := range db.freeConn {
 		if conn != wanted {
@@ -466,9 +524,14 @@ func (db *DB) connIfFree(wanted *driverConn) (conn *driverConn, ok bool) {
 		db.freeConn[i] = db.freeConn[len(db.freeConn)-1]
 		db.freeConn = db.freeConn[:len(db.freeConn)-1]
 		wanted.inUse = true
-		return wanted, true
+		return wanted, nil
 	}
-	return nil, false
+	// TODO(bradfitz): shouldn't get here. After Go 1.1, change this to:
+	// panic("connIfFree call requested a non-closed, non-busy, non-free conn")
+	// Which passes all the tests, but I'm too paranoid to include this
+	// late in Go 1.1.
+	// Instead, treat it like a busy connection:
+	return nil, errConnBusy
 }
 
 // putConnHook is a hook for testing.
@@ -485,7 +548,11 @@ func (db *DB) noteUnusedDriverStatement(c *driverConn, si driver.Stmt) {
 			si.Close()
 		})
 	} else {
-		si.Close()
+		c.Lock()
+		defer c.Unlock()
+		if !c.finalClosed {
+			si.Close()
+		}
 	}
 }
 
@@ -526,8 +593,6 @@ func (db *DB) putConn(dc *driverConn, err error) {
 		db.mu.Unlock()
 		return
 	}
-	// TODO: check to see if we need this Conn for any prepared
-	// statements which are still active?
 	db.mu.Unlock()
 
 	dc.Close()
@@ -560,7 +625,7 @@ func (db *DB) prepare(query string) (*Stmt, error) {
 		return nil, err
 	}
 	dc.Lock()
-	si, err := dc.ci.Prepare(query)
+	si, err := dc.prepareLocked(query)
 	dc.Unlock()
 	if err != nil {
 		db.putConn(dc, err)
@@ -572,7 +637,6 @@ func (db *DB) prepare(query string) (*Stmt, error) {
 		css:   []connStmt{{dc, si}},
 	}
 	db.addDep(stmt, stmt)
-	db.addDep(dc, stmt)
 	db.putConn(dc, nil)
 	return stmt, nil
 }
@@ -623,7 +687,6 @@ func (db *DB) exec(query string, args []interface{}) (res Result, err error) {
 		return nil, err
 	}
 	defer withLock(dc, func() { si.Close() })
-
 	return resultFromStatement(driverStmt{dc, si}, args...)
 }
 
@@ -1049,13 +1112,21 @@ func (s *Stmt) connStmt() (ci *driverConn, releaseConn func(error), si driver.St
 
 	var cs connStmt
 	match := false
-	for _, v := range s.css {
-		// TODO(bradfitz): lazily clean up entries in this
-		// list with dead conns while enumerating
-		if _, match = s.db.connIfFree(v.dc); match {
+	for i := 0; i < len(s.css); i++ {
+		v := s.css[i]
+		_, err := s.db.connIfFree(v.dc)
+		if err == nil {
+			match = true
 			cs = v
 			break
 		}
+		if err == errConnClosed {
+			// Lazily remove dead conn from our freelist.
+			s.css[i] = s.css[len(s.css)-1]
+			s.css = s.css[:len(s.css)-1]
+			i--
+		}
+
 	}
 	s.mu.Unlock()
 
@@ -1068,7 +1139,7 @@ func (s *Stmt) connStmt() (ci *driverConn, releaseConn func(error), si driver.St
 				return nil, nil, nil, err
 			}
 			dc.Lock()
-			si, err := dc.ci.Prepare(s.query)
+			si, err := dc.prepareLocked(s.query)
 			dc.Unlock()
 			if err == driver.ErrBadConn && i < 10 {
 				continue
@@ -1076,7 +1147,6 @@ func (s *Stmt) connStmt() (ci *driverConn, releaseConn func(error), si driver.St
 			if err != nil {
 				return nil, nil, nil, err
 			}
-			s.db.addDep(dc, s)
 			s.mu.Lock()
 			cs = connStmt{dc, si}
 			s.css = append(s.css, cs)
@@ -1195,6 +1265,7 @@ func (s *Stmt) Close() error {
 func (s *Stmt) finalClose() error {
 	for _, v := range s.css {
 		s.db.noteUnusedDriverStatement(v.dc, v.si)
+		v.dc.removeOpenStmt(v.si)
 		s.db.removeDep(v.dc, s)
 	}
 	s.css = nil
@@ -1389,7 +1460,7 @@ func (dr driverResult) RowsAffected() (int64, error) {
 }
 
 func stack() string {
-	var buf [1024]byte
+	var buf [2 << 10]byte
 	return string(buf[:runtime.Stack(buf[:], false)])
 }