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>

src/database/sql/sql_test.go

@@ -141,10 +141,7 @@ func closeDB(t testing.TB, db *DB) {
 	if err != nil {
 		t.Fatalf("error closing DB: %v", err)
 	}
-	db.mu.Lock()
-	count := db.numOpen
-	db.mu.Unlock()
-	if count != 0 {
+	if count := db.numOpenConns(); count != 0 {
 		t.Fatalf("%d connections still open after closing DB", count)
 	}
 }
@@ -183,6 +180,12 @@ func (db *DB) numFreeConns() int {
 	return len(db.freeConn)
 }
 
+func (db *DB) numOpenConns() int {
+	db.mu.Lock()
+	defer db.mu.Unlock()
+	return db.numOpen
+}
+
 // clearAllConns closes all connections in db.
 func (db *DB) clearAllConns(t *testing.T) {
 	db.SetMaxIdleConns(0)
@@ -320,6 +323,75 @@ func TestQueryContext(t *testing.T) {
 	}
 }
 
+func waitCondition(waitFor, checkEvery time.Duration, fn func() bool) bool {
+	deadline := time.Now().Add(waitFor)
+	for time.Now().Before(deadline) {
+		if fn() {
+			return true
+		}
+		time.Sleep(checkEvery)
+	}
+	return false
+}
+
+func TestQueryContextWait(t *testing.T) {
+	db := newTestDB(t, "people")
+	defer closeDB(t, db)
+	prepares0 := numPrepares(t, db)
+
+	ctx, _ := context.WithTimeout(context.Background(), time.Millisecond*15)
+
+	// This will trigger the *fakeConn.Prepare method which will take time
+	// performing the query. The ctxDriverPrepare func will check the context
+	// after this and close the rows and return an error.
+	_, err := db.QueryContext(ctx, "WAIT|30ms|SELECT|people|age,name|")
+	if err != context.DeadlineExceeded {
+		t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
+	}
+
+	// Verify closed rows connection after error condition.
+	if n := db.numFreeConns(); n != 1 {
+		t.Fatalf("free conns after query hitting EOF = %d; want 1", n)
+	}
+	if prepares := numPrepares(t, db) - prepares0; prepares != 1 {
+		t.Errorf("executed %d Prepare statements; want 1", prepares)
+	}
+}
+
+func TestTxContextWait(t *testing.T) {
+	db := newTestDB(t, "people")
+	defer closeDB(t, db)
+
+	ctx, _ := context.WithTimeout(context.Background(), time.Millisecond*15)
+
+	tx, err := db.BeginContext(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// This will trigger the *fakeConn.Prepare method which will take time
+	// performing the query. The ctxDriverPrepare func will check the context
+	// after this and close the rows and return an error.
+	_, err = tx.QueryContext(ctx, "WAIT|30ms|SELECT|people|age,name|")
+	if err != context.DeadlineExceeded {
+		t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
+	}
+
+	var numFree int
+	if !waitCondition(5*time.Second, 5*time.Millisecond, func() bool {
+		numFree = db.numFreeConns()
+		return numFree == 0
+	}) {
+		t.Fatalf("free conns after hitting EOF = %d; want 0", numFree)
+	}
+
+	// Ensure the dropped connection allows more connections to be made.
+	// Checked on DB Close.
+	waitCondition(5*time.Second, 5*time.Millisecond, func() bool {
+		return db.numOpenConns() == 0
+	})
+}
+
 func TestMultiResultSetQuery(t *testing.T) {
 	db := newTestDB(t, "people")
 	defer closeDB(t, db)