diff --git a/pkg/frontend/migrate.go b/pkg/frontend/migrate.go index cdd7780be6ad1..6606883386042 100644 --- a/pkg/frontend/migrate.go +++ b/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() { diff --git a/pkg/frontend/migrate_test.go b/pkg/frontend/migrate_test.go index ef276bc03ac80..7f752ce4270c8 100644 --- a/pkg/frontend/migrate_test.go +++ b/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()) diff --git a/pkg/frontend/mysql_cmd_executor.go b/pkg/frontend/mysql_cmd_executor.go index 3b5934de8a101..1163719f23245 100644 --- a/pkg/frontend/mysql_cmd_executor.go +++ b/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) diff --git a/pkg/frontend/mysql_cmd_executor_test.go b/pkg/frontend/mysql_cmd_executor_test.go index 6ffcc9bc6f2de..d1f071635672f 100644 --- a/pkg/frontend/mysql_cmd_executor_test.go +++ b/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) diff --git a/pkg/frontend/routine.go b/pkg/frontend/routine.go index 75dc877f3febf..9379f945d1705 100644 --- a/pkg/frontend/routine.go +++ b/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") diff --git a/pkg/frontend/routine_manager.go b/pkg/frontend/routine_manager.go index 4efb265f299e3..158c0b1b2faa8 100644 --- a/pkg/frontend/routine_manager.go +++ b/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 { diff --git a/pkg/frontend/routine_manager_test.go b/pkg/frontend/routine_manager_test.go index c8281e57c46fb..b122ca3931f39 100644 --- a/pkg/frontend/routine_manager_test.go +++ b/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,267 @@ func TestRoutineManagerMigrationAndResetErrorBranches(t *testing.T) { }, &query.MigrateConnFromResponse{}), "cannot migrate connection while user-level locks are held") } +func TestRoutineManagerLegacyMigrationActionsWaitForRequest(t *testing.T) { + t.Run("skip user lock release rechecks after request", func(t *testing.T) { + rm, routine, ses := newLegacyMigrationActionTestFixture(t, 1011) + require.True(t, routine.mc.tryBeginRequest()) + var releaseRequest sync.Once + t.Cleanup(func() { releaseRequest.Do(routine.mc.endRequest) }) + + result := startLegacyMigrationAction(rm, 1011, + query.MigrateConnFromAction_MigrateConnFromSkipUserLevelLockRelease) + requireLegacyMigrationActionPending(t, result) + + function.RestoreUserLevelLocksFromMigration(ses.proc, []function.UserLevelLockState{ + {Name: "request_acquired_lock", Count: 1}, + }) + defer function.DiscardMigratedUserLevelLocks(ses.proc) + releaseRequest.Do(routine.mc.endRequest) + + require.ErrorContains(t, receiveLegacyMigrationActionResult(t, result), + "cannot migrate connection while user-level locks are held") + require.False(t, ses.userLevelLocksMigrated) + }) + + t.Run("enable user lock release waits for request", func(t *testing.T) { + rm, routine, ses := newLegacyMigrationActionTestFixture(t, 1012) + ses.userLevelLocksMigrated = true + require.True(t, routine.mc.tryBeginRequest()) + var releaseRequest sync.Once + t.Cleanup(func() { releaseRequest.Do(routine.mc.endRequest) }) + + result := startLegacyMigrationAction(rm, 1012, + query.MigrateConnFromAction_MigrateConnFromEnableUserLevelLockRelease) + requireLegacyMigrationActionPending(t, result) + require.True(t, ses.userLevelLocksMigrated) + + releaseRequest.Do(routine.mc.endRequest) + require.NoError(t, receiveLegacyMigrationActionResult(t, result)) + require.False(t, ses.userLevelLocksMigrated) + }) +} + +func TestRoutineManagerLegacyMigrationActionsWaitForReset(t *testing.T) { + t.Run("skip user lock release checks replacement session", func(t *testing.T) { + rm, routine, oldSession := newLegacyMigrationActionTestFixture(t, 1013) + require.True(t, routine.mc.tryBeginOperation()) + var releaseReset sync.Once + t.Cleanup(func() { releaseReset.Do(routine.mc.endOperation) }) + + result := startLegacyMigrationAction(rm, 1013, + query.MigrateConnFromAction_MigrateConnFromSkipUserLevelLockRelease) + requireLegacyMigrationActionPending(t, result) + + newSession := newLegacyMigrationActionTestSession(t, 1013) + function.RestoreUserLevelLocksFromMigration(newSession.proc, []function.UserLevelLockState{ + {Name: "replacement_session_lock", Count: 1}, + }) + defer function.DiscardMigratedUserLevelLocks(newSession.proc) + routine.setSession(newSession) + releaseReset.Do(routine.mc.endOperation) + + require.ErrorContains(t, receiveLegacyMigrationActionResult(t, result), + "cannot migrate connection while user-level locks are held") + require.False(t, oldSession.userLevelLocksMigrated) + require.False(t, newSession.userLevelLocksMigrated) + }) + + t.Run("enable user lock release mutates replacement session", func(t *testing.T) { + rm, routine, oldSession := newLegacyMigrationActionTestFixture(t, 1014) + oldSession.userLevelLocksMigrated = true + require.True(t, routine.mc.tryBeginOperation()) + var releaseReset sync.Once + t.Cleanup(func() { releaseReset.Do(routine.mc.endOperation) }) + + result := startLegacyMigrationAction(rm, 1014, + query.MigrateConnFromAction_MigrateConnFromEnableUserLevelLockRelease) + requireLegacyMigrationActionPending(t, result) + require.True(t, oldSession.userLevelLocksMigrated) + + newSession := newLegacyMigrationActionTestSession(t, 1014) + newSession.userLevelLocksMigrated = true + routine.setSession(newSession) + releaseReset.Do(routine.mc.endOperation) + + require.NoError(t, receiveLegacyMigrationActionResult(t, result)) + require.True(t, oldSession.userLevelLocksMigrated) + require.False(t, newSession.userLevelLocksMigrated) + }) +} + +func newLegacyMigrationActionTestFixture( + t *testing.T, + connID uint32, +) (*RoutineManager, *Routine, *Session) { + t.Helper() + routine := NewRoutine(context.Background(), &testMysqlWriter{}, &config.FrontendParameters{}) + t.Cleanup(routine.cancelRoutineFunc) + ses := newLegacyMigrationActionTestSession(t, connID) + routine.setSession(ses) + return &RoutineManager{ + ctx: context.Background(), + routinesByConnID: map[uint32]*Routine{connID: routine}, + }, routine, ses +} + +func newLegacyMigrationActionTestSession(t *testing.T, connID uint32) *Session { + t.Helper() + proc := testutil.NewProc(t) + proc.GetSessionInfo().Account = "legacy_migration_action" + proc.GetSessionInfo().ConnectionID = uint64(connID) + return &Session{proc: proc} +} + +func startLegacyMigrationAction( + rm *RoutineManager, + connID uint32, + action query.MigrateConnFromAction, +) <-chan error { + result := make(chan error, 1) + go func() { + result <- rm.MigrateConnectionFromWithContext(context.Background(), + &query.MigrateConnFromRequest{ConnID: connID, Action: action}, + &query.MigrateConnFromResponse{}) + }() + return result +} + +func requireLegacyMigrationActionPending(t *testing.T, result <-chan error) { + t.Helper() + select { + case err := <-result: + require.Failf(t, "legacy migration action bypassed lifecycle admission", "err=%v", err) + case <-time.After(100 * time.Millisecond): + } +} + +func receiveLegacyMigrationActionResult(t *testing.T, result <-chan error) error { + t.Helper() + select { + case err := <-result: + return err + case <-time.After(time.Second): + require.FailNow(t, "legacy migration action did not finish after lifecycle release") + return nil + } +} + +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)