code review 8836045: database/sql: fix driver Conn refcounting with prepared...

src/pkg/database/sql/sql.go

 // Copyright 2011 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
 // Package sql provides a generic interface around SQL (or SQL-like)
 // databases.
 //
 // The sql package must be used in conjunction with a database driver.
 // See http://golang.org/s/sqldrivers for a list of drivers.
 package sql
@@ skipping 189 matching lines @@
         maxIdle  int                    // zero means defaultMaxIdleConns; negative means 0
 }
 
 // driverConn wraps a driver.Conn with a mutex, to
 // be held during all calls into the Conn. (including any calls onto
 // interfaces returned via that Conn, such as calls on Tx, Stmt,
 // Result, Rows)
 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

[r, 2013/04/25 20:09:17]

s/-/ / (it's a verb not an adjective)

[bradfitz, 2013/04/25 21:46:10]

Done.

+                // 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.
 func (dc *driverConn) closeDBLocked() error {
         dc.Lock()
         if dc.closed {
                 dc.Unlock()
                 return errors.New("sql: duplicate driverConn close")
         }
         dc.closed = true
         dc.Unlock() // not defer; removeDep finalClose calls may need to lock
         return dc.db.removeDepLocked(dc, dc)()
 }
 
 func (dc *driverConn) Close() error {
         dc.Lock()
         if dc.closed {
                 dc.Unlock()
                 return errors.New("sql: duplicate driverConn close")
         }
         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
 }
 
 // driverStmt associates a driver.Stmt with the
 // *driverConn from which it came, so the driverConn's lock can be
 // held during calls.
 type driverStmt struct {
         sync.Locker // the *driverConn
         si          driver.Stmt
 }
 
 func (ds *driverStmt) Close() error {
         ds.Lock()
         defer ds.Unlock()
         return ds.si.Close()
 }
 
 // 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.
         finalClose() error
 }
 
 // addDep notes that x now depends on dep, and x's finalClose won't be
 // called until all of x's dependencies are removed with removeDep.
 func (db *DB) addDep(x finalCloser, dep interface{}) {
         //println(fmt.Sprintf("addDep(%T %p, %T %p)", x, x, dep, dep))
@@ skipping 163 matching lines @@
                 db: db,
                 ci: ci,
         }
         db.mu.Lock()
         db.addDepLocked(dc, dc)
         dc.inUse = true
         db.mu.Unlock()
         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 {
                         continue
                 }
                 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.
 var putConnHook func(*DB, *driverConn)
 
 // noteUnusedDriverStatement notes that si is no longer used and should
 // be closed whenever possible (when c is next not in use), unless c is
 // already closed.
 func (db *DB) noteUnusedDriverStatement(c *driverConn, si driver.Stmt) {
         db.mu.Lock()
         defer db.mu.Unlock()
         if c.inUse {
                 c.onPut = append(c.onPut, func() {
                         si.Close()
                 })
         } else {
-                si.Close()
+                c.Lock()
+                defer c.Unlock()
+                if !c.finalClosed {
+                        si.Close()
+                }
         }
 }
 
 // debugGetPut determines whether getConn & putConn calls' stack traces
 // are returned for more verbose crashes.
 const debugGetPut = false
 
 // putConn adds a connection to the db's free pool.
 // err is optionally the last error that occurred on this connection.
 func (db *DB) putConn(dc *driverConn, err error) {
@@ skipping 20 matching lines @@
                 return
         }
         if putConnHook != nil {
                 putConnHook(db, dc)
         }
         if n := len(db.freeConn); !db.closed && n < db.maxIdleConnsLocked() {
                 db.freeConn = append(db.freeConn, dc)
                 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()
 }
 
 // Prepare creates a prepared statement for later queries or executions.
 // Multiple queries or executions may be run concurrently from the
 // returned statement.
 func (db *DB) Prepare(query string) (*Stmt, error) {
         var stmt *Stmt
@@ skipping 12 matching lines @@
         // driver.Preparer interface and call that instead, if so,
         // otherwise we make a prepared statement that's bound
         // to a connection, and to execute this prepared statement
         // we either need to use this connection (if it's free), else
         // get a new connection + re-prepare + execute on that one.
         dc, err := db.conn()
         if err != nil {
                 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)
                 return nil, err
         }
         stmt := &Stmt{
                 db:    db,
                 query: query,
                 css:   []connStmt{{dc, si}},
         }
         db.addDep(stmt, stmt)
-        db.addDep(dc, stmt)
         db.putConn(dc, nil)
         return stmt, nil
 }
 
 // Exec executes a query without returning any rows.
 // The args are for any placeholder parameters in the query.
 func (db *DB) Exec(query string, args ...interface{}) (Result, error) {
         var res Result
         var err error
         for i := 0; i < 10; i++ {
@@ skipping 30 matching lines @@
                 }
         }
 
         dc.Lock()
         si, err := dc.ci.Prepare(query)
         dc.Unlock()
         if err != nil {
                 return nil, err
         }
         defer withLock(dc, func() { si.Close() })
-
         return resultFromStatement(driverStmt{dc, si}, args...)
 }
 
 // Query executes a query that returns rows, typically a SELECT.
 // The args are for any placeholder parameters in the query.
 func (db *DB) Query(query string, args ...interface{}) (*Rows, error) {
         var rows *Rows
         var err error
         for i := 0; i < 10; i++ {
                 rows, err = db.query(query, args)
@@ skipping 405 matching lines @@
                 ci, err = s.tx.grabConn() // blocks, waiting for the connection.
                 if err != nil {
                         return
                 }
                 releaseConn = func(error) {}
                 return ci, releaseConn, s.txsi.si, nil
         }
 
         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()
 
         // Make a new conn if all are busy.
         // TODO(bradfitz): or wait for one? make configurable later?
         if !match {
                 for i := 0; ; i++ {
                         dc, err := s.db.conn()
                         if err != nil {
                                 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
                         }
                         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)
                         s.mu.Unlock()
                         break
                 }
         }
 
         conn := cs.dc
         releaseConn = func(err error) { s.db.putConn(conn, err) }
@@ skipping 98 matching lines @@
                 s.txsi.Close()
                 return nil
         }
 
         return s.db.removeDep(s, s)
 }
 
 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
         return nil
 }
 
 // Rows is the result of a query. Its cursor starts before the first row
 // of the result set. Use Next to advance through the rows:
 //
 //     rows, err := db.Query("SELECT ...")
@@ skipping 174 matching lines @@
         return dr.resi.LastInsertId()
 }
 
 func (dr driverResult) RowsAffected() (int64, error) {
         dr.Lock()
         defer dr.Unlock()
         return dr.resi.RowsAffected()
 }
 
 func stack() string {
-        var buf [1024]byte
+        var buf [2 << 10]byte
         return string(buf[:runtime.Stack(buf[:], false)])
 }
 
 // withLock runs while holding lk.
 func withLock(lk sync.Locker, fn func()) {
         lk.Lock()
         fn()
         lk.Unlock()
 }