func (s *KVSnapshot) SetReadReplicaScope(scope string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.mu.readReplicaScope = scope
}
SetReadReplicaScope set read replica scope
package txnsnapshot
import (
	"bytes"
	"context"
	"fmt"
	"math"
	"sync"
	"sync/atomic"
	"time"
	"github.com/opentracing/opentracing-go"
	"github.com/pingcap/errors"
	"github.com/pingcap/failpoint"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/kvproto/pkg/metapb"
	tikverr "github.com/tikv/client-go/v2/error"
	"github.com/tikv/client-go/v2/internal/client"
	"github.com/tikv/client-go/v2/internal/locate"
	"github.com/tikv/client-go/v2/internal/logutil"
	"github.com/tikv/client-go/v2/internal/retry"
	"github.com/tikv/client-go/v2/internal/unionstore"
	"github.com/tikv/client-go/v2/kv"
	"github.com/tikv/client-go/v2/metrics"
	"github.com/tikv/client-go/v2/oracle"
	"github.com/tikv/client-go/v2/tikvrpc"
	"github.com/tikv/client-go/v2/txnkv/txnlock"
	"github.com/tikv/client-go/v2/txnkv/txnutil"
	"github.com/tikv/client-go/v2/util"
	"go.uber.org/zap"
)
const (
	defaultScanBatchSize = 256
	batchGetSize         = 5120
	maxTimestamp         = math.MaxUint64
)
type IsoLevel kvrpcpb.IsolationLevel
const (
	SI IsoLevel = IsoLevel(kvrpcpb.IsolationLevel_SI)
	RC IsoLevel = IsoLevel(kvrpcpb.IsolationLevel_RC)
)
func (l IsoLevel) ToPB() kvrpcpb.IsolationLevel {
	return kvrpcpb.IsolationLevel(l)
}
type kvstore interface {
	CheckVisibility(startTime uint64) error
	GetRegionCache() *locate.RegionCache
	GetLockResolver() *txnlock.LockResolver
	GetTiKVClient() (client client.Client)
	SendReq(bo *retry.Backoffer, req *tikvrpc.Request, regionID locate.RegionVerID, timeout time.Duration) (*tikvrpc.Response, error)
	GetOracle() oracle.Oracle
}
type KVSnapshot struct {
	store           kvstore
	version         uint64
	isolationLevel  IsoLevel
	priority        txnutil.Priority
	notFillCache    bool
	keyOnly         bool
	vars            *kv.Variables
	replicaReadSeed uint32
	resolvedLocks   util.TSSet
	scanBatchSize   int
	mu struct {
		sync.RWMutex
		hitCnt           int64
		cached           map[string][]byte
		cachedSize       int
		stats            *SnapshotRuntimeStats
		replicaRead      kv.ReplicaReadType
		taskID           uint64
		isStaleness      bool
		readReplicaScope string
		matchStoreLabels []*metapb.StoreLabel
	}
	sampleStep uint32
	resourceGroupTag []byte
}
func NewTiKVSnapshot(store kvstore, ts uint64, replicaReadSeed uint32) *KVSnapshot {
	if ts >= math.MaxInt64 && ts != math.MaxUint64 {
		err := errors.Errorf("try to get snapshot with a large ts %d", ts)
		panic(err)
	}
	return &KVSnapshot{
		store:           store,
		version:         ts,
		scanBatchSize:   defaultScanBatchSize,
		priority:        txnutil.PriorityNormal,
		vars:            kv.DefaultVars,
		replicaReadSeed: replicaReadSeed,
	}
}
const batchGetMaxBackoff = 600000  
func (s *KVSnapshot) SetSnapshotTS(ts uint64) {
	if ts >= math.MaxInt64 && ts != math.MaxUint64 {
		err := errors.Errorf("try to get snapshot with a large ts %d", ts)
		panic(err)
	}
	s.version = ts
	s.mu.Lock()
	s.mu.cached = nil
	s.mu.Unlock()
	s.resolvedLocks = util.TSSet{}
}
func (s *KVSnapshot) BatchGet(ctx context.Context, keys [][]byte) (map[string][]byte, error) {
	m := make(map[string][]byte)
	s.mu.RLock()
	if s.mu.cached != nil {
		tmp := make([][]byte, 0, len(keys))
		for _, key := range keys {
			if val, ok := s.mu.cached[string(key)]; ok {
				atomic.AddInt64(&s.mu.hitCnt, 1)
				if len(val) > 0 {
					m[string(key)] = val
				}
			} else {
				tmp = append(tmp, key)
			}
		}
		keys = tmp
	}
	s.mu.RUnlock()
	if len(keys) == 0 {
		return m, nil
	}
	ctx = context.WithValue(ctx, retry.TxnStartKey, s.version)
	bo := retry.NewBackofferWithVars(ctx, batchGetMaxBackoff, s.vars)
	var mu sync.Mutex
	err := s.batchGetKeysByRegions(bo, keys, func(k, v []byte) {
		if len(v) == 0 {
			return
		}
		mu.Lock()
		m[string(k)] = v
		mu.Unlock()
	})
	s.recordBackoffInfo(bo)
	if err != nil {
		return nil, err
	}
	err = s.store.CheckVisibility(s.version)
	if err != nil {
		return nil, err
	}
	s.mu.Lock()
	if s.mu.cached == nil {
		s.mu.cached = make(map[string][]byte, len(m))
	}
	for _, key := range keys {
		val := m[string(key)]
		s.mu.cachedSize += len(key) + len(val)
		s.mu.cached[string(key)] = val
	}
	const cachedSizeLimit = 10 << 30
	if s.mu.cachedSize >= cachedSizeLimit {
		for k, v := range s.mu.cached {
			if _, needed := m[k]; needed {
				continue
			}
			delete(s.mu.cached, k)
			s.mu.cachedSize -= len(k) + len(v)
			if s.mu.cachedSize < cachedSizeLimit {
				break
			}
		}
	}
	s.mu.Unlock()
	return m, nil
}
type batchKeys struct {
	region locate.RegionVerID
	keys   [][]byte
}
func (b *batchKeys) relocate(bo *retry.Backoffer, c *locate.RegionCache) (bool, error) {
	loc, err := c.LocateKey(bo, b.keys[0])
	if err != nil {
		return false, err
	}
	for i := 1; i < len(b.keys); i++ {
		if !loc.Contains(b.keys[i]) {
			return false, nil
		}
	}
	b.region = loc.Region
	return true, nil
}
func appendBatchKeysBySize(b []batchKeys, region locate.RegionVerID, keys [][]byte, sizeFn func([]byte) int, limit int) []batchKeys {
	var start, end int
	for start = 0; start < len(keys); start = end {
		var size int
		for end = start; end < len(keys) && size < limit; end++ {
			size += sizeFn(keys[end])
		}
		b = append(b, batchKeys{
			region: region,
			keys:   keys[start:end],
		})
	}
	return b
}
func (s *KVSnapshot) batchGetKeysByRegions(bo *retry.Backoffer, keys [][]byte, collectF func(k, v []byte)) error {
	defer func(start time.Time) {
		metrics.TxnCmdHistogramWithBatchGet.Observe(time.Since(start).Seconds())
	}(time.Now())
	groups, _, err := s.store.GetRegionCache().GroupKeysByRegion(bo, keys, nil)
	if err != nil {
		return err
	}
	metrics.TxnRegionsNumHistogramWithSnapshot.Observe(float64(len(groups)))
	var batches []batchKeys
	for id, g := range groups {
		batches = appendBatchKeysBySize(batches, id, g, func([]byte) int { return 1 }, batchGetSize)
	}
	if len(batches) == 0 {
		return nil
	}
	if len(batches) == 1 {
		return s.batchGetSingleRegion(bo, batches[0], collectF)
	}
	ch := make(chan error)
	for _, batch1 := range batches {
		batch := batch1
		go func() {
			backoffer, cancel := bo.Fork()
			defer cancel()
			ch <- s.batchGetSingleRegion(backoffer, batch, collectF)
		}()
	}
	for i := 0; i < len(batches); i++ {
		if e := <-ch; e != nil {
			logutil.BgLogger().Debug("snapshot batchGet failed",
				zap.Error(e),
				zap.Uint64("txnStartTS", s.version))
			err = errors.WithStack(e)
		}
	}
	return err
}
func (s *KVSnapshot) batchGetSingleRegion(bo *retry.Backoffer, batch batchKeys, collectF func(k, v []byte)) error {
	cli := NewClientHelper(s.store, &s.resolvedLocks, false)
	s.mu.RLock()
	if s.mu.stats != nil {
		cli.Stats = make(map[tikvrpc.CmdType]*locate.RPCRuntimeStats)
		defer func() {
			s.mergeRegionRequestStats(cli.Stats)
		}()
	}
	s.mu.RUnlock()
	pending := batch.keys
	for {
		s.mu.RLock()
		req := tikvrpc.NewReplicaReadRequest(tikvrpc.CmdBatchGet, &kvrpcpb.BatchGetRequest{
			Keys:    pending,
			Version: s.version,
		}, s.mu.replicaRead, &s.replicaReadSeed, kvrpcpb.Context{
			Priority:         s.priority.ToPB(),
			NotFillCache:     s.notFillCache,
			TaskId:           s.mu.taskID,
			ResourceGroupTag: s.resourceGroupTag,
		})
		scope := s.mu.readReplicaScope
		isStaleness := s.mu.isStaleness
		matchStoreLabels := s.mu.matchStoreLabels
		s.mu.RUnlock()
		req.TxnScope = scope
		req.ReadReplicaScope = scope
		if isStaleness {
			req.EnableStaleRead()
		}
		ops := make([]locate.StoreSelectorOption, 0, 2)
		if len(matchStoreLabels) > 0 {
			ops = append(ops, locate.WithMatchLabels(matchStoreLabels))
		}
		resp, _, _, err := cli.SendReqCtx(bo, req, batch.region, client.ReadTimeoutMedium, tikvrpc.TiKV, "", ops...)
		if err != nil {
			return err
		}
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return err
		}
		if regionErr != nil {
			if regionErr.GetEpochNotMatch() == nil || locate.IsFakeRegionError(regionErr) {
				err = bo.Backoff(retry.BoRegionMiss, errors.New(regionErr.String()))
				if err != nil {
					return err
				}
			}
			same, err := batch.relocate(bo, cli.regionCache)
			if err != nil {
				return err
			}
			if same {
				continue
			}
			return s.batchGetKeysByRegions(bo, pending, collectF)
		}
		if resp.Resp == nil {
			return errors.WithStack(tikverr.ErrBodyMissing)
		}
		batchGetResp := resp.Resp.(*kvrpcpb.BatchGetResponse)
		var (
			lockedKeys [][]byte
			locks      []*txnlock.Lock
		)
		if keyErr := batchGetResp.GetError(); keyErr != nil {
			lock, err := txnlock.ExtractLockFromKeyErr(keyErr)
			if err != nil {
				return err
			}
			lockedKeys = append(lockedKeys, lock.Key)
			locks = append(locks, lock)
		} else {
			for _, pair := range batchGetResp.Pairs {
				keyErr := pair.GetError()
				if keyErr == nil {
					collectF(pair.GetKey(), pair.GetValue())
					continue
				}
				lock, err := txnlock.ExtractLockFromKeyErr(keyErr)
				if err != nil {
					return err
				}
				lockedKeys = append(lockedKeys, lock.Key)
				locks = append(locks, lock)
			}
		}
		if batchGetResp.ExecDetailsV2 != nil {
			readKeys := len(batchGetResp.Pairs)
			readTime := float64(batchGetResp.ExecDetailsV2.GetTimeDetail().GetKvReadWallTimeMs() / 1000)
			readSize := float64(batchGetResp.ExecDetailsV2.GetScanDetailV2().GetProcessedVersionsSize())
			metrics.ObserveReadSLI(uint64(readKeys), readTime, readSize)
			s.mergeExecDetail(batchGetResp.ExecDetailsV2)
		}
		if len(lockedKeys) > 0 {
			msBeforeExpired, err := cli.ResolveLocks(bo, s.version, locks)
			if err != nil {
				return err
			}
			if msBeforeExpired > 0 {
				err = bo.BackoffWithMaxSleepTxnLockFast(int(msBeforeExpired), errors.Errorf("batchGet lockedKeys: %d", len(lockedKeys)))
				if err != nil {
					return err
				}
			}
			if batchGetResp.GetError() == nil {
				pending = lockedKeys
			}
			continue
		}
		return nil
	}
}
const getMaxBackoff = 600000  
func (s *KVSnapshot) Get(ctx context.Context, k []byte) ([]byte, error) {
	defer func(start time.Time) {
		metrics.TxnCmdHistogramWithGet.Observe(time.Since(start).Seconds())
	}(time.Now())
	ctx = context.WithValue(ctx, retry.TxnStartKey, s.version)
	bo := retry.NewBackofferWithVars(ctx, getMaxBackoff, s.vars)
	val, err := s.get(ctx, bo, k)
	s.recordBackoffInfo(bo)
	if err != nil {
		return nil, err
	}
	err = s.store.CheckVisibility(s.version)
	if err != nil {
		return nil, err
	}
	if len(val) == 0 {
		return nil, tikverr.ErrNotExist
	}
	return val, nil
}
func (s *KVSnapshot) get(ctx context.Context, bo *retry.Backoffer, k []byte) ([]byte, error) {
	s.mu.RLock()
	if s.mu.cached != nil {
		if value, ok := s.mu.cached[string(k)]; ok {
			atomic.AddInt64(&s.mu.hitCnt, 1)
			s.mu.RUnlock()
			return value, nil
		}
	}
	s.mu.RUnlock()
	if span := opentracing.SpanFromContext(ctx); span != nil && span.Tracer() != nil {
		span1 := span.Tracer().StartSpan("tikvSnapshot.get", opentracing.ChildOf(span.Context()))
		defer span1.Finish()
		opentracing.ContextWithSpan(ctx, span1)
	}
	failpoint.Inject("snapshot-get-cache-fail", func(_ failpoint.Value) {
		if bo.GetCtx().Value("TestSnapshotCache") != nil {
			panic("cache miss")
		}
	})
	cli := NewClientHelper(s.store, &s.resolvedLocks, true)
	s.mu.RLock()
	if s.mu.stats != nil {
		cli.Stats = make(map[tikvrpc.CmdType]*locate.RPCRuntimeStats)
		defer func() {
			s.mergeRegionRequestStats(cli.Stats)
		}()
	}
	req := tikvrpc.NewReplicaReadRequest(tikvrpc.CmdGet,
		&kvrpcpb.GetRequest{
			Key:     k,
			Version: s.version,
		}, s.mu.replicaRead, &s.replicaReadSeed, kvrpcpb.Context{
			Priority:         s.priority.ToPB(),
			NotFillCache:     s.notFillCache,
			TaskId:           s.mu.taskID,
			ResourceGroupTag: s.resourceGroupTag,
		})
	isStaleness := s.mu.isStaleness
	matchStoreLabels := s.mu.matchStoreLabels
	scope := s.mu.readReplicaScope
	s.mu.RUnlock()
	req.TxnScope = scope
	req.ReadReplicaScope = scope
	var ops []locate.StoreSelectorOption
	if isStaleness {
		req.EnableStaleRead()
	}
	if len(matchStoreLabels) > 0 {
		ops = append(ops, locate.WithMatchLabels(matchStoreLabels))
	}
	var firstLock *txnlock.Lock
	for {
		util.EvalFailpoint("beforeSendPointGet")
		loc, err := s.store.GetRegionCache().LocateKey(bo, k)
		if err != nil {
			return nil, err
		}
		resp, _, _, err := cli.SendReqCtx(bo, req, loc.Region, client.ReadTimeoutShort, tikvrpc.TiKV, "", ops...)
		if err != nil {
			return nil, err
		}
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return nil, err
		}
		if regionErr != nil {
			if regionErr.GetEpochNotMatch() == nil || locate.IsFakeRegionError(regionErr) {
				err = bo.Backoff(retry.BoRegionMiss, errors.New(regionErr.String()))
				if err != nil {
					return nil, err
				}
			}
			continue
		}
		if resp.Resp == nil {
			return nil, errors.WithStack(tikverr.ErrBodyMissing)
		}
		cmdGetResp := resp.Resp.(*kvrpcpb.GetResponse)
		if cmdGetResp.ExecDetailsV2 != nil {
			readKeys := len(cmdGetResp.Value)
			readTime := float64(cmdGetResp.ExecDetailsV2.GetTimeDetail().GetKvReadWallTimeMs() / 1000)
			readSize := float64(cmdGetResp.ExecDetailsV2.GetScanDetailV2().GetProcessedVersionsSize())
			metrics.ObserveReadSLI(uint64(readKeys), readTime, readSize)
			s.mergeExecDetail(cmdGetResp.ExecDetailsV2)
		}
		val := cmdGetResp.GetValue()
		if keyErr := cmdGetResp.GetError(); keyErr != nil {
			lock, err := txnlock.ExtractLockFromKeyErr(keyErr)
			if err != nil {
				return nil, err
			}
			if firstLock == nil {
				firstLock = lock
			} else if s.version == maxTimestamp && firstLock.TxnID != lock.TxnID {
				cli.resolvedLocks.Put(lock.TxnID)
				continue
			}
			msBeforeExpired, err := cli.ResolveLocks(bo, s.version, []*txnlock.Lock{lock})
			if err != nil {
				return nil, err
			}
			if msBeforeExpired > 0 {
				err = bo.BackoffWithMaxSleepTxnLockFast(int(msBeforeExpired), errors.New(keyErr.String()))
				if err != nil {
					return nil, err
				}
			}
			continue
		}
		return val, nil
	}
}
func (s *KVSnapshot) mergeExecDetail(detail *kvrpcpb.ExecDetailsV2) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if detail == nil || s.mu.stats == nil {
		return
	}
	if s.mu.stats.scanDetail == nil {
		s.mu.stats.scanDetail = &util.ScanDetail{}
	}
	if s.mu.stats.timeDetail == nil {
		s.mu.stats.timeDetail = &util.TimeDetail{}
	}
	s.mu.stats.scanDetail.MergeFromScanDetailV2(detail.ScanDetailV2)
	s.mu.stats.timeDetail.MergeFromTimeDetail(detail.TimeDetail)
}
func (s *KVSnapshot) Iter(k []byte, upperBound []byte) (unionstore.Iterator, error) {
	scanner, err := newScanner(s, k, upperBound, s.scanBatchSize, false)
	return scanner, err
}
func (s *KVSnapshot) IterReverse(k []byte) (unionstore.Iterator, error) {
	scanner, err := newScanner(s, nil, k, s.scanBatchSize, true)
	return scanner, err
}
func (s *KVSnapshot) SetNotFillCache(b bool) {
	s.notFillCache = b
}
func (s *KVSnapshot) SetKeyOnly(b bool) {
	s.keyOnly = b
}
func (s *KVSnapshot) SetScanBatchSize(batchSize int) {
	s.scanBatchSize = batchSize
}
func (s *KVSnapshot) SetReplicaRead(readType kv.ReplicaReadType) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.mu.replicaRead = readType
}
func (s *KVSnapshot) SetIsolationLevel(level IsoLevel) {
	s.isolationLevel = level
}
func (s *KVSnapshot) SetSampleStep(step uint32) {
	s.sampleStep = step
}
func (s *KVSnapshot) SetPriority(pri txnutil.Priority) {
	s.priority = pri
}
func (s *KVSnapshot) SetTaskID(id uint64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.mu.taskID = id
}
func (s *KVSnapshot) SetRuntimeStats(stats *SnapshotRuntimeStats) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.mu.stats = stats
}
func (s *KVSnapshot) SetTxnScope(scope string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.mu.readReplicaScope = scope
}