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

@@ -4864,13 +4864,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

@@ -5944,6 +5944,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_manager.go

@@ -418,6 +418,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

pkg/frontend/routine_manager_test.go

@@ -27,6 +27,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 
 	"github.com/matrixorigin/matrixone/pkg/common/moerr"
@@ -135,6 +136,22 @@ type blockingCloseConn struct {
 	startOnce    sync.Once
 }
 
+type blockingResponseProtocol struct {
+	MysqlRrWr
+	responseWritten chan struct{}
+	releaseResponse chan struct{}
+	writtenOnce     sync.Once
+}
+
+func (p *blockingResponseProtocol) WriteResponse(ctx context.Context, response *Response) error {
+	err := p.MysqlRrWr.WriteResponse(ctx, response)
+	p.writtenOnce.Do(func() {
+		close(p.responseWritten)
+	})
+	<-p.releaseResponse
+	return err
+}
+
 func (tc *blockingCloseConn) Close() error {
 	tc.startOnce.Do(func() {
 		close(tc.closeStarted)
@@ -756,6 +773,122 @@ func TestRoutineManagerMigrationAndResetErrorBranches(t *testing.T) {
 	}, &query.MigrateConnFromResponse{}), "cannot migrate connection while user-level locks are held")
 }
 
+func TestRoutineManagerResetSessionRejectsRequestAfterResponseWrite(t *testing.T) {
+	const connID = uint32(1009)
+	ctrl := gomock.NewController(t)
+	oldSession := newTestSession(t, ctrl)
+	protocol := &blockingResponseProtocol{
+		MysqlRrWr:       oldSession.GetResponser().MysqlRrWr(),
+		responseWritten: make(chan struct{}),
+		releaseResponse: make(chan struct{}),
+	}
+	routine := NewRoutine(context.Background(), protocol, getPu("").SV)
+	rm, err := NewRoutineManager(context.Background(), "")
+	require.NoError(t, err)
+	rm.sessionManager = queryservice.NewSessionManager()
+
+	oldSession.respr = NewMysqlResp(protocol)
+	oldSession.setRoutineManager(rm)
+	oldSession.setRoutine(routine)
+	routine.setSession(oldSession)
+	rm.sessionManager.AddSession(oldSession)
+	conn := &Conn{id: uint64(connID), conn: &testConn{}, remoteAddr: "remote"}
+	rm.setRoutine(conn, connID, routine)
+
+	var releaseOnce sync.Once
+	handlerFinished := make(chan struct{})
+	handlerResult := make(chan struct {
+		err       error
+		recovered any
+	}, 1)
+	t.Cleanup(func() {
+		releaseOnce.Do(func() {
+			close(protocol.releaseResponse)
+		})
+		select {
+		case <-handlerFinished:
+		case <-time.After(time.Second):
+		}
+		if current := routine.getSession(); current != nil && current.GetProc() != nil {
+			rm.sessionManager.RemoveSession(current)
+			current.Close()
+		}
+		routine.cancelRoutineFunc()
+		rm.cancelCtx()
+	})
+
+	go func() {
+		var result struct {
+			err       error
+			recovered any
+		}
+		defer func() {
+			result.recovered = recover()
+			handlerResult <- result
+			close(handlerFinished)
+		}()
+		result.err = rm.Handler(conn, []byte{byte(COM_PING)})
+	}()
+
+	select {
+	case <-protocol.responseWritten:
+	case <-time.After(time.Second):
+		t.Fatal("request did not write its terminal response")
+	}
+
+	oldProc := oldSession.GetProc()
+	oldTxnHandler := oldSession.GetTxnHandler()
+	err = routine.resetSession("", &query.ResetSessionResponse{})
+	require.ErrorContains(t, err, "cannot reset session as routine is closed or busy")
+	require.Same(t, oldSession, routine.getSession())
+	require.Same(t, oldProc, oldSession.GetProc())
+	require.Same(t, oldTxnHandler, oldSession.GetTxnHandler())
+	registered := rm.sessionManager.GetAllSessions()
+	require.Len(t, registered, 1)
+	require.Same(t, oldSession, registered[0])
+
+	releaseOnce.Do(func() {
+		close(protocol.releaseResponse)
+	})
+	select {
+	case result := <-handlerResult:
+		require.Nil(t, result.recovered)
+		require.NoError(t, result.err)
+	case <-time.After(time.Second):
+		t.Fatal("request handler did not finish after response release")
+	}
+
+	require.NoError(t, routine.resetSession("", &query.ResetSessionResponse{}))
+	newSession := routine.getSession()
+	require.NotSame(t, oldSession, newSession)
+	require.Nil(t, oldSession.GetProc())
+	require.Nil(t, oldSession.GetTxnHandler())
+	registered = rm.sessionManager.GetAllSessions()
+	require.Len(t, registered, 1)
+	require.Same(t, newSession, registered[0])
+	require.NoError(t, rm.Handler(conn, []byte{byte(COM_PING)}))
+}
+
+func TestRoutineManagerHandlerRejectsLifecycleConflictBeforeSessionRead(t *testing.T) {
+	routine := NewRoutine(context.Background(), &testMysqlWriter{}, &config.FrontendParameters{})
+	t.Cleanup(routine.cancelRoutineFunc)
+	require.True(t, routine.mc.tryBeginOperation())
+	defer routine.mc.endOperation()
+
+	conn := &Conn{id: 1010, conn: &testConn{}, remoteAddr: "remote"}
+	rm := &RoutineManager{
+		ctx:              context.Background(),
+		clients:          map[*Conn]*Routine{conn: routine},
+		routinesByConnID: map[uint32]*Routine{1010: routine},
+	}
+
+	require.ErrorContains(
+		t,
+		rm.Handler(conn, []byte{byte(COM_PING)}),
+		"cannot process request as routine is closed or busy",
+	)
+}
+
 func TestRoutineMigrateConnectionFromRejectsUserLevelLocks(t *testing.T) {
 	rt, proto := newUnitTestRoutine(t, 1006)
 	proc := testutil.NewProc(t)