Merge pull request #618 from ainar-g/master

Avoid goroutine leak on error in query with context

Author: tamird

conn_go18.go

@@ -19,6 +19,9 @@ func (cn *conn) QueryContext(ctx context.Context, query string, args []driver.Na
 	finish := cn.watchCancel(ctx)
 	r, err := cn.query(query, list)
 	if err != nil {
+		if finish != nil {
+			finish()
+		}
 		return nil, err
 	}
 	r.finish = finish

go18_test.go

@@ -5,6 +5,7 @@ package pq
 import (
 	"context"
 	"database/sql"
+	"runtime"
 	"testing"
 	"time"
 )
@@ -155,6 +156,36 @@ func TestContextCancelQuery(t *testing.T) {
 	}
 }
 
+// TestIssue617 tests that a failed query in QueryContext doesn't lead to a
+// goroutine leak.
+func TestIssue617(t *testing.T) {
+	db := openTestConn(t)
+	defer db.Close()
+
+	const N = 10
+
+	numGoroutineStart := runtime.NumGoroutine()
+	for i := 0; i < N; i++ {
+		func() {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			_, err := db.QueryContext(ctx, `SELECT * FROM DOESNOTEXIST`)
+			pqErr, _ := err.(*Error)
+			// Expecting "pq: relation \"doesnotexist\" does not exist" error.
+			if err == nil || pqErr == nil || pqErr.Code != "42P01" {
+				t.Fatalf("expected undefined table error, got %v", err)
+			}
+		}()
+	}
+	numGoroutineFinish := runtime.NumGoroutine()
+
+	// We use N/2 and not N because the GC and other actors may increase or
+	// decrease the number of goroutines.
+	if numGoroutineFinish-numGoroutineStart >= N/2 {
+		t.Errorf("goroutine leak detected, was %d, now %d", numGoroutineStart, numGoroutineFinish)
+	}
+}
+
 func TestContextCancelBegin(t *testing.T) {
 	db := openTestConn(t)
 	defer db.Close()