Sync with Community

vendor/github.com/lib/pq/conn_go18.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"sync/atomic"
 	"time"
 )
 
@@ -89,10 +90,21 @@ func (cn *conn) Ping(ctx context.Context) error {
 
 func (cn *conn) watchCancel(ctx context.Context) func() {
 	if done := ctx.Done(); done != nil {
-		finished := make(chan struct{})
+		finished := make(chan struct{}, 1)
 		go func() {
 			select {
 			case <-done:
+				select {
+				case finished <- struct{}{}:
+				default:
+					// We raced with the finish func, let the next query handle this with the
+					// context.
+					return
+				}
+
+				// Set the connection state to bad so it does not get reused.
+				cn.setBad()
+
 				// At this point the function level context is canceled,
 				// so it must not be used for the additional network
 				// request to cancel the query.
@@ -101,13 +113,14 @@ func (cn *conn) watchCancel(ctx context.Context) func() {
 				defer cancel()
 
 				_ = cn.cancel(ctxCancel)
-				finished <- struct{}{}
 			case <-finished:
 			}
 		}()
 		return func() {
 			select {
 			case <-finished:
+				cn.setBad()
+				cn.Close()
 			case finished <- struct{}{}:
 			}
 		}
@@ -116,17 +129,29 @@ func (cn *conn) watchCancel(ctx context.Context) func() {
 }
 
 func (cn *conn) cancel(ctx context.Context) error {
-	c, err := dial(ctx, cn.dialer, cn.opts)
+	// Create a new values map (copy). This makes sure the connection created
+	// in this method cannot write to the same underlying data, which could
+	// cause a concurrent map write panic. This is necessary because cancel
+	// is called from a goroutine in watchCancel.
+	o := make(values)
+	for k, v := range cn.opts {
+		o[k] = v
+	}
+
+	c, err := dial(ctx, cn.dialer, o)
 	if err != nil {
 		return err
 	}
 	defer c.Close()
 
 	{
+		bad := &atomic.Value{}
+		bad.Store(false)
 		can := conn{
-			c: c,
+			c:   c,
+			bad: bad,
 		}
-		err = can.ssl(cn.opts)
+		err = can.ssl(o)
 		if err != nil {
 			return err
 		}