Fix a bug for sending cancel request to terminate a long running query

base.go

@@ -142,15 +142,17 @@ func (db *baseDB) withConn(
 	if ctx != nil && ctx.Done() != nil {
 		fnDone = make(chan struct{})
 		go func() {
-			select {
-			case <-fnDone: // fn has finished, skip cancel
-			case <-ctx.Done():
-				err := db.cancelRequest(cn.ProcessID, cn.SecretKey)
-				if err != nil {
-					internal.Logger.Printf(ctx, "cancelRequest failed: %s", err)
+			for {
+				select {
+				case <-fnDone: // fn has finished, skip cancel
+				case <-ctx.Done():
+					err := db.cancelRequest(cn.ProcessID, cn.SecretKey)
+					if err != nil {
+						internal.Logger.Printf(ctx, "cancelRequest failed: %s", err)
+					}
+					// Signal end of conn use.
+					fnDone <- struct{}{}
 				}
-				// Signal end of conn use.
-				fnDone <- struct{}{}
 			}
 		}()
 	}
@@ -188,8 +190,8 @@ func (db *baseDB) shouldRetry(err error) bool {
 			"53300", // too_many_connections
 			"55000": // attempted to delete invisible tuple
 			return true
-		case "57014": // statement_timeout
-			return db.opt.RetryStatementTimeout
+		case "57014": // query_canceled
+			return db.opt.RetryCancelledQuery
 		default:
 			return false
 		}

options.go

@@ -57,8 +57,8 @@ type Options struct {
 	// Maximum number of retries before giving up.
 	// Default is to not retry failed queries.
 	MaxRetries int
-	// Whether to retry queries cancelled because of statement_timeout.
-	RetryStatementTimeout bool
+	// Whether to retry queries cancelled because of statement_timeout or cancel request.
+	RetryCancelledQuery bool
 	// Minimum backoff between each retry.
 	// Default is 250 milliseconds; -1 disables backoff.
 	MinRetryBackoff time.Duration