database/sql: deflake query cancel tests

Rather then using a sleep in the fake DB, go to a channel
select and wait for the context to be done.

Fixes #18115

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

diff --git a/src/database/sql/fakedb_test.go b/src/database/sql/fakedb_test.go
index 9de9289644b41c2f01ee8c234c5919889158e9f6..416b97d501024c41982218e54c3c6962dc725268 100644 (file)
--- a/src/database/sql/fakedb_test.go
+++ b/src/database/sql/fakedb_test.go
@@ -511,6 +511,10 @@ func (c *fakeConn) prepareInsert(stmt *fakeStmt, parts []string) (*fakeStmt, err
 var hookPrepareBadConn func() bool
 
 func (c *fakeConn) Prepare(query string) (driver.Stmt, error) {
+       panic("use PrepareContext")
+}
+
+func (c *fakeConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
        c.numPrepare++
        if c.db == nil {
                panic("nil c.db; conn = " + fmt.Sprintf("%#v", c))
@@ -549,7 +553,13 @@ func (c *fakeConn) Prepare(query string) (driver.Stmt, error) {
                parts = parts[1:]
 
                if stmt.wait > 0 {
-                       time.Sleep(stmt.wait)
+                       wait := time.NewTimer(stmt.wait)
+                       select {
+                       case <-wait.C:
+                       case <-ctx.Done():
+                               wait.Stop()
+                               return nil, ctx.Err()
+                       }
                }
 
                c.incrStat(&c.stmtsMade)

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index c0f2cf29095ca1924435e95aa09a4db59fb4760f..27fb765cde87d757518f7d8e2ca970d62d760bd5 100644 (file)
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -24,6 +24,17 @@ func init() {
                c  *driverConn
        }
        freedFrom := make(map[dbConn]string)
+       var mu sync.Mutex
+       getFreedFrom := func(c dbConn) string {
+               mu.Lock()
+               defer mu.Unlock()
+               return freedFrom[c]
+       }
+       setFreedFrom := func(c dbConn, s string) {
+               mu.Lock()
+               defer mu.Unlock()
+               freedFrom[c] = s
+       }
        putConnHook = func(db *DB, c *driverConn) {
                idx := -1
                for i, v := range db.freeConn {
@@ -36,10 +47,10 @@ func init() {
                        // print before panic, as panic may get lost due to conflicting panic
                        // (all goroutines asleep) elsewhere, since we might not unlock
                        // the mutex in freeConn here.
-                       println("double free of conn. conflicts are:\nA) " + freedFrom[dbConn{db, c}] + "\n\nand\nB) " + stack())
+                       println("double free of conn. conflicts are:\nA) " + getFreedFrom(dbConn{db, c}) + "\n\nand\nB) " + stack())
                        panic("double free of conn.")
                }
-               freedFrom[dbConn{db, c}] = stack()
+               setFreedFrom(dbConn{db, c}, stack())
        }
 }
 
@@ -344,7 +355,7 @@ func TestQueryContextWait(t *testing.T) {
        // 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|")
+       _, err := db.QueryContext(ctx, "WAIT|1s|SELECT|people|age,name|")
        if err != context.DeadlineExceeded {
                t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
        }
@@ -372,7 +383,7 @@ func TestTxContextWait(t *testing.T) {
        // 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|")
+       _, err = tx.QueryContext(ctx, "WAIT|1s|SELECT|people|age,name|")
        if err != context.DeadlineExceeded {
                t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
        }