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)])
 }