Repository: pingcap/tidb
File: pkg/session/syssession/session.go (Go)

Code before:
	return s.inUse
}

// CheckNoPendingTxn checks whether some txn is not committed or pending for TSO.
// If it happens, it returns an error.
func (s *session) CheckNoPendingTxn() error {
	if s.sctx.GetPreparedTxnFuture() != nil {
		return errors.New("txn is pending for TSO")
	}

	txn, err := s.sctx.Txn(false)
	if err != nil {
		return err
	}

	if txn.Valid() {
		return errors.New("txn is still valid")
	}

	return nil
}

// OwnerResetState resets the state of the session by owner
func (s *session) OwnerResetState(ctx context.Context, caller sessionOwner) error {
	return s.OwnerWithSctx(caller, func(sctx SessionContext) error {
		sctx.RollbackTxn(ctx)
		// Reset time-related session vars to avoid cross-task contamination when sessions are reused.
		globalTZ, err := sctx.GetSessionVars().GetGlobalSystemVar(ctx, vardef.TimeZone)
		if err != nil {
			return err
		}
		if err := sctx.GetSessionVars().SetSystemVar(vardef.TimeZone, globalTZ); err != nil {
			return err
		}
		sctx.GetSessionVars().StmtCtx.SetTimeZone(sctx.GetSessionVars().Location())
		return nil
	})
}

// OwnerWithSctx executes the input function with the session context.
func (s *session) OwnerWithSctx(caller sessionOwner, fn func(SessionContext) error) error {
	intest.AssertNotNil(fn)
	sctx, exit, err := s.EnterOperation(caller, false)
	if err != nil {
		return err
	}
	defer exit()
	return fn(sctx)
}

// EnterOperation enters an operation on the session.

Review comment (suggestion):
OwnerResetState now calls GetSessionVars()/GetGlobalSystemVar() in addition to RollbackTxn(). This introduces new runtime requirements for any mocked SessionContext used in tests (e.g. mocks embedding a nil sessionctx.Context will now panic on GetSessionVars). Please update the syssession unit-test mocks to provide a real SessionVars (or implement GetSessionVars) so OwnerResetState can run without panicking, or adjust the reset logic to avoid calling into a nil embedded Context in tests.

Code after:
}

// CheckNoPendingTxn checks whether some txn is not committed or pending for TSO.
// If it happens, it returns an error.
func (s *session) CheckNoPendingTxn() error {
	if s.sctx.GetPreparedTxnFuture() != nil {
		return errors.New("txn is pending for TSO")
	}

	txn, err := s.sctx.Txn(false)
	if err != nil {
		return err
	}

	if txn.Valid() {
		return errors.New("txn is still valid")
	}

	return nil
}

// OwnerResetState resets the state of the session by owner
func (s *session) OwnerResetState(ctx context.Context, caller sessionOwner) error {
	return s.OwnerWithSctx(caller, func(sctx SessionContext) error {
		sctx.RollbackTxn(ctx)
		return nil
	})
}

// OwnerWithSctx executes the input function with the session context.
func (s *session) OwnerWithSctx(caller sessionOwner, fn func(SessionContext) error) error {
	intest.AssertNotNil(fn)
	sctx, exit, err := s.EnterOperation(caller, false)
	if err != nil {
		return err
	}
	defer exit()
	return fn(sctx)
}

// EnterOperation enters an operation on the session.
// It receives a caller with the type sessionOwner as the argument,
// and if the session is closed or the caller is not the owner, it will return an error.
// If the operation is entered successfully,
// it will return the internal SessionContext for further usage and a function to exit the operation.
func (s *session) EnterOperation(caller sessionOwner, threadSafe bool) (SessionContext, func(), error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	seqStart := s.seqIncWithoutLock()
	if err := s.checkOwnerWithoutLock(caller, "EnterOperation error: "); err != nil {
		s.reportErrorWithoutLock(