Repository: pingcap/tidb
File: pkg/sessiontxn/isolation/base.go (Go)

Code before:
	})

	oracleStore := sctx.GetStore().GetOracle()
	option := &oracle.Option{TxnScope: scope}

	var innerFuture oracle.Future
	if sctx.GetSessionVars().UseLowResolutionTSO() {
		innerFuture = oracleStore.GetLowResolutionTimestampAsync(ctx, option)
	} else {
		innerFuture = oracleStore.GetTimestampAsync(ctx, option)
	}

	// Wrap the future to inject delay for testing maxExecutionTime behavior
	return &delayedOracleFuture{inner: innerFuture}
}

// funcFuture implements oracle.Future
type funcFuture func() (uint64, error)

// Wait returns a ts got from the func
func (f funcFuture) Wait() (uint64, error) {
	return f()
}

// delayedOracleFuture wraps an oracle.Future to inject delay before Wait() for testing
type delayedOracleFuture struct {
	inner oracle.Future
}

// Wait injects delay before calling inner Wait() if failpoint is enabled
func (f *delayedOracleFuture) Wait() (uint64, error) {
	// Inject delay before waiting for TSO to test maxExecutionTime behavior
	failpoint.Inject("injectTSOWaitDelay", func(val failpoint.Value) {
		if delayMs, ok := val.(int); ok {
			time.Sleep(time.Duration(delayMs) * time.Millisecond)
		}
	})
	return f.inner.Wait()
}

// basePessimisticTxnContextProvider extends baseTxnContextProvider with some functionalities that are commonly used in
// pessimistic transactions.
type basePessimisticTxnContextProvider struct {
	baseTxnContextProvider
}

// OnPessimisticStmtStart is the hook that should be called when starts handling a pessimistic DML or
// a pessimistic select-for-update statements.
func (p *basePessimisticTxnContextProvider) OnPessimisticStmtStart(ctx context.Context) error {
	if err := p.baseTxnContextProvider.OnPessimisticStmtStart(ctx); err != nil {
		return err

Review comment (bug):
The `delayedOracleFuture` wrapper introduces test-only logic into the production path, adding an extra indirection for every TSO fetch.

A cleaner approach might be to inject the failpoint directly at the `Wait()` call site in `ActivateTxn()`, instead of wrapping the entire oracle. This would keep the production code simpler and better isolate test-only behavior. The code is as follows:
```
func (p *baseTxnContextProvider) ActivateTxn() (kv.Transaction, error) {
    // ...
    txnFuture := p.sctx.GetPreparedTxnFuture()
    
    // Inject delay before TSO wait for testing maxExecutionTime
    failpoint.Inject("injectTSOWaitDelay", func(val failpoint.Value) {
        if delayMs, ok := val.(int); ok {
            time.Sleep(time.Duration(delayMs) * time.Millisecond)
        }
    })
    
    txn, err := txnFuture.Wait(p.ctx, p.sctx)
    // ...
}
```

Code after:

// newOracleFuture creates new future according to the scope and the session context
func newOracleFuture(ctx context.Context, sctx sessionctx.Context, scope string) oracle.Future {
	r, ctx := tracing.StartRegionEx(ctx, "isolation.newOracleFuture")
	defer r.End()

	failpoint.Inject("requestTsoFromPD", func() {
		sessiontxn.TsoRequestCountInc(sctx)
	})

	oracleStore := sctx.GetStore().GetOracle()
	option := &oracle.Option{TxnScope: scope}

	if sctx.GetSessionVars().UseLowResolutionTSO() {
		return oracleStore.GetLowResolutionTimestampAsync(ctx, option)
	}
	return oracleStore.GetTimestampAsync(ctx, option)
}

// funcFuture implements oracle.Future
type funcFuture func() (uint64, error)

// Wait returns a ts got from the func
func (f funcFuture) Wait() (uint64, error) {
	return f()
}

// basePessimisticTxnContextProvider extends baseTxnContextProvider with some functionalities that are commonly used in
// pessimistic transactions.
type basePessimisticTxnContextProvider struct {
	baseTxnContextProvider
}

// OnPessimisticStmtStart is the hook that should be called when starts handling a pessimistic DML or
// a pessimistic select-for-update statements.
func (p *basePessimisticTxnContextProvider) OnPessimisticStmtStart(ctx context.Context) error {
	if err := p.baseTxnContextProvider.OnPessimisticStmtStart(ctx); err != nil {
		return err
	}
	if p.sctx.GetSessionVars().PessimisticTransactionFairLocking &&
		p.txn != nil &&
		p.sctx.GetSessionVars().ConnectionID != 0 &&
		!p.sctx.GetSessionVars().InRestrictedSQL {
		if err := p.txn.StartFairLocking(); err != nil {
			return err
		}
	}
	return nil
}

// OnPessimisticStmtEnd is the hook that should be called when finishes handling a pessimistic DML or