database/sql: do not bypass the driver locks with Context methods

When context methods were initially added it was attempted to unify
behavior between drivers without Context methods and those with
Context methods to always return right away when the Context expired.
However in doing so the driver call could be executed outside of the
scope of the driver connection lock and thus bypassing thread safety.

The new behavior waits until the driver operation is complete. It then
checks to see if the context has expired and if so returns that error.

Change-Id: I4a5c7c3263420c57778f36a5ed6fa0ef8cb32b20
Reviewed-on: https://go-review.googlesource.com/32422
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

1 files changed, 48 insertions, 184 deletions

diff --git a/src/database/sql/ctxutil.go b/src/database/sql/ctxutil.go
index ddc4b7228f..7c05ce2448 100644
--- a/src/database/sql/ctxutil.go
+++ b/src/database/sql/ctxutil.go
@@ -14,40 +14,16 @@ func ctxDriverPrepare(ctx context.Context, ci driver.Conn, query string) (driver
 	if ciCtx, is := ci.(driver.ConnPrepareContext); is {
 		return ciCtx.PrepareContext(ctx, query)
 	}
-	if ctx.Done() == context.Background().Done() {
-		return ci.Prepare(query)
-	}
-
-	type R struct {
-		err   error
-		panic interface{}
-		si    driver.Stmt
-	}
-
-	rc := make(chan R, 1)
-	go func() {
-		r := R{}
-		defer func() {
-			if v := recover(); v != nil {
-				r.panic = v
-			}
-			rc <- r
-		}()
-		r.si, r.err = ci.Prepare(query)
-	}()
-	select {
-	case <-ctx.Done():
-		go func() {
-			<-rc
-			close(rc)
-		}()
-		return nil, ctx.Err()
-	case r := <-rc:
-		if r.panic != nil {
-			panic(r.panic)
+	si, err := ci.Prepare(query)
+	if err == nil {
+		select {
+		default:
+		case <-ctx.Done():
+			si.Close()
+			return nil, ctx.Err()
 		}
-		return r.si, r.err
 	}
+	return si, err
 }
 
 func ctxDriverExec(ctx context.Context, execer driver.Execer, query string, nvdargs []driver.NamedValue) (driver.Result, error) {
@@ -58,84 +34,38 @@ func ctxDriverExec(ctx context.Context, execer driver.Execer, query string, nvda
 	if err != nil {
 		return nil, err
 	}
-	if ctx.Done() == context.Background().Done() {
-		return execer.Exec(query, dargs)
-	}
-
-	type R struct {
-		err   error
-		panic interface{}
-		resi  driver.Result
-	}
 
-	rc := make(chan R, 1)
-	go func() {
-		r := R{}
-		defer func() {
-			if v := recover(); v != nil {
-				r.panic = v
-			}
-			rc <- r
-		}()
-		r.resi, r.err = execer.Exec(query, dargs)
-	}()
-	select {
-	case <-ctx.Done():
-		go func() {
-			<-rc
-			close(rc)
-		}()
-		return nil, ctx.Err()
-	case r := <-rc:
-		if r.panic != nil {
-			panic(r.panic)
+	resi, err := execer.Exec(query, dargs)
+	if err == nil {
+		select {
+		default:
+		case <-ctx.Done():
+			return resi, ctx.Err()
 		}
-		return r.resi, r.err
 	}
+	return resi, err
 }
 
 func ctxDriverQuery(ctx context.Context, queryer driver.Queryer, query string, nvdargs []driver.NamedValue) (driver.Rows, error) {
 	if queryerCtx, is := queryer.(driver.QueryerContext); is {
-		return queryerCtx.QueryContext(ctx, query, nvdargs)
+		ret, err := queryerCtx.QueryContext(ctx, query, nvdargs)
+		return ret, err
 	}
 	dargs, err := namedValueToValue(nvdargs)
 	if err != nil {
 		return nil, err
 	}
-	if ctx.Done() == context.Background().Done() {
-		return queryer.Query(query, dargs)
-	}
 
-	type R struct {
-		err   error
-		panic interface{}
-		rowsi driver.Rows
-	}
-
-	rc := make(chan R, 1)
-	go func() {
-		r := R{}
-		defer func() {
-			if v := recover(); v != nil {
-				r.panic = v
-			}
-			rc <- r
-		}()
-		r.rowsi, r.err = queryer.Query(query, dargs)
-	}()
-	select {
-	case <-ctx.Done():
-		go func() {
-			<-rc
-			close(rc)
-		}()
-		return nil, ctx.Err()
-	case r := <-rc:
-		if r.panic != nil {
-			panic(r.panic)
+	rowsi, err := queryer.Query(query, dargs)
+	if err == nil {
+		select {
+		default:
+		case <-ctx.Done():
+			rowsi.Close()
+			return nil, ctx.Err()
 		}
-		return r.rowsi, r.err
 	}
+	return rowsi, err
 }
 
 func ctxDriverStmtExec(ctx context.Context, si driver.Stmt, nvdargs []driver.NamedValue) (driver.Result, error) {
@@ -146,40 +76,16 @@ func ctxDriverStmtExec(ctx context.Context, si driver.Stmt, nvdargs []driver.Nam
 	if err != nil {
 		return nil, err
 	}
-	if ctx.Done() == context.Background().Done() {
-		return si.Exec(dargs)
-	}
-
-	type R struct {
-		err   error
-		panic interface{}
-		resi  driver.Result
-	}
 
-	rc := make(chan R, 1)
-	go func() {
-		r := R{}
-		defer func() {
-			if v := recover(); v != nil {
-				r.panic = v
-			}
-			rc <- r
-		}()
-		r.resi, r.err = si.Exec(dargs)
-	}()
-	select {
-	case <-ctx.Done():
-		go func() {
-			<-rc
-			close(rc)
-		}()
-		return nil, ctx.Err()
-	case r := <-rc:
-		if r.panic != nil {
-			panic(r.panic)
+	resi, err := si.Exec(dargs)
+	if err == nil {
+		select {
+		default:
+		case <-ctx.Done():
+			return resi, ctx.Err()
 		}
-		return r.resi, r.err
 	}
+	return resi, err
 }
 
 func ctxDriverStmtQuery(ctx context.Context, si driver.Stmt, nvdargs []driver.NamedValue) (driver.Rows, error) {
@@ -190,40 +96,17 @@ func ctxDriverStmtQuery(ctx context.Context, si driver.Stmt, nvdargs []driver.Na
 	if err != nil {
 		return nil, err
 	}
-	if ctx.Done() == context.Background().Done() {
-		return si.Query(dargs)
-	}
 
-	type R struct {
-		err   error
-		panic interface{}
-		rowsi driver.Rows
-	}
-
-	rc := make(chan R, 1)
-	go func() {
-		r := R{}
-		defer func() {
-			if v := recover(); v != nil {
-				r.panic = v
-			}
-			rc <- r
-		}()
-		r.rowsi, r.err = si.Query(dargs)
-	}()
-	select {
-	case <-ctx.Done():
-		go func() {
-			<-rc
-			close(rc)
-		}()
-		return nil, ctx.Err()
-	case r := <-rc:
-		if r.panic != nil {
-			panic(r.panic)
+	rowsi, err := si.Query(dargs)
+	if err == nil {
+		select {
+		default:
+		case <-ctx.Done():
+			rowsi.Close()
+			return nil, ctx.Err()
 		}
-		return r.rowsi, r.err
 	}
+	return rowsi, err
 }
 
 var errLevelNotSupported = errors.New("sql: selected isolation level is not supported")
@@ -249,35 +132,16 @@ func ctxDriverBegin(ctx context.Context, ci driver.Conn) (driver.Tx, error) {
 		return nil, errors.New("sql: driver does not support read-only transactions")
 	}
 
-	type R struct {
-		err   error
-		panic interface{}
-		txi   driver.Tx
-	}
-	rc := make(chan R, 1)
-	go func() {
-		r := R{}
-		defer func() {
-			if v := recover(); v != nil {
-				r.panic = v
-			}
-			rc <- r
-		}()
-		r.txi, r.err = ci.Begin()
-	}()
-	select {
-	case <-ctx.Done():
-		go func() {
-			<-rc
-			close(rc)
-		}()
-		return nil, ctx.Err()
-	case r := <-rc:
-		if r.panic != nil {
-			panic(r.panic)
+	txi, err := ci.Begin()
+	if err == nil {
+		select {
+		default:
+		case <-ctx.Done():
+			txi.Rollback()
+			return nil, ctx.Err()
 		}
-		return r.txi, r.err
 	}
+	return txi, err
 }
 
 func namedValueToValue(named []driver.NamedValue) ([]driver.Value, error) {