fix(frontend): fence session reset from active requests

pkg/frontend/migrate.go

@@ -30,6 +30,8 @@ type migrateController struct {
 	closed bool
 	// inProgress indicates if a lifecycle operation is in progress.
 	inProgress bool
+	// requestInProgress indicates if a SQL request owns the routine session.
+	requestInProgress bool
 	// operationCancel cancels the active lifecycle operation. It is published
 	// together with inProgress while holding the controller lock.
 	operationCancel context.CancelFunc
@@ -94,7 +96,7 @@ func (mc *migrateController) beginOperationWithContext(ctx context.Context) (con
 
 	mc.Lock()
 	defer mc.Unlock()
-	for mc.inProgress && !mc.closed && ctx.Err() == nil {
+	for (mc.inProgress || mc.requestInProgress) && !mc.closed && ctx.Err() == nil {
 		mc.cond.Wait()
 	}
 	return mc.startOperationLocked(ctx)
@@ -120,7 +122,7 @@ func (mc *migrateController) startOperationLocked(ctx context.Context) (context.
 	if mc.closed || ctx.Err() != nil {
 		return nil, false
 	}
-	if mc.inProgress {
+	if mc.inProgress || mc.requestInProgress {
 		return nil, false
 	}
 	operationCtx, cancel := context.WithCancel(ctx)
@@ -130,6 +132,27 @@ func (mc *migrateController) startOperationLocked(ctx context.Context) (context.
 	return operationCtx, true
 }
 
+// tryBeginRequest acquires the routine session for a request only when no
+// lifecycle operation is active. Requests never wait here: if reset or
+// migration already owns the routine, the caller must fail before reading the
+// session.
+func (mc *migrateController) tryBeginRequest() bool {
+	mc.Lock()
+	defer mc.Unlock()
+	if mc.closed || mc.inProgress || mc.requestInProgress {
+		return false
+	}
+	mc.requestInProgress = true
+	return true
+}
+
+func (mc *migrateController) endRequest() {
+	mc.Lock()
+	defer mc.Unlock()
+	mc.requestInProgress = false
+	mc.cond.Broadcast()
+}
+
 // endOperation completes a lifecycle operation and wakes a routine waiting
 // to close.
 func (mc *migrateController) endOperation() {

pkg/frontend/migrate_test.go

@@ -84,6 +84,95 @@ func TestLifecycleControllerRejectsBusyTryOperation(t *testing.T) {
 	mc.endOperation()
 }
 
+func TestLifecycleControllerRequestAndOperationAreMutuallyExclusive(t *testing.T) {
+	t.Run("request first", func(t *testing.T) {
+		mc := newMigrateController()
+		assert.True(t, mc.tryBeginRequest())
+		assert.False(t, mc.tryBeginRequest())
+		assert.False(t, mc.tryBeginOperation())
+		mc.endRequest()
+		assert.True(t, mc.tryBeginOperation())
+		mc.endOperation()
+	})
+
+	t.Run("operation first", func(t *testing.T) {
+		mc := newMigrateController()
+		assert.True(t, mc.tryBeginOperation())
+		assert.False(t, mc.tryBeginRequest())
+		mc.endOperation()
+		assert.True(t, mc.tryBeginRequest())
+		mc.endRequest()
+	})
+}
+
+func TestLifecycleControllerOperationWaitsForRequest(t *testing.T) {
+	mc := newMigrateController()
+	assert.True(t, mc.tryBeginRequest())
+
+	operationStarted := make(chan bool, 1)
+	go func() {
+		_, ok := mc.beginOperationWithContext(context.Background())
+		operationStarted <- ok
+		if ok {
+			mc.endOperation()
+		}
+	}()
+
+	select {
+	case <-operationStarted:
+		t.Fatal("lifecycle operation started while a request owned the session")
+	case <-time.After(50 * time.Millisecond):
+	}
+	mc.endRequest()
+
+	select {
+	case ok := <-operationStarted:
+		assert.True(t, ok)
+	case <-time.After(time.Second):
+		t.Fatal("lifecycle operation did not start after request completion")
+	}
+}
+
+func TestLifecycleControllerWaitingForRequestHonorsContext(t *testing.T) {
+	mc := newMigrateController()
+	assert.True(t, mc.tryBeginRequest())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	result := make(chan bool, 1)
+	go func() {
+		_, ok := mc.beginOperationWithContext(ctx)
+		result <- ok
+	}()
+	cancel()
+
+	select {
+	case ok := <-result:
+		assert.False(t, ok)
+	case <-time.After(time.Second):
+		t.Fatal("lifecycle operation waiting for a request ignored caller cancellation")
+	}
+	mc.endRequest()
+}
+
+func TestLifecycleControllerCloseDoesNotWaitForRequest(t *testing.T) {
+	mc := newMigrateController()
+	assert.True(t, mc.tryBeginRequest())
+
+	closed := make(chan struct{})
+	go func() {
+		mc.waitAndClose()
+		close(closed)
+	}()
+
+	select {
+	case <-closed:
+	case <-time.After(time.Second):
+		t.Fatal("routine close waited for request completion")
+	}
+	assert.False(t, mc.tryBeginRequest())
+	mc.endRequest()
+}
+
 func TestLifecycleControllerCloseRejectsQueuedOperation(t *testing.T) {
 	mc := newMigrateController()
 	assert.True(t, mc.beginOperation())

pkg/frontend/mysql_cmd_executor.go

@@ -4940,13 +4940,17 @@ func ExecRequest(ses *Session, execCtx *ExecCtx, req *Request) (resp *Response,
 	defer func() {
 		if e := recover(); e != nil {
 			markRowCountFailed(ses, ses.GetProc())
+			var serverStatus uint16
+			if txnHandler := ses.GetTxnHandler(); txnHandler != nil {
+				serverStatus = txnHandler.GetServerStatus()
+			}
 			moe, ok := e.(*moerr.Error)
 			if !ok {
 				err = errors.Join(err, moerr.ConvertPanicError(execCtx.reqCtx, e))
-				resp = NewGeneralErrorResponse(COM_QUERY, ses.txnHandler.GetServerStatus(), err)
+				resp = NewGeneralErrorResponse(COM_QUERY, serverStatus, err)
 			} else {
 				err = errors.Join(err, moe)
-				resp = NewGeneralErrorResponse(COM_QUERY, ses.txnHandler.GetServerStatus(), moe)
+				resp = NewGeneralErrorResponse(COM_QUERY, serverStatus, moe)
 			}
 			// log the query's statement and error info.
 			logStatementStatus(execCtx.reqCtx, ses, execCtx.stmt, fail, err)

pkg/frontend/mysql_cmd_executor_test.go

@@ -6031,6 +6031,35 @@ func Test_panic(t *testing.T) {
 	runPanic(fault.PanicUseNonMoErr)
 }
 
+func TestExecRequestRecoverWithNilTxnHandler(t *testing.T) {
+	fault.EnableDomain(fault.DomainFrontend)
+	defer fault.DisableDomain(fault.DomainFrontend)
+	fault.AddFaultPointInDomain(
+		context.Background(),
+		fault.DomainFrontend,
+		"exec_request_panic",
+		":::",
+		"panic",
+		fault.PanicUseNonMoErr,
+		"has panic",
+		false,
+	)
+	defer fault.RemoveFaultPointFromDomain(context.Background(), fault.DomainFrontend, "exec_request_panic")
+
+	ctrl := gomock.NewController(t)
+	ses := newTestSession(t, ctrl)
+	t.Cleanup(ses.Close)
+	ses.mu.Lock()
+	ses.txnHandler = nil
+	ses.mu.Unlock()
+
+	resp, err := ExecRequest(ses, &ExecCtx{reqCtx: context.Background(), ses: ses}, &Request{cmd: COM_PING})
+	require.Error(t, err)
+	require.NotNil(t, resp)
+	require.Equal(t, ErrorResponse, resp.GetCategory())
+	require.Zero(t, resp.GetStatus())
+}
+
 func Test_run_panic(t *testing.T) {
 	fault.EnableDomain(fault.DomainFrontend)
 	defer fault.DisableDomain(fault.DomainFrontend)

pkg/frontend/routine.go

@@ -556,6 +556,18 @@ func (rt *Routine) migrateConnectionFrom(resp *query.MigrateConnFromResponse) er
 func (rt *Routine) migrateConnectionFromWithContext(
 	ctx context.Context,
 	resp *query.MigrateConnFromResponse,
+) error {
+	action := query.MigrateConnFromAction_MigrateConnFromExport
+	if resp == nil {
+		action = query.MigrateConnFromAction_MigrateConnFromSkipUserLevelLockRelease
+	}
+	return rt.migrateConnectionFromActionWithContext(ctx, action, resp)
+}
+
+func (rt *Routine) migrateConnectionFromActionWithContext(
+	ctx context.Context,
+	action query.MigrateConnFromAction,
+	resp *query.MigrateConnFromResponse,
 ) error {
 	operationCtx, ok := rt.mc.beginOperationWithContext(ctx)
 	if !ok {
@@ -568,13 +580,20 @@ func (rt *Routine) migrateConnectionFromWithContext(
 	}
 	defer rt.mc.endOperation()
 
+	if cause := context.Cause(operationCtx); cause != nil {
+		return cause
+	}
 	ses := rt.getSession()
-	if resp == nil {
+	switch action {
+	case query.MigrateConnFromAction_MigrateConnFromSkipUserLevelLockRelease:
+		if states := function.UserLevelLocksForMigration(ses.proc); len(states) > 0 {
+			return moerr.NewInternalErrorNoCtx("cannot migrate connection while user-level locks are held")
+		}
 		ses.userLevelLocksMigrated = true
 		return nil
-	}
-	if cause := context.Cause(operationCtx); cause != nil {
-		return cause
+	case query.MigrateConnFromAction_MigrateConnFromEnableUserLevelLockRelease:
+		ses.userLevelLocksMigrated = false
+		return nil
 	}
 	if states := function.UserLevelLocksForMigration(ses.proc); len(states) > 0 {
 		return moerr.NewInternalErrorNoCtx("cannot migrate connection while user-level locks are held")

pkg/frontend/routine_manager.go

@@ -33,7 +33,6 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/logutil"
 	"github.com/matrixorigin/matrixone/pkg/pb/query"
 	"github.com/matrixorigin/matrixone/pkg/queryservice"
-	"github.com/matrixorigin/matrixone/pkg/sql/plan/function"
 	"github.com/matrixorigin/matrixone/pkg/util/metric"
 	v2 "github.com/matrixorigin/matrixone/pkg/util/metric/v2"
 	"github.com/matrixorigin/matrixone/pkg/util/trace"
@@ -418,6 +417,10 @@ func (rm *RoutineManager) Handler(rs *Conn, msg []byte) error {
 		logutil.Errorf("%s error:%v", connectionInfo, err)
 		return err
 	}
+	if !routine.mc.tryBeginRequest() {
+		return moerr.NewInternalError(ctx, "cannot process request as routine is closed or busy")
+	}
+	defer routine.mc.endRequest()
 	routine.setInProcessRequest(true)
 	defer routine.setInProcessRequest(false)
 	payload := msg
@@ -506,18 +509,7 @@ func (rm *RoutineManager) MigrateConnectionFromWithContext(
 	if routine == nil {
 		return moerr.NewInternalErrorf(rm.ctx, "cannot get routine to migrate connection %d", req.ConnID)
 	}
-	switch req.Action {
-	case query.MigrateConnFromAction_MigrateConnFromSkipUserLevelLockRelease:
-		if states := function.UserLevelLocksForMigration(routine.getSession().proc); len(states) > 0 {
-			return moerr.NewInternalErrorNoCtx("cannot migrate connection while user-level locks are held")
-		}
-		return routine.migrateConnectionFromWithContext(ctx, nil)
-	case query.MigrateConnFromAction_MigrateConnFromEnableUserLevelLockRelease:
-		routine.getSession().userLevelLocksMigrated = false
-		return nil
-	default:
-		return routine.migrateConnectionFromWithContext(ctx, resp)
-	}
+	return routine.migrateConnectionFromActionWithContext(ctx, req.Action, resp)
 }
 
 func (rm *RoutineManager) ResetSession(req *query.ResetSessionRequest, resp *query.ResetSessionResponse) error {