func (c *client) tryConnect(dispatcherCtx context.Context, dc string) (connectionContext, error) {
	var (
		url           string
		networkErrNum uint64
		err           error
		cc            *grpc.ClientConn
		stream        pdpb.PD_TsoClient
	)
	for i := 0; i < maxRetryTimes; i++ {
		cc, url = c.getAllocatorClientConnByDCLocation(dc)
		cctx, cancel := context.WithCancel(dispatcherCtx)
		stream, err = c.createTsoStream(cctx, cancel, pdpb.NewPDClient(cc))
		failpoint.Inject("unreachableNetwork", func() {
			stream = nil
			err = status.New(codes.Unavailable, "unavailable").Err()
		})
		if stream != nil && err == nil {
			return connectionContext{stream, cancel, nil}, nil
		}
		if err != nil && c.enableForwarding {
			if rpcErr, ok := status.FromError(err); ok && (isNetworkError(rpcErr.Code()) || rpcErr.Code() == codes.Canceled) {
				networkErrNum++
			}
		}
		cancel()
		select {
		case <-dispatcherCtx.Done():
			return connectionContext{}, err
		case <-time.After(retryInterval):
		}
	}
	if networkErrNum == maxRetryTimes {
		followerClient, addr := c.followerClient()
		if followerClient != nil {
			log.Info("fall back to use follower to forward tso stream", zap.String("dc", dc), zap.String("addr", addr))
			forwardedHost, ok := c.getAllocatorLeaderAddrByDCLocation(dc)
			if !ok {
				return connectionContext{}, errors.Errorf("cannot find the allocator leader in %s", dc)
			}
			cctx, cancel := context.WithCancel(dispatcherCtx)
			cctx = grpcutil.BuildForwardContext(cctx, forwardedHost)
			stream, err = c.createTsoStream(cctx, cancel, followerClient)
			if err == nil {
				streamCh := make(streamCh)
				forwardedHostTrim := trimHTTPPrefix(forwardedHost)
				addrTrim := trimHTTPPrefix(addr)
				go c.checkAllocator(dispatcherCtx, cancel, dc, forwardedHostTrim, addrTrim, url, streamCh)
				requestForwarded.WithLabelValues(forwardedHostTrim, addrTrim).Set(1)
				return connectionContext{stream, cancel, streamCh}, nil
			}
			cancel()
			return connectionContext{}, err
		}
	}
	return connectionContext{}, err
}
tryConnect will try to connect to the TSO allocator leader. If the connection becomes unreachable
and enableForwarding is true, it will create a new connection to a follower to do the forwarding,
while a new daemon will be created also to switch back to a normal leader connection ASAP the
connection comes back to normal.
package pd
import (
	"context"
	"fmt"
	"math/rand"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"github.com/opentracing/opentracing-go"
	"github.com/pingcap/errors"
	"github.com/pingcap/failpoint"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/log"
	"github.com/tikv/pd/pkg/errs"
	"github.com/tikv/pd/pkg/grpcutil"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/status"
)
type Region struct {
	Meta         *metapb.Region
	Leader       *metapb.Peer
	DownPeers    []*metapb.Peer
	PendingPeers []*metapb.Peer
}
type Client interface {
	GetClusterID(ctx context.Context) uint64
	GetAllMembers(ctx context.Context) ([]*pdpb.Member, error)
	GetLeaderAddr() string
	GetTS(ctx context.Context) (int64, int64, error)
	GetTSAsync(ctx context.Context) TSFuture
	GetLocalTS(ctx context.Context, dcLocation string) (int64, int64, error)
	GetLocalTSAsync(ctx context.Context, dcLocation string) TSFuture
	GetRegion(ctx context.Context, key []byte) (*Region, error)
	GetRegionFromMember(ctx context.Context, key []byte, memberURLs []string) (*Region, error)
	GetPrevRegion(ctx context.Context, key []byte) (*Region, error)
	GetRegionByID(ctx context.Context, regionID uint64) (*Region, error)
	ScanRegions(ctx context.Context, key, endKey []byte, limit int) ([]*Region, error)
	GetStore(ctx context.Context, storeID uint64) (*metapb.Store, error)
	GetAllStores(ctx context.Context, opts ...GetStoreOption) ([]*metapb.Store, error)
	UpdateGCSafePoint(ctx context.Context, safePoint uint64) (uint64, error)
	UpdateServiceGCSafePoint(ctx context.Context, serviceID string, ttl int64, safePoint uint64) (uint64, error)
	ScatterRegion(ctx context.Context, regionID uint64) error
	ScatterRegions(ctx context.Context, regionsID []uint64, opts ...RegionsOption) (*pdpb.ScatterRegionResponse, error)
	SplitRegions(ctx context.Context, splitKeys [][]byte, opts ...RegionsOption) (*pdpb.SplitRegionsResponse, error)
	GetOperator(ctx context.Context, regionID uint64) (*pdpb.GetOperatorResponse, error)
	Close()
}
type GetStoreOp struct {
	excludeTombstone bool
}
type GetStoreOption func(*GetStoreOp)
func WithExcludeTombstone() GetStoreOption {
	return func(op *GetStoreOp) { op.excludeTombstone = true }
}
type RegionsOp struct {
	group      string
	retryLimit uint64
}
type RegionsOption func(op *RegionsOp)
func WithGroup(group string) RegionsOption {
	return func(op *RegionsOp) { op.group = group }
}
func WithRetry(retry uint64) RegionsOption {
	return func(op *RegionsOp) { op.retryLimit = retry }
}
type tsoRequest struct {
	start      time.Time
	clientCtx  context.Context
	requestCtx context.Context
	done       chan error
	physical   int64
	logical    int64
	dcLocation string
}
type tsoBatchController struct {
	maxBatchSize         int
	maxBatchWaitInterval time.Duration
	bestBatchSize int
	tsoRequestCh          chan *tsoRequest
	collectedRequests     []*tsoRequest
	collectedRequestCount int
	batchStartTime time.Time
}
func newTSOBatchController(tsoRequestCh chan *tsoRequest, maxBatchSize int, maxBatchWaitInterval time.Duration) *tsoBatchController {
	return &tsoBatchController{
		maxBatchSize:          maxBatchSize,
		maxBatchWaitInterval:  maxBatchWaitInterval,
		bestBatchSize:         8,  
		tsoRequestCh:          tsoRequestCh,
		collectedRequests:     make([]*tsoRequest, maxBatchSize+1),
		collectedRequestCount: 0,
	}
}
func (tbc *tsoBatchController) fetchPendingRequests(ctx context.Context) error {
	var firstTSORequest *tsoRequest
	select {
	case <-ctx.Done():
		return ctx.Err()
	case firstTSORequest = <-tbc.tsoRequestCh:
	}
	tbc.batchStartTime = time.Now()
	tbc.collectedRequestCount = 0
	tbc.pushRequest(firstTSORequest)
fetchPendingRequestsLoop:
	for tbc.collectedRequestCount < tbc.maxBatchSize {
		select {
		case tsoReq := <-tbc.tsoRequestCh:
			tbc.pushRequest(tsoReq)
		case <-ctx.Done():
			return ctx.Err()
		default:
			break fetchPendingRequestsLoop
		}
	}
	if tbc.collectedRequestCount >= tbc.maxBatchSize || tbc.maxBatchWaitInterval <= 0 {
		return nil
	}
	if tbc.collectedRequestCount < tbc.bestBatchSize {
		after := time.NewTimer(tbc.maxBatchWaitInterval)
		defer after.Stop()
		for tbc.collectedRequestCount < tbc.bestBatchSize {
			select {
			case tsoReq := <-tbc.tsoRequestCh:
				tbc.pushRequest(tsoReq)
			case <-ctx.Done():
				return ctx.Err()
			case <-after.C:
				return nil
			}
		}
	}
	for tbc.collectedRequestCount < tbc.maxBatchSize {
		select {
		case tsoReq := <-tbc.tsoRequestCh:
			tbc.pushRequest(tsoReq)
		case <-ctx.Done():
			return ctx.Err()
		default:
			return nil
		}
	}
	return nil
}
func (tbc *tsoBatchController) pushRequest(tsoReq *tsoRequest) {
	tbc.collectedRequests[tbc.collectedRequestCount] = tsoReq
	tbc.collectedRequestCount++
}
func (tbc *tsoBatchController) adjustBestBatchSize() {
	tsoBestBatchSize.Observe(float64(tbc.bestBatchSize))
	length := tbc.collectedRequestCount
	if length < tbc.bestBatchSize && tbc.bestBatchSize > 1 {
		tbc.bestBatchSize--
	} else if length > tbc.bestBatchSize+4   &&
		tbc.bestBatchSize < tbc.maxBatchSize {
		tbc.bestBatchSize++
	}
}
type tsoDispatcher struct {
	dispatcherCtx      context.Context
	dispatcherCancel   context.CancelFunc
	tsoBatchController *tsoBatchController
}
type lastTSO struct {
	physical int64
	logical  int64
}
const (
	defaultPDTimeout       = 3 * time.Second
	dialTimeout            = 3 * time.Second
	updateMemberTimeout    = time.Second  
	tsLoopDCCheckInterval  = time.Minute
	defaultMaxTSOBatchSize = 10000  
	maxInitClusterRetries  = 100
	retryInterval          = 1 * time.Second
	maxRetryTimes          = 5
)
var LeaderHealthCheckInterval = time.Second
var (
	errFailInitClusterID = errors.New("[pd] failed to get cluster id")
	errClosing = errors.New("[pd] closing")
	errTSOLength = errors.New("[pd] tso length in rpc response is incorrect")
)
type ClientOption func(c *client)
func WithGRPCDialOptions(opts ...grpc.DialOption) ClientOption {
	return func(c *client) {
		c.gRPCDialOptions = append(c.gRPCDialOptions, opts...)
	}
}
func WithCustomTimeoutOption(timeout time.Duration) ClientOption {
	return func(c *client) {
		c.timeout = timeout
	}
}
func WithForwardingOption(enableForwarding bool) ClientOption {
	return func(c *client) {
		c.enableForwarding = enableForwarding
	}
}
func WithMaxErrorRetry(count int) ClientOption {
	return func(c *client) {
		c.maxRetryTimes = count
	}
}
func WithTSOFollowerProxy(enableFollowerTSOBacth bool) ClientOption {
	return func(c *client) {
		c.enableTSOFollowerProxy = enableFollowerTSOBacth
	}
}
func WithMaxTSOBatchWaitInterval(maxTSOBatchWaitInterval time.Duration) ClientOption {
	return func(c *client) {
		c.maxTSOBatchWaitInterval = maxTSOBatchWaitInterval
	}
}
type client struct {
	*baseClient
	tsoDispatcher sync.Map  
	tsDeadline sync.Map  
	lastTSMap sync.Map  
	checkTSDeadlineCh    chan struct{}
	leaderNetworkFailure int32
	enableForwarding       bool
	enableTSOFollowerProxy bool
	maxTSOBatchWaitInterval time.Duration
}
func NewClient(pdAddrs []string, security SecurityOption, opts ...ClientOption) (Client, error) {
	return NewClientWithContext(context.Background(), pdAddrs, security, opts...)
}
func NewClientWithContext(ctx context.Context, pdAddrs []string, security SecurityOption, opts ...ClientOption) (Client, error) {
	log.Info("[pd] create pd client with endpoints", zap.Strings("pd-address", pdAddrs))
	c := &client{
		baseClient:        newBaseClient(ctx, addrsToUrls(pdAddrs), security),
		checkTSDeadlineCh: make(chan struct{}),
	}
	for _, opt := range opts {
		opt(c)
	}
	if err := c.init(); err != nil {
		return nil, err
	}
	c.updateTSODispatcher()
	c.wg.Add(3)
	go c.tsLoop()
	go c.tsCancelLoop()
	go c.leaderCheckLoop()
	return c, nil
}
func (c *client) updateTSODispatcher() {
	c.allocators.Range(func(dcLocationKey, _ interface{}) bool {
		dcLocation := dcLocationKey.(string)
		if !c.checkTSODispatcher(dcLocation) {
			c.createTSODispatcher(dcLocation)
		}
		return true
	})
	c.tsoDispatcher.Range(func(dcLocationKey, _ interface{}) bool {
		dcLocation := dcLocationKey.(string)
		if dcLocation == globalDCLocation {
			return true
		}
		if dispatcher, exist := c.allocators.Load(dcLocation); !exist {
			log.Info("[pd] delete unused tso dispatcher", zap.String("dc-location", dcLocation))
			dispatcher.(*tsoDispatcher).dispatcherCancel()
			c.tsoDispatcher.Delete(dcLocation)
		}
		return true
	})
}
func (c *client) leaderCheckLoop() {
	defer c.wg.Done()
	leaderCheckLoopCtx, leaderCheckLoopCancel := context.WithCancel(c.ctx)
	defer leaderCheckLoopCancel()
	ticker := time.NewTicker(LeaderHealthCheckInterval)
	defer ticker.Stop()
	for {
		select {
		case <-c.ctx.Done():
			return
		case <-ticker.C:
			c.checkLeaderHealth(leaderCheckLoopCtx)
		}
	}
}
func (c *client) checkLeaderHealth(ctx context.Context) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()
	if cc, ok := c.clientConns.Load(c.GetLeaderAddr()); ok {
		healthCli := healthpb.NewHealthClient(cc.(*grpc.ClientConn))
		resp, err := healthCli.Check(ctx, &healthpb.HealthCheckRequest{Service: ""})
		rpcErr, ok := status.FromError(err)
		failpoint.Inject("unreachableNetwork1", func() {
			resp = nil
			err = status.New(codes.Unavailable, "unavailable").Err()
		})
		if (ok && isNetworkError(rpcErr.Code())) || resp.GetStatus() != healthpb.HealthCheckResponse_SERVING {
			atomic.StoreInt32(&(c.leaderNetworkFailure), int32(1))
		} else {
			atomic.StoreInt32(&(c.leaderNetworkFailure), int32(0))
		}
	}
}
type deadline struct {
	timer  <-chan time.Time
	done   chan struct{}
	cancel context.CancelFunc
}
func (c *client) tsCancelLoop() {
	defer c.wg.Done()
	tsCancelLoopCtx, tsCancelLoopCancel := context.WithCancel(c.ctx)
	defer tsCancelLoopCancel()
	ticker := time.NewTicker(tsLoopDCCheckInterval)
	defer ticker.Stop()
	for {
		c.allocators.Range(func(dcLocation, _ interface{}) bool {
			c.watchTSDeadline(tsCancelLoopCtx, dcLocation.(string))
			return true
		})
		select {
		case <-c.checkTSDeadlineCh:
			continue
		case <-ticker.C:
			continue
		case <-tsCancelLoopCtx.Done():
			return
		}
	}
}
func (c *client) watchTSDeadline(ctx context.Context, dcLocation string) {
	if _, exist := c.tsDeadline.Load(dcLocation); !exist {
		tsDeadlineCh := make(chan deadline, 1)
		c.tsDeadline.Store(dcLocation, tsDeadlineCh)
		go func(dc string, tsDeadlineCh <-chan deadline) {
			for {
				select {
				case d := <-tsDeadlineCh:
					select {
					case <-d.timer:
						log.Error("tso request is canceled due to timeout", zap.String("dc-location", dc), errs.ZapError(errs.ErrClientGetTSOTimeout))
						d.cancel()
					case <-d.done:
						continue
					case <-ctx.Done():
						return
					}
				case <-ctx.Done():
					return
				}
			}
		}(dcLocation, tsDeadlineCh)
	}
}
func (c *client) scheduleCheckTSDeadline() {
	select {
	case c.checkTSDeadlineCh <- struct{}{}:
	default:
	}
}
func (c *client) checkStreamTimeout(streamCtx context.Context, cancel context.CancelFunc, done chan struct{}) {
	select {
	case <-done:
		return
	case <-time.After(c.timeout):
		cancel()
	case <-streamCtx.Done():
	}
	<-done
}
func (c *client) GetAllMembers(ctx context.Context) ([]*pdpb.Member, error) {
	start := time.Now()
	defer func() { cmdDurationGetAllMembers.Observe(time.Since(start).Seconds()) }()
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	req := &pdpb.GetMembersRequest{Header: c.requestHeader()}
	ctx = grpcutil.BuildForwardContext(ctx, c.GetLeaderAddr())
	resp, err := c.getClient().GetMembers(ctx, req)
	cancel()
	if err != nil {
		cmdFailDurationGetAllMembers.Observe(time.Since(start).Seconds())
		c.ScheduleCheckLeader()
		return nil, errors.WithStack(err)
	}
	return resp.GetMembers(), nil
}
func (c *client) tsLoop() {
	defer c.wg.Done()
	loopCtx, loopCancel := context.WithCancel(c.ctx)
	defer loopCancel()
	ticker := time.NewTicker(tsLoopDCCheckInterval)
	defer ticker.Stop()
	for {
		c.updateTSODispatcher()
		select {
		case <-ticker.C:
		case <-c.checkTSODispatcherCh:
		case <-loopCtx.Done():
			return
		}
	}
}
func (c *client) createTsoStream(ctx context.Context, cancel context.CancelFunc, client pdpb.PDClient) (pdpb.PD_TsoClient, error) {
	done := make(chan struct{})
	go c.checkStreamTimeout(ctx, cancel, done)
	stream, err := client.Tso(ctx)
	done <- struct{}{}
	return stream, err
}
func (c *client) checkAllocator(dispatcherCtx context.Context, forwardCancel context.CancelFunc, dc, forwardedHostTrim, addrTrim, url string, streamCh streamCh) {
	defer func() {
		forwardCancel()
		close(streamCh)
		requestForwarded.WithLabelValues(forwardedHostTrim, addrTrim).Set(0)
	}()
	cc, u := c.getAllocatorClientConnByDCLocation(dc)
	healthCli := healthpb.NewHealthClient(cc)
	for {
		if u != url {
			log.Info("the leader of the allocator leader is changed", zap.String("dc", dc), zap.String("origin", url), zap.String("new", u))
			return
		}
		healthCtx, healthCancel := context.WithTimeout(dispatcherCtx, c.timeout)
		resp, err := healthCli.Check(healthCtx, &healthpb.HealthCheckRequest{Service: ""})
		failpoint.Inject("unreachableNetwork", func() {
			resp.Status = healthpb.HealthCheckResponse_UNKNOWN
		})
		healthCancel()
		if err == nil && resp.GetStatus() == healthpb.HealthCheckResponse_SERVING {
			cctx, cancel := context.WithCancel(dispatcherCtx)
			stream, err := c.createTsoStream(cctx, cancel, pdpb.NewPDClient(cc))
			if err == nil && stream != nil {
				log.Info("recover the original tso stream since the network has become normal", zap.String("dc", dc), zap.String("url", url))
				streamCh <- struct {
					cancel context.CancelFunc
					stream pdpb.PD_TsoClient
				}{
					cancel: cancel,
					stream: stream,
				}
				return
			}
		}
		select {
		case <-dispatcherCtx.Done():
			return
		case <-time.After(time.Second):
			_, u = c.getAllocatorClientConnByDCLocation(dc)
		}
	}
}
func (c *client) checkTSODispatcher(dcLocation string) bool {
	dispatcher, ok := c.tsoDispatcher.Load(dcLocation)
	if !ok || dispatcher == nil {
		return false
	}
	return true
}
func (c *client) createTSODispatcher(dcLocation string) {
	dispatcherCtx, dispatcherCancel := context.WithCancel(c.ctx)
	dispatcher := &tsoDispatcher{
		dispatcherCtx:      dispatcherCtx,
		dispatcherCancel:   dispatcherCancel,
		tsoBatchController: newTSOBatchController(make(chan *tsoRequest, defaultMaxTSOBatchSize*2), defaultMaxTSOBatchSize, c.maxTSOBatchWaitInterval),
	}
	go c.handleDispatcher(dispatcherCtx, dcLocation, dispatcher.tsoBatchController)
	c.tsoDispatcher.Store(dcLocation, dispatcher)
	log.Info("[pd] tso dispatcher created", zap.String("dc-location", dcLocation))
}
type streamCh chan struct {
	cancel context.CancelFunc
	stream pdpb.PD_TsoClient
}
func (c *client) handleDispatcher(dispatcherCtx context.Context, dc string, tbc *tsoBatchController) {
	var (
		err           error
		cancel        context.CancelFunc
		stream        pdpb.PD_TsoClient
		opts          []opentracing.StartSpanOption
		needUpdate    = false
		streamCh      streamCh
		connectionCtx connectionContext
	)
	defer func() {
		log.Info("[pd] exit tso dispatcher", zap.String("dc-location", dc))
		if cancel != nil {
			cancel()
		}
	}()
	createTSOConnection := c.tryConnect
	if c.enableTSOFollowerProxy && dc == globalDCLocation   {
		createTSOConnection = c.tryConnectToFollowerProxy
	}
	for {
		if stream == nil {
			if needUpdate {
				err = c.updateMember()
				if err != nil {
					select {
					case <-dispatcherCtx.Done():
						return
					default:
					}
					continue
				}
				needUpdate = false
			}
			connectionCtx, err = createTSOConnection(dispatcherCtx, dc)
			stream, cancel, streamCh = connectionCtx.stream, connectionCtx.cancel, connectionCtx.streamCh
			if err != nil || stream == nil {
				select {
				case <-dispatcherCtx.Done():
					return
				default:
				}
				log.Error("[pd] create tso stream error", zap.String("dc-location", dc), errs.ZapError(errs.ErrClientCreateTSOStream, err))
				c.ScheduleCheckLeader()
				c.revokeTSORequest(errors.WithStack(err), tbc.tsoRequestCh)
				select {
				case <-time.After(time.Second):
				case <-dispatcherCtx.Done():
					return
				}
				continue
			}
		}
		if err = tbc.fetchPendingRequests(dispatcherCtx); err != nil {
			log.Error("[pd] fetch pending tso requests error", zap.String("dc-location", dc), errs.ZapError(errs.ErrClientGetTSO, err))
			return
		}
		if tbc.maxBatchWaitInterval >= 0 {
			tbc.adjustBestBatchSize()
		}
		done := make(chan struct{})
		dl := deadline{
			timer:  time.After(c.timeout),
			done:   done,
			cancel: cancel,
		}
		tsDeadlineCh, ok := c.tsDeadline.Load(dc)
		for !ok || tsDeadlineCh == nil {
			c.scheduleCheckTSDeadline()
			time.Sleep(time.Millisecond * 100)
			tsDeadlineCh, ok = c.tsDeadline.Load(dc)
		}
		select {
		case tsDeadlineCh.(chan deadline) <- dl:
		case <-dispatcherCtx.Done():
			return
		}
		opts = extractSpanReference(tbc, opts[:0])
		select {
		case newStream, ok := <-streamCh:
			if ok {
				stream = newStream.stream
				cancel = newStream.cancel
				log.Info("tso stream has changed back to pd or tso allocator leader")
			}
		default:
		}
		err = c.processTSORequests(stream, dc, tbc, opts)
		close(done)
		if err != nil {
			select {
			case <-dispatcherCtx.Done():
				return
			default:
			}
			log.Error("[pd] getTS error", zap.String("dc-location", dc), errs.ZapError(errs.ErrClientGetTSO, err))
			c.ScheduleCheckLeader()
			cancel()
			stream = nil
			if IsLeaderChange(err) {
				needUpdate = true
			}
		}
	}
}
type connectionContext struct {
	stream pdpb.PD_TsoClient
	cancel context.CancelFunc
	streamCh streamCh
}