Make connection killing resilient to MySQL hangs

go/mysql/fakesqldb/server.go

@@ -376,11 +376,11 @@
 	}
 	key := strings.ToLower(query)
 	db.mu.Lock()
-	defer db.mu.Unlock()
 	db.queryCalled[key]++
 	db.querylog = append(db.querylog, key)
 	// Check if we should close the connection and provoke errno 2013.
 	if db.shouldClose.Load() {
+		defer db.mu.Unlock()
 		c.Close()
 
 		// log error
@@ -394,6 +394,8 @@
 	// The driver may send this at connection time, and we don't want it to
 	// interfere.
 	if key == "set names utf8" || strings.HasPrefix(key, "set collation_connection = ") {
+		defer db.mu.Unlock()
+
 		// log error
 		if err := callback(&sqltypes.Result{}); err != nil {
 			log.Errorf("callback failed : %v", err)
@@ -403,12 +405,14 @@
 
 	// check if we should reject it.
 	if err, ok := db.rejectedData[key]; ok {
+		db.mu.Unlock()
 		return err
 	}
 
 	// Check explicit queries from AddQuery().
 	result, ok := db.data[key]
 	if ok {
+		db.mu.Unlock()
 		if f := result.BeforeFunc; f != nil {
 			f()
 		}
@@ -419,12 +423,9 @@
 	for _, pat := range db.patternData {
 		if pat.expr.MatchString(query) {
 			userCallback, ok := db.queryPatternUserCallback[pat.expr]
+			db.mu.Unlock()
 			if ok {
-				// Since the user call back can be indefinitely stuck, we shouldn't hold the lock indefinitely.
-				// This is only test code, so no actual cause for concern.
-				db.mu.Unlock()
 				userCallback(query)
-				db.mu.Lock()
 			}
 			if pat.err != "" {
 				return fmt.Errorf(pat.err)
@@ -433,6 +434,8 @@
 		}
 	}
 
+	defer db.mu.Unlock()
+
 	if db.neverFail.Load() {
 		return callback(&sqltypes.Result{})
 	}

go/vt/vttablet/tabletserver/connpool/dbconn.go

@@ -40,6 +40,8 @@
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 )
 
+const defaultKillTimeout = 5 * time.Second
+
 // Conn is a db connection for tabletserver.
 // It performs automatic reconnects as needed.
 // Its Execute function has a timeout that can kill
@@ -52,11 +54,13 @@
 	env     tabletenv.Env
 	dbaPool *dbconnpool.ConnectionPool
 	stats   *tabletenv.Stats
-	current atomic.Value
+	current atomic.Pointer[string]
 
 	// err will be set if a query is killed through a Kill.
 	errmu sync.Mutex
 	err   error
+
+	killTimeout time.Duration
 }
 
 // NewConnection creates a new DBConn. It triggers a CheckMySQL if creation fails.
@@ -71,12 +75,12 @@
 		return nil, err
 	}
 	db := &Conn{
-		conn:    c,
-		env:     pool.env,
-		stats:   pool.env.Stats(),
-		dbaPool: pool.dbaPool,
+		conn:        c,
+		env:         pool.env,
+		stats:       pool.env.Stats(),
+		dbaPool:     pool.dbaPool,
+		killTimeout: defaultKillTimeout,
 	}
-	db.current.Store("")
 	return db, nil
 }
 
@@ -87,12 +91,12 @@
 		return nil, err
 	}
 	dbconn := &Conn{
-		conn:    c,
-		dbaPool: dbaPool,
-		stats:   tabletenv.NewStats(servenv.NewExporter("Temp", "Tablet")),
-		env:     env,
+		conn:        c,
+		dbaPool:     dbaPool,
+		stats:       tabletenv.NewStats(servenv.NewExporter("Temp", "Tablet")),
+		env:         env,
+		killTimeout: defaultKillTimeout,
 	}
-	dbconn.current.Store("")
 	if setting == nil {
 		return dbconn, nil
 	}
@@ -153,28 +157,42 @@
 }
 
 func (dbc *Conn) execOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
-	dbc.current.Store(query)
-	defer dbc.current.Store("")
+	dbc.current.Store(&query)
+	defer dbc.current.Store(nil)
 
 	// Check if the context is already past its deadline before
 	// trying to execute the query.
 	if err := ctx.Err(); err != nil {
 		return nil, fmt.Errorf("%v before execution started", err)
 	}
 
-	defer dbc.stats.MySQLTimings.Record("Exec", time.Now())
-
-	done, wg := dbc.setDeadline(ctx)
-	qr, err := dbc.conn.ExecuteFetch(query, maxrows, wantfields)
+	now := time.Now()
+	defer dbc.stats.MySQLTimings.Record("Exec", now)
 
-	if done != nil {
-		close(done)
-		wg.Wait()
+	type execResult struct {
+		result *sqltypes.Result
+		err    error
 	}
-	if dbcerr := dbc.Err(); dbcerr != nil {
-		return nil, dbcerr
+
+	ch := make(chan execResult)
+	go func() {
+		result, err := dbc.conn.ExecuteFetch(query, maxrows, wantfields)
+		ch <- execResult{result, err}
+	}()
+
+	select {
+	case <-ctx.Done():
+		killCtx, cancel := context.WithTimeout(context.Background(), dbc.killTimeout)
+		defer cancel()
+
+		_ = dbc.KillWithContext(killCtx, ctx.Err().Error(), time.Since(now))
+		return nil, dbc.Err()
+	case r := <-ch:
+		if dbcErr := dbc.Err(); dbcErr != nil {
+			return nil, dbcErr
+		}
+		return r.result, r.err
 	}
-	return qr, err
 }
 
 // ExecOnce executes the specified query, but does not retry on connection errors.
@@ -250,22 +268,30 @@
 }
 
 func (dbc *Conn) streamOnce(ctx context.Context, query string, callback func(*sqltypes.Result) error, alloc func() *sqltypes.Result, streamBufferSize int) error {
-	defer dbc.stats.MySQLTimings.Record("ExecStream", time.Now())
+	dbc.current.Store(&query)
+	defer dbc.current.Store(nil)
 
-	dbc.current.Store(query)
-	defer dbc.current.Store("")
+	now := time.Now()
+	defer dbc.stats.MySQLTimings.Record("ExecStream", now)
 
-	done, wg := dbc.setDeadline(ctx)
-	err := dbc.conn.ExecuteStreamFetch(query, callback, alloc, streamBufferSize)
+	ch := make(chan error)
+	go func() {
+		ch <- dbc.conn.ExecuteStreamFetch(query, callback, alloc, streamBufferSize)
+	}()
 
-	if done != nil {
-		close(done)
-		wg.Wait()
-	}
-	if dbcerr := dbc.Err(); dbcerr != nil {
-		return dbcerr
+	select {
+	case <-ctx.Done():
+		killCtx, cancel := context.WithTimeout(context.Background(), dbc.killTimeout)
+		defer cancel()
+
+		_ = dbc.KillWithContext(killCtx, ctx.Err().Error(), time.Since(now))
+		return dbc.Err()
+	case err := <-ch:
+		if dbcErr := dbc.Err(); dbcErr != nil {
+			return dbcErr
+		}
+		return err
 	}
-	return err
 }
 
 // StreamOnce executes the query and streams the results. But, does not retry on connection errors.
@@ -363,10 +389,19 @@
 	return dbc.conn.IsClosed()
 }
 
-// Kill kills the currently executing query both on MySQL side
+// Kill wraps KillWithContext using context.Background.
+func (dbc *Conn) Kill(reason string, elapsed time.Duration) error {
+	return dbc.KillWithContext(context.Background(), reason, elapsed)
+}
+
+// KillWithContext kills the currently executing query both on MySQL side
 // and on the connection side. If no query is executing, it's a no-op.
 // Kill will also not kill a query more than once.
-func (dbc *Conn) Kill(reason string, elapsed time.Duration) error {
+func (dbc *Conn) KillWithContext(ctx context.Context, reason string, elapsed time.Duration) error {
+	if cause := context.Cause(ctx); cause != nil {
+		return cause
+	}
+
 	dbc.stats.KillCounters.Add("Queries", 1)
 	log.Infof("Due to %s, elapsed time: %v, killing query ID %v %s", reason, elapsed, dbc.conn.ID(), dbc.CurrentForLogging())
 
@@ -377,25 +412,43 @@
 	dbc.conn.Close()
 
 	// Server side action. Kill the session.
-	killConn, err := dbc.dbaPool.Get(context.TODO())
+	killConn, err := dbc.dbaPool.Get(ctx)
 	if err != nil {
 		log.Warningf("Failed to get conn from dba pool: %v", err)
 		return err
 	}
 	defer killConn.Recycle()
+
+	ch := make(chan error)
 	sql := fmt.Sprintf("kill %d", dbc.conn.ID())
-	_, err = killConn.Conn.ExecuteFetch(sql, 10000, false)
-	if err != nil {
-		log.Errorf("Could not kill query ID %v %s: %v", dbc.conn.ID(),
-			dbc.CurrentForLogging(), err)
-		return err
+	go func() {
+		_, err := killConn.Conn.ExecuteFetch(sql, -1, false)
+		ch <- err
+	}()
+
+	select {
+	case <-ctx.Done():
+		killConn.Close()
+
+		dbc.stats.InternalErrors.Add("HungQuery", 1)
+		log.Warningf("Query may be hung: %s", dbc.CurrentForLogging())
+
+		return context.Cause(ctx)
+	case err := <-ch:
+		if err != nil {
+			log.Errorf("Could not kill query ID %v %s: %v", dbc.conn.ID(), dbc.CurrentForLogging(), err)
+			return err
+		}
+		return nil
 	}
-	return nil
 }
 
 // Current returns the currently executing query.
 func (dbc *Conn) Current() string {
-	return dbc.current.Load().(string)
+	if q := dbc.current.Load(); q != nil {
+		return *q
+	}
+	return ""
 }
 
 // ID returns the connection id.
@@ -437,45 +490,6 @@
 	return nil
 }
 
-// setDeadline starts a goroutine that will kill the currently executing query
-// if the deadline is exceeded. It returns a channel and a waitgroup. After the
-// query is done executing, the caller is required to close the done channel
-// and wait for the waitgroup to make sure that the necessary cleanup is done.
-func (dbc *Conn) setDeadline(ctx context.Context) (chan bool, *sync.WaitGroup) {
-	if ctx.Done() == nil {
-		return nil, nil
-	}
-	done := make(chan bool)
-	var wg sync.WaitGroup
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		startTime := time.Now()
-		select {
-		case <-ctx.Done():
-			dbc.Kill(ctx.Err().Error(), time.Since(startTime))
-		case <-done:
-			return
-		}
-		elapsed := time.Since(startTime)
-
-		// Give 2x the elapsed time and some buffer as grace period
-		// for the query to get killed.
-		tmr2 := time.NewTimer(2*elapsed + 5*time.Second)
-		defer tmr2.Stop()
-		select {
-		case <-tmr2.C:
-			dbc.stats.InternalErrors.Add("HungQuery", 1)
-			log.Warningf("Query may be hung: %s", dbc.CurrentForLogging())
-		case <-done:
-			return
-		}
-		<-done
-		log.Warningf("Hung query returned")
-	}()
-	return done, &wg
-}
-
 // CurrentForLogging applies transformations to the query making it suitable to log.
 // It applies sanitization rules based on tablet settings and limits the max length of
 // queries.