sql.go.diff

diff --git a/src/database/sql/sql.go b/src/database/sql/sql.go
index 8db9c78..caadc56 100644
--- a/src/database/sql/sql.go
+++ b/src/database/sql/sql.go
@@ -246,7 +246,7 @@ 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 connIfFree
+	dbmuClosed bool     // same as closed, but guarded by db.mu, for connIfFreePrepared
 }
 
 func (dc *driverConn) releaseConn(err error) {
@@ -683,40 +683,35 @@ var (
 	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.
+// connIfFreePrepared returns (dc, nil) if there are dc isn't in use in wanted.
 //
-// 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) {
+// The error is errConnBusy if no free connections in wanted.
+func (db *DB) connIfFreePrepared(wanted map[*driverConn]driver.Stmt) (*driverConn, driver.Stmt, error) {
 	db.mu.Lock()
 	defer db.mu.Unlock()
-	if wanted.dbmuClosed {
-		return nil, errConnClosed
-	}
-	if wanted.inUse {
-		return nil, errConnBusy
+	// remove dead connections
+	for dc := range wanted {
+		if dc.dbmuClosed {
+			delete(wanted, dc)
+		}
 	}
 	idx := -1
+	var dc *driverConn
+	var si driver.Stmt
 	for ii, v := range db.freeConn {
-		if v == wanted {
+		si = wanted[v]
+		if si != nil {
 			idx = ii
+			dc = v
 			break
 		}
 	}
 	if idx >= 0 {
 		db.freeConn = append(db.freeConn[:idx], db.freeConn[idx+1:]...)
-		wanted.inUse = true
-		return wanted, nil
+		dc.inUse = true
+		return dc, si, nil
 	}
-	// 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
+	return nil, nil, errConnBusy
 }
 
 // putConnHook is a hook for testing.
@@ -858,8 +853,9 @@ func (db *DB) prepare(query string) (*Stmt, error) {
 	stmt := &Stmt{
 		db:    db,
 		query: query,
-		css:   []connStmt{{dc, si}},
+		css:   make(map[*driverConn]driver.Stmt),
 	}
+	stmt.css[dc] = si
 	db.addDep(stmt, stmt)
 	db.putConn(dc, nil)
 	return stmt, nil
@@ -1266,12 +1262,6 @@ func (tx *Tx) QueryRow(query string, args ...interface{}) *Row {
 	return &Row{rows: rows, err: err}
 }
 
-// connStmt is a prepared statement on a particular connection.
-type connStmt struct {
-	dc *driverConn
-	si driver.Stmt
-}
-
 // Stmt is a prepared statement. Stmt is safe for concurrent use by multiple goroutines.
 type Stmt struct {
 	// Immutable:
@@ -1288,11 +1278,11 @@ type Stmt struct {
 	mu     sync.Mutex // protects the rest of the fields
 	closed bool
 
-	// css is a list of underlying driver statement interfaces
-	// that are valid on particular connections.  This is only
-	// used if tx == nil and one is found that has idle
+	// css is a map of underlying driver statement interfaces
+	// from particular connection the statement is valid on.
+	// This is only used if tx == nil and one is found that has idle
 	// connections.  If tx != nil, txsi is always used.
-	css []connStmt
+	css map[*driverConn]driver.Stmt
 }
 
 // Exec executes a prepared statement with the given arguments and
@@ -1372,29 +1362,18 @@ func (s *Stmt) connStmt() (ci *driverConn, releaseConn func(error), si driver.St
 		return ci, releaseConn, s.txsi.si, nil
 	}
 
-	for i := 0; i < len(s.css); i++ {
-		v := s.css[i]
-		_, err := s.db.connIfFree(v.dc)
-		if err == nil {
-			s.mu.Unlock()
-			return v.dc, v.dc.releaseConn, v.si, nil
-		}
-		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--
-		}
-
-	}
+	dc, si, err := s.db.connIfFreePrepared(s.css)
 	s.mu.Unlock()
+	if err == nil {
+		return dc, dc.releaseConn, si, nil
+	}
 
 	// If all connections are busy, either wait for one to become available (if
 	// we've already hit the maximum number of open connections) or create a
 	// new one.
 	//
 	// TODO(bradfitz): or always wait for one? make configurable later?
-	dc, err := s.db.conn()
+	dc, err = s.db.conn()
 	if err != nil {
 		return nil, nil, nil, err
 	}
@@ -1402,13 +1381,11 @@ func (s *Stmt) connStmt() (ci *driverConn, releaseConn func(error), si driver.St
 	// Do another pass over the list to see whether this statement has
 	// already been prepared on the connection assigned to us.
 	s.mu.Lock()
-	for _, v := range s.css {
-		if v.dc == dc {
-			s.mu.Unlock()
-			return dc, dc.releaseConn, v.si, nil
-		}
-	}
+	si = s.css[dc]
 	s.mu.Unlock()
+	if si != nil {
+		return dc, dc.releaseConn, si, nil
+	}
 
 	// No luck; we need to prepare the statement on this connection
 	dc.Lock()
@@ -1419,8 +1396,10 @@ func (s *Stmt) connStmt() (ci *driverConn, releaseConn func(error), si driver.St
 		return nil, nil, nil, err
 	}
 	s.mu.Lock()
-	cs := connStmt{dc, si}
-	s.css = append(s.css, cs)
+	if s.css == nil {
+		s.css = make(map[*driverConn]driver.Stmt)
+	}
+	s.css[dc] = si
 	s.mu.Unlock()
 
 	return dc, dc.releaseConn, si, nil
@@ -1541,9 +1520,9 @@ func (s *Stmt) finalClose() error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.css != nil {
-		for _, v := range s.css {
-			s.db.noteUnusedDriverStatement(v.dc, v.si)
-			v.dc.removeOpenStmt(v.si)
+		for dc, si := range s.css {
+			s.db.noteUnusedDriverStatement(dc, si)
+			dc.removeOpenStmt(si)
 		}
 		s.css = nil
 	}