1 files changed, 102 insertions, 43 deletions

diff --git a/libgo/go/database/sql/sql.go b/libgo/go/database/sql/sql.go
index 62b551d89b5..51a357b37de 100644
--- a/libgo/go/database/sql/sql.go
+++ b/libgo/go/database/sql/sql.go
@@ -175,6 +175,16 @@ var ErrNoRows = errors.New("sql: no rows in result set")
 
 // DB is a database handle. It's safe for concurrent use by multiple
 // goroutines.
+//
+// If the underlying database driver has the concept of a connection
+// and per-connection session state, the sql package manages creating
+// and freeing connections automatically, including maintaining a free
+// pool of idle connections. If observing session state is required,
+// either do not share a *DB between multiple concurrent goroutines or
+// create and observe all state only within a transaction. Once
+// DB.Open is called, the returned Tx is bound to a single isolated
+// connection. Once Tx.Commit or Tx.Rollback is called, that
+// connection is returned to DB's idle connection pool.
 type DB struct {
 	driver driver.Driver
 	dsn    string
@@ -241,34 +251,56 @@ func (db *DB) conn() (driver.Conn, error) {
 func (db *DB) connIfFree(wanted driver.Conn) (conn driver.Conn, ok bool) {
 	db.mu.Lock()
 	defer db.mu.Unlock()
-	for n, conn := range db.freeConn {
-		if conn == wanted {
-			db.freeConn[n] = db.freeConn[len(db.freeConn)-1]
-			db.freeConn = db.freeConn[:len(db.freeConn)-1]
-			return wanted, true
+	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]
+		return wanted, true
 	}
 	return nil, false
 }
 
-func (db *DB) putConn(c driver.Conn) {
+// putConnHook is a hook for testing.
+var putConnHook func(*DB, driver.Conn)
+
+// putConn adds a connection to the db's free pool.
+// err is optionally the last error that occured on this connection.
+func (db *DB) putConn(c driver.Conn, err error) {
+	if err == driver.ErrBadConn {
+		// Don't reuse bad connections.
+		return
+	}
 	db.mu.Lock()
-	defer db.mu.Unlock()
+	if putConnHook != nil {
+		putConnHook(db, c)
+	}
 	if n := len(db.freeConn); !db.closed && n < db.maxIdleConns() {
 		db.freeConn = append(db.freeConn, c)
+		db.mu.Unlock()
 		return
 	}
-	db.closeConn(c) // TODO(bradfitz): release lock before calling this?
-}
-
-func (db *DB) closeConn(c driver.Conn) {
-	// TODO: check to see if we need this Conn for any prepared statements
-	// that are active.
+	// TODO: check to see if we need this Conn for any prepared
+	// statements which are still active?
+	db.mu.Unlock()
 	c.Close()
 }
 
 // Prepare creates a prepared statement for later execution.
 func (db *DB) Prepare(query string) (*Stmt, error) {
+	var stmt *Stmt
+	var err error
+	for i := 0; i < 10; i++ {
+		stmt, err = db.prepare(query)
+		if err != driver.ErrBadConn {
+			break
+		}
+	}
+	return stmt, err
+}
+
+func (db *DB) prepare(query string) (stmt *Stmt, err error) {
 	// TODO: check if db.driver supports an optional
 	// driver.Preparer interface and call that instead, if so,
 	// otherwise we make a prepared statement that's bound
@@ -279,12 +311,12 @@ func (db *DB) Prepare(query string) (*Stmt, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer db.putConn(ci)
+	defer db.putConn(ci, err)
 	si, err := ci.Prepare(query)
 	if err != nil {
 		return nil, err
 	}
-	stmt := &Stmt{
+	stmt = &Stmt{
 		db:    db,
 		query: query,
 		css:   []connStmt{{ci, si}},
@@ -295,15 +327,22 @@ func (db *DB) Prepare(query string) (*Stmt, error) {
 // Exec executes a query without returning any rows.
 func (db *DB) Exec(query string, args ...interface{}) (Result, error) {
 	sargs, err := subsetTypeArgs(args)
-	if err != nil {
-		return nil, err
+	var res Result
+	for i := 0; i < 10; i++ {
+		res, err = db.exec(query, sargs)
+		if err != driver.ErrBadConn {
+			break
+		}
 	}
+	return res, err
+}
 
+func (db *DB) exec(query string, sargs []driver.Value) (res Result, err error) {
 	ci, err := db.conn()
 	if err != nil {
 		return nil, err
 	}
-	defer db.putConn(ci)
+	defer db.putConn(ci, err)
 
 	if execer, ok := ci.(driver.Execer); ok {
 		resi, err := execer.Exec(query, sargs)
@@ -354,13 +393,25 @@ func (db *DB) QueryRow(query string, args ...interface{}) *Row {
 // Begin starts a transaction. The isolation level is dependent on
 // the driver.
 func (db *DB) Begin() (*Tx, error) {
+	var tx *Tx
+	var err error
+	for i := 0; i < 10; i++ {
+		tx, err = db.begin()
+		if err != driver.ErrBadConn {
+			break
+		}
+	}
+	return tx, err
+}
+
+func (db *DB) begin() (tx *Tx, err error) {
 	ci, err := db.conn()
 	if err != nil {
 		return nil, err
 	}
 	txi, err := ci.Begin()
 	if err != nil {
-		db.putConn(ci)
+		db.putConn(ci, err)
 		return nil, fmt.Errorf("sql: failed to Begin transaction: %v", err)
 	}
 	return &Tx{
@@ -406,7 +457,7 @@ func (tx *Tx) close() {
 		panic("double close") // internal error
 	}
 	tx.done = true
-	tx.db.putConn(tx.ci)
+	tx.db.putConn(tx.ci, nil)
 	tx.ci = nil
 	tx.txi = nil
 }
@@ -561,9 +612,11 @@ func (tx *Tx) Query(query string, args ...interface{}) (*Rows, error) {
 		return nil, err
 	}
 	rows, err := stmt.Query(args...)
-	if err == nil {
-		rows.closeStmt = stmt
+	if err != nil {
+		stmt.Close()
+		return nil, err
 	}
+	rows.closeStmt = stmt
 	return rows, err
 }
 
@@ -609,7 +662,7 @@ func (s *Stmt) Exec(args ...interface{}) (Result, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer releaseConn()
+	defer releaseConn(nil)
 
 	// -1 means the driver doesn't know how to count the number of
 	// placeholders, so we won't sanity check input here and instead let the
@@ -672,7 +725,7 @@ func (s *Stmt) Exec(args ...interface{}) (Result, error) {
 // connStmt returns a free driver connection on which to execute the
 // statement, a function to call to release the connection, and a
 // statement bound to that connection.
-func (s *Stmt) connStmt() (ci driver.Conn, releaseConn func(), si driver.Stmt, err error) {
+func (s *Stmt) connStmt() (ci driver.Conn, releaseConn func(error), si driver.Stmt, err error) {
 	if err = s.stickyErr; err != nil {
 		return
 	}
@@ -691,7 +744,7 @@ func (s *Stmt) connStmt() (ci driver.Conn, releaseConn func(), si driver.Stmt, e
 		if err != nil {
 			return
 		}
-		releaseConn = func() { s.tx.releaseConn() }
+		releaseConn = func(error) { s.tx.releaseConn() }
 		return ci, releaseConn, s.txsi, nil
 	}
 
@@ -700,7 +753,7 @@ func (s *Stmt) connStmt() (ci driver.Conn, releaseConn func(), si driver.Stmt, e
 	for _, v := range s.css {
 		// TODO(bradfitz): lazily clean up entries in this
 		// list with dead conns while enumerating
-		if _, match = s.db.connIfFree(cs.ci); match {
+		if _, match = s.db.connIfFree(v.ci); match {
 			cs = v
 			break
 		}
@@ -710,22 +763,28 @@ func (s *Stmt) connStmt() (ci driver.Conn, releaseConn func(), si driver.Stmt, e
 	// Make a new conn if all are busy.
 	// TODO(bradfitz): or wait for one? make configurable later?
 	if !match {
-		ci, err := s.db.conn()
-		if err != nil {
-			return nil, nil, nil, err
-		}
-		si, err := ci.Prepare(s.query)
-		if err != nil {
-			return nil, nil, nil, err
+		for i := 0; ; i++ {
+			ci, err := s.db.conn()
+			if err != nil {
+				return nil, nil, nil, err
+			}
+			si, err := ci.Prepare(s.query)
+			if err == driver.ErrBadConn && i < 10 {
+				continue
+			}
+			if err != nil {
+				return nil, nil, nil, err
+			}
+			s.mu.Lock()
+			cs = connStmt{ci, si}
+			s.css = append(s.css, cs)
+			s.mu.Unlock()
+			break
 		}
-		s.mu.Lock()
-		cs = connStmt{ci, si}
-		s.css = append(s.css, cs)
-		s.mu.Unlock()
 	}
 
 	conn := cs.ci
-	releaseConn = func() { s.db.putConn(conn) }
+	releaseConn = func(err error) { s.db.putConn(conn, err) }
 	return conn, releaseConn, cs.si, nil
 }
 
@@ -749,7 +808,7 @@ func (s *Stmt) Query(args ...interface{}) (*Rows, error) {
 	}
 	rowsi, err := si.Query(sargs)
 	if err != nil {
-		s.db.putConn(ci)
+		releaseConn(err)
 		return nil, err
 	}
 	// Note: ownership of ci passes to the *Rows, to be freed
@@ -800,7 +859,7 @@ func (s *Stmt) Close() error {
 		for _, v := range s.css {
 			if ci, match := s.db.connIfFree(v.ci); match {
 				v.si.Close()
-				s.db.putConn(ci)
+				s.db.putConn(ci, nil)
 			} else {
 				// TODO(bradfitz): care that we can't close
 				// this statement because the statement's
@@ -827,7 +886,7 @@ func (s *Stmt) Close() error {
 type Rows struct {
 	db          *DB
 	ci          driver.Conn // owned; must call putconn when closed to release
-	releaseConn func()
+	releaseConn func(error)
 	rowsi       driver.Rows
 
 	closed    bool
@@ -939,7 +998,7 @@ func (rs *Rows) Close() error {
 	}
 	rs.closed = true
 	err := rs.rowsi.Close()
-	rs.releaseConn()
+	rs.releaseConn(err)
 	if rs.closeStmt != nil {
 		rs.closeStmt.Close()
 	}
@@ -963,7 +1022,7 @@ func (r *Row) Scan(dest ...interface{}) error {
 	}
 
 	// TODO(bradfitz): for now we need to defensively clone all
-	// []byte that the driver returned (not permitting 
+	// []byte that the driver returned (not permitting
 	// *RawBytes in Rows.Scan), since we're about to close
 	// the Rows in our defer, when we return from this function.
 	// the contract with the driver.Next(...) interface is that it