From 2fc38e0018af31d4774917d97db5a03d5602ba02 Mon Sep 17 00:00:00 2001
From: Bin Shi <binshi.bing@gmail.com>
Date: Thu, 16 Mar 2023 17:46:40 -0700
Subject: [PATCH] Add general tso forward/dispatcher for independent
 pd(tso)/tso services and cross cluster forwarding.

Signed-off-by: Bin Shi <binshi.bing@gmail.com>

more change

Signed-off-by: Bin Shi <binshi.bing@gmail.com>
---
 pkg/errs/errno.go                        |   1 +
 pkg/mcs/tso/server/grpc_service.go       | 197 +---------------
 pkg/mcs/tso/server/server.go             |  16 +-
 pkg/utils/tsoutil/tso_dispatcher.go      | 207 +++++++++++++++++
 pkg/utils/tsoutil/tso_proto_factory.go   | 109 +++++++++
 pkg/utils/tsoutil/tso_request.go         | 160 +++++++++++++
 pkg/utils/tsoutil/{tso.go => tsoutil.go} |   0
 server/grpc_service.go                   | 279 ++---------------------
 server/server.go                         |  11 +-
 9 files changed, 528 insertions(+), 452 deletions(-)
 create mode 100644 pkg/utils/tsoutil/tso_dispatcher.go
 create mode 100644 pkg/utils/tsoutil/tso_proto_factory.go
 create mode 100644 pkg/utils/tsoutil/tso_request.go
 rename pkg/utils/tsoutil/{tso.go => tsoutil.go} (100%)

diff --git a/pkg/errs/errno.go b/pkg/errs/errno.go
index 155d7ef45a4b..706de8764d70 100644
--- a/pkg/errs/errno.go
+++ b/pkg/errs/errno.go
@@ -39,6 +39,7 @@ var (
 	ErrGenerateTimestamp  = errors.Normalize("generate timestamp failed, %s", errors.RFCCodeText("PD:tso:ErrGenerateTimestamp"))
 	ErrLogicOverflow      = errors.Normalize("logic part overflow", errors.RFCCodeText("PD:tso:ErrLogicOverflow"))
 	ErrProxyTSOTimeout    = errors.Normalize("proxy tso timeout", errors.RFCCodeText("PD:tso:ErrProxyTSOTimeout"))
+	ErrTSOTypeConversion  = errors.Normalize("type conversion error", errors.RFCCodeText("PD:tso:ErrTypeConversion"))
 )
 
 // member errors
diff --git a/pkg/mcs/tso/server/grpc_service.go b/pkg/mcs/tso/server/grpc_service.go
index c650c4910ad3..31f7abb50c41 100644
--- a/pkg/mcs/tso/server/grpc_service.go
+++ b/pkg/mcs/tso/server/grpc_service.go
@@ -20,17 +20,14 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/kvproto/pkg/tsopb"
 	"github.com/pingcap/log"
 	"github.com/pkg/errors"
 	bs "github.com/tikv/pd/pkg/basicserver"
-	"github.com/tikv/pd/pkg/errs"
 	"github.com/tikv/pd/pkg/mcs/registry"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/pkg/utils/grpcutil"
-	"github.com/tikv/pd/pkg/utils/logutil"
-	"go.uber.org/zap"
+	"github.com/tikv/pd/pkg/utils/tsoutil"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -116,16 +113,20 @@ func (s *Service) Tso(stream tsopb.TSO_TsoServer) error {
 		streamCtx := stream.Context()
 		forwardedHost := grpcutil.GetForwardedHost(streamCtx)
 		if !s.IsLocalRequest(forwardedHost) {
+			clientConn, err := s.GetDelegateClient(s.ctx, forwardedHost)
+			if err != nil {
+				return err
+			}
+
 			if errCh == nil {
 				doneCh = make(chan struct{})
 				defer close(doneCh)
 				errCh = make(chan error)
 			}
-			s.dispatchTSORequest(ctx, &tsoRequest{
-				forwardedHost,
-				request,
-				stream,
-			}, forwardedHost, doneCh, errCh)
+
+			tsoProtoFactory := s.tsoProtoTSOFactory
+			tsoRequest := tsoutil.NewTSOProtoTSORequest(forwardedHost, clientConn, request, stream)
+			s.tsoDispatcher.DispatchRequest(ctx, tsoRequest, tsoProtoFactory, doneCh, errCh)
 			continue
 		}
 
@@ -174,181 +175,3 @@ func (s *Service) errorHeader(err *tsopb.Error) *tsopb.ResponseHeader {
 		Error:     err,
 	}
 }
-
-type tsoRequest struct {
-	forwardedHost string
-	request       *tsopb.TsoRequest
-	stream        tsopb.TSO_TsoServer
-}
-
-func (s *Service) dispatchTSORequest(ctx context.Context, request *tsoRequest, forwardedHost string, doneCh <-chan struct{}, errCh chan<- error) {
-	tsoRequestChInterface, loaded := s.tsoDispatcher.LoadOrStore(forwardedHost, make(chan *tsoRequest, maxMergeTSORequests))
-	if !loaded {
-		tsDeadlineCh := make(chan deadline, 1)
-		go s.handleDispatcher(ctx, forwardedHost, tsoRequestChInterface.(chan *tsoRequest), tsDeadlineCh, doneCh, errCh)
-		go watchTSDeadline(ctx, tsDeadlineCh)
-	}
-	tsoRequestChInterface.(chan *tsoRequest) <- request
-}
-
-func (s *Service) handleDispatcher(ctx context.Context, forwardedHost string, tsoRequestCh <-chan *tsoRequest, tsDeadlineCh chan<- deadline, doneCh <-chan struct{}, errCh chan<- error) {
-	defer logutil.LogPanic()
-	dispatcherCtx, ctxCancel := context.WithCancel(ctx)
-	defer ctxCancel()
-	defer s.tsoDispatcher.Delete(forwardedHost)
-
-	var (
-		forwardStream tsopb.TSO_TsoClient
-		cancel        context.CancelFunc
-	)
-	client, err := s.GetDelegateClient(ctx, forwardedHost)
-	if err != nil {
-		goto errHandling
-	}
-	log.Info("create tso forward stream", zap.String("forwarded-host", forwardedHost))
-	forwardStream, cancel, err = s.CreateTsoForwardStream(client)
-errHandling:
-	if err != nil || forwardStream == nil {
-		log.Error("create tso forwarding stream error", zap.String("forwarded-host", forwardedHost), errs.ZapError(errs.ErrGRPCCreateStream, err))
-		select {
-		case <-dispatcherCtx.Done():
-			return
-		case _, ok := <-doneCh:
-			if !ok {
-				return
-			}
-		case errCh <- err:
-			close(errCh)
-			return
-		}
-	}
-	defer cancel()
-
-	requests := make([]*tsoRequest, maxMergeTSORequests+1)
-	for {
-		select {
-		case first := <-tsoRequestCh:
-			pendingTSOReqCount := len(tsoRequestCh) + 1
-			requests[0] = first
-			for i := 1; i < pendingTSOReqCount; i++ {
-				requests[i] = <-tsoRequestCh
-			}
-			done := make(chan struct{})
-			dl := deadline{
-				timer:  time.After(defaultTSOProxyTimeout),
-				done:   done,
-				cancel: cancel,
-			}
-			select {
-			case tsDeadlineCh <- dl:
-			case <-dispatcherCtx.Done():
-				return
-			}
-			err = s.processTSORequests(forwardStream, requests[:pendingTSOReqCount])
-			close(done)
-			if err != nil {
-				log.Error("proxy forward tso error", zap.String("forwarded-host", forwardedHost), errs.ZapError(errs.ErrGRPCSend, err))
-				select {
-				case <-dispatcherCtx.Done():
-					return
-				case _, ok := <-doneCh:
-					if !ok {
-						return
-					}
-				case errCh <- err:
-					close(errCh)
-					return
-				}
-			}
-		case <-dispatcherCtx.Done():
-			return
-		}
-	}
-}
-
-func (s *Service) processTSORequests(forwardStream tsopb.TSO_TsoClient, requests []*tsoRequest) error {
-	start := time.Now()
-	// Merge the requests
-	count := uint32(0)
-	for _, request := range requests {
-		count += request.request.GetCount()
-	}
-	req := &tsopb.TsoRequest{
-		Header: requests[0].request.GetHeader(),
-		Count:  count,
-		// TODO: support Local TSO proxy forwarding.
-		DcLocation: requests[0].request.GetDcLocation(),
-	}
-	// Send to the leader stream.
-	if err := forwardStream.Send(req); err != nil {
-		return err
-	}
-	resp, err := forwardStream.Recv()
-	if err != nil {
-		return err
-	}
-	tsoProxyHandleDuration.Observe(time.Since(start).Seconds())
-	tsoProxyBatchSize.Observe(float64(count))
-	// Split the response
-	physical, logical, suffixBits := resp.GetTimestamp().GetPhysical(), resp.GetTimestamp().GetLogical(), resp.GetTimestamp().GetSuffixBits()
-	// `logical` is the largest ts's logical part here, we need to do the subtracting before we finish each TSO request.
-	// This is different from the logic of client batch, for example, if we have a largest ts whose logical part is 10,
-	// count is 5, then the splitting results should be 5 and 10.
-	firstLogical := addLogical(logical, -int64(count), suffixBits)
-	return s.finishTSORequest(requests, physical, firstLogical, suffixBits)
-}
-
-// Because of the suffix, we need to shift the count before we add it to the logical part.
-func addLogical(logical, count int64, suffixBits uint32) int64 {
-	return logical + count<<suffixBits
-}
-
-func (s *Service) finishTSORequest(requests []*tsoRequest, physical, firstLogical int64, suffixBits uint32) error {
-	countSum := int64(0)
-	for i := 0; i < len(requests); i++ {
-		count := requests[i].request.GetCount()
-		countSum += int64(count)
-		response := &tsopb.TsoResponse{
-			Header: s.header(),
-			Count:  count,
-			Timestamp: &pdpb.Timestamp{
-				Physical:   physical,
-				Logical:    addLogical(firstLogical, countSum, suffixBits),
-				SuffixBits: suffixBits,
-			},
-		}
-		// Send back to the client.
-		if err := requests[i].stream.Send(response); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-type deadline struct {
-	timer  <-chan time.Time
-	done   chan struct{}
-	cancel context.CancelFunc
-}
-
-func watchTSDeadline(ctx context.Context, tsDeadlineCh <-chan deadline) {
-	defer logutil.LogPanic()
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-	for {
-		select {
-		case d := <-tsDeadlineCh:
-			select {
-			case <-d.timer:
-				log.Error("tso proxy request processing is canceled due to timeout", errs.ZapError(errs.ErrProxyTSOTimeout))
-				d.cancel()
-			case <-d.done:
-				continue
-			case <-ctx.Done():
-				return
-			}
-		case <-ctx.Done():
-			return
-		}
-	}
-}
diff --git a/pkg/mcs/tso/server/server.go b/pkg/mcs/tso/server/server.go
index a6aed144d37d..fae4f7f56111 100644
--- a/pkg/mcs/tso/server/server.go
+++ b/pkg/mcs/tso/server/server.go
@@ -53,6 +53,7 @@ import (
 	"github.com/tikv/pd/pkg/utils/logutil"
 	"github.com/tikv/pd/pkg/utils/memberutil"
 	"github.com/tikv/pd/pkg/utils/metricutil"
+	"github.com/tikv/pd/pkg/utils/tsoutil"
 	"github.com/tikv/pd/pkg/versioninfo"
 	"go.etcd.io/etcd/clientv3"
 	"go.etcd.io/etcd/pkg/types"
@@ -105,8 +106,12 @@ type Server struct {
 	tsoAllocatorManager *tso.AllocatorManager
 	// Store as map[string]*grpc.ClientConn
 	clientConns sync.Map
-	// Store as map[string]chan *tsoRequest
-	tsoDispatcher sync.Map
+	// tsoDispatcher is used to dispatch different TSO requests to
+	// the corresponding forwarding TSO channel.
+	tsoDispatcher *tsoutil.TSODispatcher
+	// tsoProtoTSOFactory is the abstract factory for creating tso
+	// related data structures defined in pd protocol
+	tsoProtoTSOFactory *tsoutil.TSOProtoTSOFactory
 
 	// Callback functions for different stages
 	// startCallbacks will be called after the server is started.
@@ -354,11 +359,6 @@ func (s *Server) GetTSOAllocatorManager() *tso.AllocatorManager {
 	return s.tsoAllocatorManager
 }
 
-// GetTSODispatcher gets the TSO Dispatcher
-func (s *Server) GetTSODispatcher() *sync.Map {
-	return &s.tsoDispatcher
-}
-
 // IsLocalRequest checks if the forwarded host is the current host
 func (s *Server) IsLocalRequest(forwardedHost string) bool {
 	// TODO: Check if the forwarded host is the current host.
@@ -578,6 +578,8 @@ func (s *Server) startServer() (err error) {
 		s.cfg.GetTLSConfig(), func() time.Duration { return s.cfg.MaxResetTSGap.Duration })
 	// Set up the Global TSO Allocator here, it will be initialized once this TSO participant campaigns leader successfully.
 	s.tsoAllocatorManager.SetUpAllocator(s.ctx, tso.GlobalDCLocation, s.participant.GetLeadership())
+	s.tsoDispatcher = tsoutil.NewTSODispatcher(tsoProxyHandleDuration, tsoProxyBatchSize)
+	s.tsoProtoTSOFactory = &tsoutil.TSOProtoTSOFactory{}
 
 	s.service = &Service{Server: s}
 
diff --git a/pkg/utils/tsoutil/tso_dispatcher.go b/pkg/utils/tsoutil/tso_dispatcher.go
new file mode 100644
index 000000000000..fedff1e0b18f
--- /dev/null
+++ b/pkg/utils/tsoutil/tso_dispatcher.go
@@ -0,0 +1,207 @@
+// Copyright 2023 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tsoutil
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/pingcap/kvproto/pkg/pdpb"
+	"github.com/pingcap/log"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/tikv/pd/pkg/errs"
+	"go.uber.org/zap"
+	"google.golang.org/grpc"
+)
+
+const (
+	maxMergeTSORequests = 10000
+	// DefaultTSOProxyTimeout defines the default timeout value of TSP Proxying
+	DefaultTSOProxyTimeout = 3 * time.Second
+)
+
+type tsoResp interface {
+	GetTimestamp() *pdpb.Timestamp
+}
+
+type TSODispatcher struct {
+	tsoProxyHandleDuration prometheus.Histogram
+	tsoProxyBatchSize      prometheus.Histogram
+
+	// dispatchChs is used to dispatch different TSO requests to the corresponding forwarding TSO channel.
+	dispatchChs sync.Map // Store as map[string]chan TSORequest
+}
+
+func NewTSODispatcher(tsoProxyHandleDuration, tsoProxyBatchSize prometheus.Histogram) *TSODispatcher {
+	tsoDispatcher := &TSODispatcher{
+		tsoProxyHandleDuration: tsoProxyHandleDuration,
+		tsoProxyBatchSize:      tsoProxyBatchSize,
+	}
+	return tsoDispatcher
+}
+
+func (s *TSODispatcher) DispatchRequest(ctx context.Context, req TSORequest, tsoProtoFactory TSOProtoFactory, doneCh <-chan struct{}, errCh chan<- error) {
+	tsoRequestChInterface, loaded := s.dispatchChs.LoadOrStore(req.getForwardedHost(), make(chan TSORequest, maxMergeTSORequests))
+	if !loaded {
+		tsDeadlineCh := make(chan deadline, 1)
+		go s.handleDispatcher(ctx, tsoProtoFactory, req.getForwardedHost(), req.getClientConn(), tsoRequestChInterface.(chan TSORequest), tsDeadlineCh, doneCh, errCh)
+		go watchTSDeadline(ctx, tsDeadlineCh)
+	}
+	tsoRequestChInterface.(chan TSORequest) <- req
+}
+
+func (s *TSODispatcher) handleDispatcher(ctx context.Context, tsoProtoFactory TSOProtoFactory, forwardedHost string, clientConn *grpc.ClientConn,
+	tsoRequestCh <-chan TSORequest, tsDeadlineCh chan<- deadline, doneCh <-chan struct{}, errCh chan<- error) {
+	dispatcherCtx, ctxCancel := context.WithCancel(ctx)
+	defer ctxCancel()
+	defer s.dispatchChs.Delete(forwardedHost)
+
+	log.Info("create tso forward stream", zap.String("forwarded-host", forwardedHost))
+	forwardStream, cancel, err := tsoProtoFactory.createTSOForwardStream(ctx, clientConn)
+	if err != nil || forwardStream == nil {
+		log.Error("create tso forwarding stream error", zap.String("forwarded-host", forwardedHost), errs.ZapError(errs.ErrGRPCCreateStream, err))
+		select {
+		case <-dispatcherCtx.Done():
+			return
+		case _, ok := <-doneCh:
+			if !ok {
+				return
+			}
+		case errCh <- err:
+			close(errCh)
+			return
+		}
+	}
+	defer cancel()
+
+	requests := make([]TSORequest, maxMergeTSORequests+1)
+	for {
+		select {
+		case first := <-tsoRequestCh:
+			pendingTSOReqCount := len(tsoRequestCh) + 1
+			requests[0] = first
+			for i := 1; i < pendingTSOReqCount; i++ {
+				requests[i] = <-tsoRequestCh
+			}
+			done := make(chan struct{})
+			dl := deadline{
+				timer:  time.After(DefaultTSOProxyTimeout),
+				done:   done,
+				cancel: cancel,
+			}
+			select {
+			case tsDeadlineCh <- dl:
+			case <-dispatcherCtx.Done():
+				return
+			}
+			err = s.processTSORequests(forwardStream, requests[:pendingTSOReqCount], tsoProtoFactory)
+			close(done)
+			if err != nil {
+				log.Error("proxy forward tso error", zap.String("forwarded-host", forwardedHost), errs.ZapError(errs.ErrGRPCSend, err))
+				select {
+				case <-dispatcherCtx.Done():
+					return
+				case _, ok := <-doneCh:
+					if !ok {
+						return
+					}
+				case errCh <- err:
+					close(errCh)
+					return
+				}
+			}
+		case <-dispatcherCtx.Done():
+			return
+		}
+	}
+}
+
+func (s *TSODispatcher) processTSORequests(forwardStream tsoStream, requests []TSORequest, tsoProtoFactory TSOProtoFactory) error {
+	start := time.Now()
+	// Merge the requests
+	count := uint32(0)
+	for _, request := range requests {
+		count += request.getCount()
+	}
+	resp, err := requests[0].process(forwardStream, count, tsoProtoFactory)
+	if err != nil {
+		return err
+	}
+	s.tsoProxyHandleDuration.Observe(time.Since(start).Seconds())
+	s.tsoProxyBatchSize.Observe(float64(count))
+	// Split the response
+	physical, logical, suffixBits := resp.GetTimestamp().GetPhysical(), resp.GetTimestamp().GetLogical(), resp.GetTimestamp().GetSuffixBits()
+	// `logical` is the largest ts's logical part here, we need to do the subtracting before we finish each TSO request.
+	// This is different from the logic of client batch, for example, if we have a largest ts whose logical part is 10,
+	// count is 5, then the splitting results should be 5 and 10.
+	firstLogical := addLogical(logical, -int64(count), suffixBits)
+	return s.finishTSORequest(requests, physical, firstLogical, suffixBits)
+}
+
+// Because of the suffix, we need to shift the count before we add it to the logical part.
+func addLogical(logical, count int64, suffixBits uint32) int64 {
+	return logical + count<<suffixBits
+}
+
+func (s *TSODispatcher) finishTSORequest(requests []TSORequest, physical, firstLogical int64, suffixBits uint32) error {
+	countSum := int64(0)
+	for i := 0; i < len(requests); i++ {
+		newCountSum, err := requests[i].postProcess(countSum, physical, firstLogical, suffixBits)
+		if err != nil {
+			return err
+		}
+		countSum = newCountSum
+	}
+	return nil
+}
+
+type deadline struct {
+	timer  <-chan time.Time
+	done   chan struct{}
+	cancel context.CancelFunc
+}
+
+func watchTSDeadline(ctx context.Context, tsDeadlineCh <-chan deadline) {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	for {
+		select {
+		case d := <-tsDeadlineCh:
+			select {
+			case <-d.timer:
+				log.Error("tso proxy request processing is canceled due to timeout", errs.ZapError(errs.ErrProxyTSOTimeout))
+				d.cancel()
+			case <-d.done:
+				continue
+			case <-ctx.Done():
+				return
+			}
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+func checkStream(streamCtx context.Context, cancel context.CancelFunc, done chan struct{}) {
+	select {
+	case <-done:
+		return
+	case <-time.After(3 * time.Second):
+		cancel()
+	case <-streamCtx.Done():
+	}
+	<-done
+}
diff --git a/pkg/utils/tsoutil/tso_proto_factory.go b/pkg/utils/tsoutil/tso_proto_factory.go
new file mode 100644
index 000000000000..b5d7408abbdc
--- /dev/null
+++ b/pkg/utils/tsoutil/tso_proto_factory.go
@@ -0,0 +1,109 @@
+// Copyright 2023 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tsoutil
+
+import (
+	"context"
+
+	"github.com/pingcap/kvproto/pkg/pdpb"
+	"github.com/pingcap/kvproto/pkg/tsopb"
+	"google.golang.org/grpc"
+)
+
+type TSOProtoFactory interface {
+	createTSOForwardStream(ctx context.Context, client *grpc.ClientConn) (tsoStream, context.CancelFunc, error)
+}
+
+type TSOProtoTSOFactory struct {
+}
+
+type TSOProtoPDFactory struct {
+}
+
+func (s *TSOProtoTSOFactory) createTSOForwardStream(ctx context.Context, clientConn *grpc.ClientConn) (tsoStream, context.CancelFunc, error) {
+	done := make(chan struct{})
+	cctx, cancel := context.WithCancel(ctx)
+	go checkStream(cctx, cancel, done)
+	forwardStream, err := tsopb.NewTSOClient(clientConn).Tso(cctx)
+	done <- struct{}{}
+	return &tsoTSOStream{forwardStream}, cancel, err
+}
+
+func (s *TSOProtoPDFactory) createTSOForwardStream(ctx context.Context, clientConn *grpc.ClientConn) (tsoStream, context.CancelFunc, error) {
+	done := make(chan struct{})
+	cctx, cancel := context.WithCancel(ctx)
+	go checkStream(cctx, cancel, done)
+	forwardStream, err := pdpb.NewPDClient(clientConn).Tso(cctx)
+	done <- struct{}{}
+	return &pdTSOStream{forwardStream}, cancel, err
+}
+
+type tsoStream interface {
+	// process sends a request and receives the response through the stream
+	process(clusterID uint64, count, keyspaceID, keyspaceGroupID uint32, dcLocation string) (tsoResponse, error)
+}
+
+type tsoTSOStream struct {
+	stream tsopb.TSO_TsoClient
+}
+
+// process sends a request and receives the response through the stream
+func (s *tsoTSOStream) process(clusterID uint64, count, keyspaceID, keyspaceGroupID uint32, dcLocation string) (tsoResponse, error) {
+	req := &tsopb.TsoRequest{
+		Header: &tsopb.RequestHeader{
+			ClusterId:       clusterID,
+			KeyspaceId:      keyspaceID,
+			KeyspaceGroupId: keyspaceGroupID,
+		},
+		Count: count,
+		// TODO: support Local TSO proxy forwarding.
+		DcLocation: dcLocation,
+	}
+
+	// Send to the tso server stream
+	if err := s.stream.Send(req); err != nil {
+		return nil, err
+	}
+	resp, err := s.stream.Recv()
+	if err != nil {
+		return nil, err
+	}
+	return resp, nil
+}
+
+type pdTSOStream struct {
+	stream pdpb.PD_TsoClient
+}
+
+// process sends a request and receives the response through the stream
+func (s *pdTSOStream) process(clusterID uint64, count, _, _ uint32, dcLocation string) (tsoResponse, error) {
+	req := &pdpb.TsoRequest{
+		Header: &pdpb.RequestHeader{
+			ClusterId: clusterID,
+		},
+		Count: count,
+		// TODO: support Local TSO proxy forwarding.
+		DcLocation: dcLocation,
+	}
+	// Send to the tso server stream
+	if err := s.stream.Send(req); err != nil {
+		return nil, err
+	}
+	resp, err := s.stream.Recv()
+	if err != nil {
+		return nil, err
+	}
+	return resp, nil
+}
diff --git a/pkg/utils/tsoutil/tso_request.go b/pkg/utils/tsoutil/tso_request.go
new file mode 100644
index 000000000000..6cbd31985191
--- /dev/null
+++ b/pkg/utils/tsoutil/tso_request.go
@@ -0,0 +1,160 @@
+// Copyright 2023 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tsoutil
+
+import (
+	"github.com/pingcap/kvproto/pkg/pdpb"
+	"github.com/pingcap/kvproto/pkg/tsopb"
+	"github.com/tikv/pd/pkg/mcs/utils"
+	"google.golang.org/grpc"
+)
+
+// TSORequest is an interface wrapping tsopb.TsoRequest and pdpb.TsoRequest so
+// they can be generally handled by the TSO dispatcher
+type TSORequest interface {
+	// getForwardedHost returns the forwarded host
+	getForwardedHost() string
+	// getClientConn returns the grpc client connection
+	getClientConn() *grpc.ClientConn
+	// getCount returns the count of timestamps to retrieve
+	getCount() uint32
+	// process sends request and receive response via stream.
+	// count defins the count of timestamps to retrieve.
+	process(forwardStream tsoStream, count uint32, tsoProtoFactory TSOProtoFactory) (tsoResp, error)
+	// postProcess sends the response back to the sender of the request
+	postProcess(countSum, physical, firstLogical int64, suffixBits uint32) (int64, error)
+}
+
+type tsoResponse interface {
+	GetTimestamp() *pdpb.Timestamp
+}
+
+type TSOProtoTSORequest struct {
+	forwardedHost   string
+	clientConn      *grpc.ClientConn
+	request         *tsopb.TsoRequest
+	stream          tsopb.TSO_TsoServer
+}
+
+func NewTSOProtoTSORequest(forwardedHost string, clientConn *grpc.ClientConn, request *tsopb.TsoRequest, stream tsopb.TSO_TsoServer) TSORequest {
+	tsoRequest := &TSOProtoTSORequest{
+		forwardedHost:   forwardedHost,
+		clientConn:      clientConn,
+		request:         request,
+		stream:          stream,
+	}
+	return tsoRequest
+}
+
+// getForwardedHost returns the forwarded host
+func (r *TSOProtoTSORequest) getForwardedHost() string {
+	return r.forwardedHost
+}
+
+// getClientConn returns the grpc client connection
+func (r *TSOProtoTSORequest) getClientConn() *grpc.ClientConn {
+	return r.clientConn
+}
+
+// getCount returns the count of timestamps to retrieve
+func (r *TSOProtoTSORequest) getCount() uint32 {
+	return r.request.GetCount()
+}
+
+// process sends request and receive response via stream.
+// count defins the count of timestamps to retrieve.
+func (r *TSOProtoTSORequest) process(forwardStream tsoStream, count uint32, tsoProtoFactory TSOProtoFactory) (tsoResp, error) {
+	return forwardStream.process(r.request.GetHeader().GetClusterId(), count,
+		r.request.GetHeader().GetKeyspaceId(), r.request.GetHeader().GetKeyspaceGroupId(), r.request.GetDcLocation())
+}
+
+// postProcess sends the response back to the sender of the request
+func (r *TSOProtoTSORequest) postProcess(countSum, physical, firstLogical int64, suffixBits uint32) (int64, error) {
+	count := r.request.GetCount()
+	countSum += int64(count)
+	response := &tsopb.TsoResponse{
+		Header: &tsopb.ResponseHeader{ClusterId: r.request.GetHeader().GetClusterId()},
+		Count:  count,
+		Timestamp: &pdpb.Timestamp{
+			Physical:   physical,
+			Logical:    addLogical(firstLogical, countSum, suffixBits),
+			SuffixBits: suffixBits,
+		},
+	}
+	// Send back to the client.
+	if err := r.stream.Send(response); err != nil {
+		return countSum, err
+	}
+	return countSum, nil
+}
+
+type PDProtoTSORequest struct {
+	forwardedHost string
+	clientConn    *grpc.ClientConn
+	request       *pdpb.TsoRequest
+	stream        pdpb.PD_TsoServer
+}
+
+func NewPDProtoTSORequest(forwardedHost string, clientConn *grpc.ClientConn, request *pdpb.TsoRequest, stream pdpb.PD_TsoServer) TSORequest {
+	tsoRequest := &PDProtoTSORequest{
+		forwardedHost: forwardedHost,
+		clientConn:    clientConn,
+		request:       request,
+		stream:        stream,
+	}
+	return tsoRequest
+}
+
+// getForwardedHost returns the forwarded host
+func (r *PDProtoTSORequest) getForwardedHost() string {
+	return r.forwardedHost
+}
+
+// getClientConn returns the grpc client connection
+func (r *PDProtoTSORequest) getClientConn() *grpc.ClientConn {
+	return r.clientConn
+}
+
+// getCount returns the count of timestamps to retrieve
+func (r *PDProtoTSORequest) getCount() uint32 {
+	return r.request.GetCount()
+}
+
+// process sends request and receive response via stream.
+// count defins the count of timestamps to retrieve.
+func (r *PDProtoTSORequest) process(forwardStream tsoStream, count uint32, tsoProtoFactory TSOProtoFactory) (tsoResp, error) {
+	return forwardStream.process(r.request.GetHeader().GetClusterId(), count,
+		utils.DefaultKeySpaceID, utils.DefaultKeySpaceGroupID, r.request.GetDcLocation())
+}
+
+// postProcess sends the response back to the sender of the request
+func (r *PDProtoTSORequest) postProcess(countSum, physical, firstLogical int64, suffixBits uint32) (int64, error) {
+	count := r.request.GetCount()
+	countSum += int64(count)
+	response := &pdpb.TsoResponse{
+		Header: &pdpb.ResponseHeader{ClusterId: r.request.GetHeader().GetClusterId()},
+		Count:  count,
+		Timestamp: &pdpb.Timestamp{
+			Physical:   physical,
+			Logical:    addLogical(firstLogical, countSum, suffixBits),
+			SuffixBits: suffixBits,
+		},
+	}
+	// Send back to the client.
+	if err := r.stream.Send(response); err != nil {
+		return countSum, err
+	}
+	return countSum, nil
+}
diff --git a/pkg/utils/tsoutil/tso.go b/pkg/utils/tsoutil/tsoutil.go
similarity index 100%
rename from pkg/utils/tsoutil/tso.go
rename to pkg/utils/tsoutil/tsoutil.go
diff --git a/server/grpc_service.go b/server/grpc_service.go
index 8ac952e706e0..3c9cd6a47630 100644
--- a/server/grpc_service.go
+++ b/server/grpc_service.go
@@ -49,10 +49,6 @@ import (
 
 const (
 	heartbeatSendTimeout = 5 * time.Second
-
-	// tso
-	maxMergeTSORequests    = 10000
-	defaultTSOProxyTimeout = 3 * time.Second
 )
 
 // gRPC errors
@@ -186,38 +182,36 @@ func (s *GrpcServer) Tso(stream pdpb.PD_TsoServer) error {
 		}
 
 		streamCtx := stream.Context()
-		if s.IsAPIServiceMode() {
-			forwardedHost, ok := s.GetServicePrimaryAddr(ctx, utils.TSOServiceName)
-			if !ok || forwardedHost == "" {
-				return ErrNotFoundTSOAddr
+		forwardedHost := grpcutil.GetForwardedHost(streamCtx)
+		if s.IsAPIServiceMode() || !s.isLocalRequest(forwardedHost) {
+			if s.IsAPIServiceMode() {
+				ok := true
+				forwardedHost, ok = s.GetServicePrimaryAddr(ctx, utils.TSOServiceName)
+				if !ok || len(forwardedHost) == 0 {
+					return ErrNotFoundTSOAddr
+				}
 			}
+
+			clientConn, err := s.getDelegateClient(s.ctx, forwardedHost)
 			if err != nil {
-				return errors.WithStack(err)
+				return err
 			}
+
 			if errCh == nil {
 				doneCh = make(chan struct{})
 				defer close(doneCh)
 				errCh = make(chan error)
 			}
-			s.dispatchTSORequest(ctx, &tsoRequest{
-				forwardedHost,
-				request,
-				stream,
-			}, forwardedHost, doneCh, errCh, true)
-			continue
-		}
-		forwardedHost := grpcutil.GetForwardedHost(streamCtx)
-		if !s.isLocalRequest(forwardedHost) {
-			if errCh == nil {
-				doneCh = make(chan struct{})
-				defer close(doneCh)
-				errCh = make(chan error)
+
+			var tsoProtoFactory tsoutil.TSOProtoFactory
+			if s.IsAPIServiceMode() {
+				tsoProtoFactory = s.tsoProtoTSOFactory				
+			} else {
+				tsoProtoFactory = s.tsoProtoPDFactory
 			}
-			s.dispatchTSORequest(ctx, &tsoRequest{
-				forwardedHost,
-				request,
-				stream,
-			}, forwardedHost, doneCh, errCh, false)
+
+			tsoRequest := tsoutil.NewPDProtoTSORequest(forwardedHost, clientConn, request, stream)
+			s.tsoDispatcher.DispatchRequest(ctx, tsoRequest, tsoProtoFactory, doneCh, errCh)
 			continue
 		}
 
@@ -246,217 +240,6 @@ func (s *GrpcServer) Tso(stream pdpb.PD_TsoServer) error {
 	}
 }
 
-type tsoRequest struct {
-	forwardedHost string
-	request       *pdpb.TsoRequest
-	stream        pdpb.PD_TsoServer
-}
-
-func (s *GrpcServer) dispatchTSORequest(ctx context.Context, request *tsoRequest, forwardedHost string, doneCh <-chan struct{}, errCh chan<- error, withTSOProto bool) {
-	tsoRequestChInterface, loaded := s.tsoDispatcher.LoadOrStore(forwardedHost, make(chan *tsoRequest, maxMergeTSORequests))
-	if !loaded {
-		tsDeadlineCh := make(chan deadline, 1)
-		go s.handleDispatcher(ctx, forwardedHost, tsoRequestChInterface.(chan *tsoRequest), tsDeadlineCh, doneCh, errCh, withTSOProto)
-		go watchTSDeadline(ctx, tsDeadlineCh)
-	}
-	tsoRequestChInterface.(chan *tsoRequest) <- request
-}
-
-func (s *GrpcServer) handleDispatcher(ctx context.Context, forwardedHost string, tsoRequestCh <-chan *tsoRequest, tsDeadlineCh chan<- deadline, doneCh <-chan struct{}, errCh chan<- error, withTSOProto bool) {
-	dispatcherCtx, ctxCancel := context.WithCancel(ctx)
-	defer ctxCancel()
-	defer s.tsoDispatcher.Delete(forwardedHost)
-
-	var (
-		forwardStream    pdpb.PD_TsoClient
-		forwardMCSStream tsopb.TSO_TsoClient
-		cancel           context.CancelFunc
-	)
-	client, err := s.getDelegateClient(ctx, forwardedHost)
-	if err != nil {
-		goto errHandling
-	}
-	log.Info("create tso forward stream", zap.String("forwarded-host", forwardedHost))
-	if withTSOProto {
-		forwardMCSStream, cancel, err = s.createMCSTSOForwardStream(client)
-	} else {
-		forwardStream, cancel, err = s.createTsoForwardStream(client)
-	}
-errHandling:
-	if err != nil || (forwardStream == nil && !withTSOProto) || (forwardMCSStream == nil && withTSOProto) {
-		log.Error("create tso forwarding stream error", zap.String("forwarded-host", forwardedHost), errs.ZapError(errs.ErrGRPCCreateStream, err))
-		select {
-		case <-dispatcherCtx.Done():
-			return
-		case _, ok := <-doneCh:
-			if !ok {
-				return
-			}
-		case errCh <- err:
-			close(errCh)
-			return
-		}
-	}
-	defer cancel()
-
-	requests := make([]*tsoRequest, maxMergeTSORequests+1)
-	for {
-		select {
-		case first := <-tsoRequestCh:
-			pendingTSOReqCount := len(tsoRequestCh) + 1
-			requests[0] = first
-			for i := 1; i < pendingTSOReqCount; i++ {
-				requests[i] = <-tsoRequestCh
-			}
-			done := make(chan struct{})
-			dl := deadline{
-				timer:  time.After(defaultTSOProxyTimeout),
-				done:   done,
-				cancel: cancel,
-			}
-			select {
-			case tsDeadlineCh <- dl:
-			case <-dispatcherCtx.Done():
-				return
-			}
-			err = s.processTSORequests(forwardStream, forwardMCSStream, requests[:pendingTSOReqCount])
-			close(done)
-			if err != nil {
-				log.Error("proxy forward tso error", zap.String("forwarded-host", forwardedHost), errs.ZapError(errs.ErrGRPCSend, err))
-				select {
-				case <-dispatcherCtx.Done():
-					return
-				case _, ok := <-doneCh:
-					if !ok {
-						return
-					}
-				case errCh <- err:
-					close(errCh)
-					return
-				}
-			}
-		case <-dispatcherCtx.Done():
-			return
-		}
-	}
-}
-
-type tsoResp interface {
-	GetTimestamp() *pdpb.Timestamp
-}
-
-func (s *GrpcServer) processTSORequests(forwardStream pdpb.PD_TsoClient, forwardMCSStream tsopb.TSO_TsoClient, requests []*tsoRequest) error {
-	start := time.Now()
-	// Merge the requests
-	count := uint32(0)
-	for _, request := range requests {
-		count += request.request.GetCount()
-	}
-	var (
-		resp tsoResp
-		err  error
-	)
-	if forwardStream != nil {
-		req := &pdpb.TsoRequest{
-			Header: requests[0].request.GetHeader(),
-			Count:  count,
-			// TODO: support Local TSO proxy forwarding.
-			DcLocation: requests[0].request.GetDcLocation(),
-		}
-		// Send to the tso server stream
-		if err := forwardStream.Send(req); err != nil {
-			return err
-		}
-		resp, err = forwardStream.Recv()
-		if err != nil {
-			return err
-		}
-	}
-	if forwardMCSStream != nil {
-		req := &tsopb.TsoRequest{
-			Header: &tsopb.RequestHeader{
-				ClusterId:       requests[0].request.GetHeader().GetClusterId(),
-				KeyspaceId:      utils.DefaultKeySpaceID,
-				KeyspaceGroupId: utils.DefaultKeySpaceGroupID,
-			},
-			Count: count,
-			// TODO: support Local TSO proxy forwarding.
-			DcLocation: requests[0].request.GetDcLocation(),
-		}
-		// Send to the tso server stream.
-		if err := forwardMCSStream.Send(req); err != nil {
-			return err
-		}
-		resp, err = forwardMCSStream.Recv()
-		if err != nil {
-			return err
-		}
-	}
-	tsoProxyHandleDuration.Observe(time.Since(start).Seconds())
-	tsoProxyBatchSize.Observe(float64(count))
-	// Split the response
-	physical, logical, suffixBits := resp.GetTimestamp().GetPhysical(), resp.GetTimestamp().GetLogical(), resp.GetTimestamp().GetSuffixBits()
-	// `logical` is the largest ts's logical part here, we need to do the subtracting before we finish each TSO request.
-	// This is different from the logic of client batch, for example, if we have a largest ts whose logical part is 10,
-	// count is 5, then the splitting results should be 5 and 10.
-	firstLogical := addLogical(logical, -int64(count), suffixBits)
-	return s.finishTSORequest(requests, physical, firstLogical, suffixBits)
-}
-
-// Because of the suffix, we need to shift the count before we add it to the logical part.
-func addLogical(logical, count int64, suffixBits uint32) int64 {
-	return logical + count<<suffixBits
-}
-
-func (s *GrpcServer) finishTSORequest(requests []*tsoRequest, physical, firstLogical int64, suffixBits uint32) error {
-	countSum := int64(0)
-	for i := 0; i < len(requests); i++ {
-		count := requests[i].request.GetCount()
-		countSum += int64(count)
-		response := &pdpb.TsoResponse{
-			Header: s.header(),
-			Count:  count,
-			Timestamp: &pdpb.Timestamp{
-				Physical:   physical,
-				Logical:    addLogical(firstLogical, countSum, suffixBits),
-				SuffixBits: suffixBits,
-			},
-		}
-		// Send back to the client.
-		if err := requests[i].stream.Send(response); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-type deadline struct {
-	timer  <-chan time.Time
-	done   chan struct{}
-	cancel context.CancelFunc
-}
-
-func watchTSDeadline(ctx context.Context, tsDeadlineCh <-chan deadline) {
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-	for {
-		select {
-		case d := <-tsDeadlineCh:
-			select {
-			case <-d.timer:
-				log.Error("tso proxy request processing is canceled due to timeout", errs.ZapError(errs.ErrProxyTSOTimeout))
-				d.cancel()
-			case <-d.done:
-				continue
-			case <-ctx.Done():
-				return
-			}
-		case <-ctx.Done():
-			return
-		}
-	}
-}
-
 // Bootstrap implements gRPC PDServer.
 func (s *GrpcServer) Bootstrap(ctx context.Context, request *pdpb.BootstrapRequest) (*pdpb.BootstrapResponse, error) {
 	fn := func(ctx context.Context, client *grpc.ClientConn) (interface{}, error) {
@@ -1885,24 +1668,6 @@ func (s *GrpcServer) isLocalRequest(forwardedHost string) bool {
 	return false
 }
 
-func (s *GrpcServer) createTsoForwardStream(client *grpc.ClientConn) (pdpb.PD_TsoClient, context.CancelFunc, error) {
-	done := make(chan struct{})
-	ctx, cancel := context.WithCancel(s.ctx)
-	go checkStream(ctx, cancel, done)
-	forwardStream, err := pdpb.NewPDClient(client).Tso(ctx)
-	done <- struct{}{}
-	return forwardStream, cancel, err
-}
-
-func (s *GrpcServer) createMCSTSOForwardStream(client *grpc.ClientConn) (tsopb.TSO_TsoClient, context.CancelFunc, error) {
-	done := make(chan struct{})
-	ctx, cancel := context.WithCancel(s.ctx)
-	go checkStream(ctx, cancel, done)
-	forwardStream, err := tsopb.NewTSOClient(client).Tso(ctx)
-	done <- struct{}{}
-	return forwardStream, cancel, err
-}
-
 func (s *GrpcServer) createHeartbeatForwardStream(client *grpc.ClientConn) (pdpb.PD_RegionHeartbeatClient, context.CancelFunc, error) {
 	done := make(chan struct{})
 	ctx, cancel := context.WithCancel(s.ctx)
@@ -1997,7 +1762,7 @@ func (s *GrpcServer) getTSOForwardStream(ctx context.Context, forwardedHost stri
 		return nil, err
 	}
 	done := make(chan struct{})
-	ctx, cancel := context.WithTimeout(s.ctx, defaultTSOProxyTimeout)
+	ctx, cancel := context.WithTimeout(s.ctx, tsoutil.DefaultTSOProxyTimeout)
 	go checkStream(ctx, cancel, done)
 	forwardStream, err := tsopb.NewTSOClient(client).Tso(ctx)
 	if err != nil {
diff --git a/server/server.go b/server/server.go
index 0f9d9f2d6a43..394f81d2a070 100644
--- a/server/server.go
+++ b/server/server.go
@@ -192,7 +192,13 @@ type Server struct {
 	tsoClients sync.Map
 	// tsoDispatcher is used to dispatch different TSO requests to
 	// the corresponding forwarding TSO channel.
-	tsoDispatcher sync.Map /* Store as map[string]chan *tsoRequest */
+	tsoDispatcher *tsoutil.TSODispatcher
+	// tsoProtoTSOFactory is the abstract factory for creating tso
+	// related data structures defined in pd protocol 
+	tsoProtoTSOFactory *tsoutil.TSOProtoTSOFactory
+	// tsoProtoTSOFactory is the abstract factory for creating tso
+	// related data structures defined in tso protocol 
+	tsoProtoPDFactory *tsoutil.TSOProtoPDFactory
 
 	serviceRateLimiter *ratelimit.Limiter
 	serviceLabels      map[string][]apiutil.AccessPath
@@ -386,6 +392,9 @@ func (s *Server) startServer(ctx context.Context) error {
 	}
 	defaultStorage := storage.NewStorageWithEtcdBackend(s.client, s.rootPath)
 	s.storage = storage.NewCoreStorage(defaultStorage, regionStorage)
+	s.tsoDispatcher = tsoutil.NewTSODispatcher(tsoProxyHandleDuration, tsoProxyBatchSize)
+	s.tsoProtoTSOFactory = &tsoutil.TSOProtoTSOFactory{}
+	s.tsoProtoPDFactory = &tsoutil.TSOProtoPDFactory{}
 	if !s.IsAPIServiceMode() {
 		s.tsoAllocatorManager = tso.NewAllocatorManager(
 			s.member, s.rootPath, s.storage, s.cfg.IsLocalTSOEnabled(), s.cfg.GetTSOSaveInterval(), s.cfg.GetTSOUpdatePhysicalInterval(), s.cfg.GetTLSConfig(),