stash

Signed-off-by: MyonKeminta <MyonKeminta@users.noreply.github.com>

Author: MyonKeminta

client/metrics.go

@@ -39,13 +39,14 @@ func initAndRegisterMetrics(constLabels prometheus.Labels) {
 }
 
 var (
-	cmdDuration         *prometheus.HistogramVec
-	cmdFailedDuration   *prometheus.HistogramVec
-	requestDuration     *prometheus.HistogramVec
-	tsoBestBatchSize    prometheus.Histogram
-	tsoBatchSize        prometheus.Histogram
-	tsoBatchSendLatency prometheus.Histogram
-	requestForwarded    *prometheus.GaugeVec
+	cmdDuration              *prometheus.HistogramVec
+	cmdFailedDuration        *prometheus.HistogramVec
+	requestDuration          *prometheus.HistogramVec
+	tsoBestBatchSize         prometheus.Histogram
+	tsoBatchSize             prometheus.Histogram
+	tsoBatchSendLatency      prometheus.Histogram
+	requestForwarded         *prometheus.GaugeVec
+	ongoingRequestCountGauge *prometheus.GaugeVec
 )
 
 func initMetrics(constLabels prometheus.Labels) {
@@ -117,6 +118,15 @@ func initMetrics(constLabels prometheus.Labels) {
 			Help:        "The status to indicate if the request is forwarded",
 			ConstLabels: constLabels,
 		}, []string{"host", "delegate"})
+
+	ongoingRequestCountGauge = prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Namespace:   "pd_client",
+			Subsystem:   "request",
+			Name:        "ongoing_requests_count",
+			Help:        "Current count of ongoing batch tso requests",
+			ConstLabels: constLabels,
+		}, []string{"stream"})
 }
 
 var (

client/tso_dispatcher.go

@@ -81,6 +81,8 @@ type tsoDispatcher struct {
 	tsDeadlineCh    chan *deadline
 	lastTSOInfo     *tsoInfo
 
+	batchBufferPool sync.Pool
+
 	updateConnectionCtxsCh chan struct{}
 }
 
@@ -186,7 +188,7 @@ func (td *tsoDispatcher) handleDispatcher(wg *sync.WaitGroup) {
 		streamCtx context.Context
 		cancel    context.CancelFunc
 		streamURL string
-		stream    *tsoStream
+		stream    *tsoStream[notifier]
 	)
 	// Loop through each batch of TSO requests and send them for processing.
 	streamLoopTimer := time.NewTimer(option.timeout)
@@ -422,26 +424,44 @@ func (td *tsoDispatcher) processRequests(
 		keyspaceID         = svcDiscovery.GetKeyspaceID()
 		reqKeyspaceGroupID = svcDiscovery.GetKeyspaceGroupID()
 	)
-	respKeyspaceGroupID, physical, logical, suffixBits, err := stream.processRequests(
+
+	cb := func(result tsoRequestResult, reqKeyspaceGroupID uint32, streamURL string, err error) {
+		curTSOInfo := &tsoInfo{
+			tsoServer:           stream.getServerURL(),
+			reqKeyspaceGroupID:  reqKeyspaceGroupID,
+			respKeyspaceGroupID: result.respKeyspaceGroupID,
+			respReceivedAt:      time.Now(),
+			physical:            result.physical,
+			logical:             result.logical,
+		}
+		// `logical` is the largest ts's logical part here, we need to do the subtracting before we finish each TSO request.
+		firstLogical := tsoutil.AddLogical(result.logical, -int64(result.count)+1, result.suffixBits)
+		td.compareAndSwapTS(curTSOInfo, firstLogical)
+		finishCollectedRequests(tbc.getCollectedRequests(), result.physical, firstLogical, result.suffixBits, err)
+		tbc.collectedRequestCount = 0
+		td.batchBufferPool.Put(tbc)
+	}
+
+	err := stream.processRequests(
 		clusterID, keyspaceID, reqKeyspaceGroupID,
-		dcLocation, count, tbc.batchStartTime)
+		dcLocation, count, tbc.batchStartTime, cb)
 	if err != nil {
 		tbc.finishCollectedRequests(0, 0, 0, err)
 		return err
 	}
-	curTSOInfo := &tsoInfo{
-		tsoServer:           stream.getServerURL(),
-		reqKeyspaceGroupID:  reqKeyspaceGroupID,
-		respKeyspaceGroupID: respKeyspaceGroupID,
-		respReceivedAt:      time.Now(),
-		physical:            physical,
-		logical:             logical,
-	}
-	// `logical` is the largest ts's logical part here, we need to do the subtracting before we finish each TSO request.
-	firstLogical := tsoutil.AddLogical(logical, -count+1, suffixBits)
-	td.compareAndSwapTS(curTSOInfo, firstLogical)
-	tbc.finishCollectedRequests(physical, firstLogical, suffixBits, nil)
-	return nil
+	//curTSOInfo := &tsoInfo{
+	//	tsoServer:           stream.getServerURL(),
+	//	reqKeyspaceGroupID:  reqKeyspaceGroupID,
+	//	respKeyspaceGroupID: respKeyspaceGroupID,
+	//	respReceivedAt:      time.Now(),
+	//	physical:            physical,
+	//	logical:             logical,
+	//}
+	//// `logical` is the largest ts's logical part here, we need to do the subtracting before we finish each TSO request.
+	//firstLogical := tsoutil.AddLogical(logical, -count+1, suffixBits)
+	//td.compareAndSwapTS(curTSOInfo, firstLogical)
+	//tbc.finishCollectedRequests(physical, firstLogical, suffixBits, nil)
+	//return nil
 }
 
 func (td *tsoDispatcher) compareAndSwapTS(
@@ -482,3 +502,24 @@ func (td *tsoDispatcher) compareAndSwapTS(
 	}
 	td.lastTSOInfo = curTSOInfo
 }
+
+type notifier struct {
+	tbc *tsoBatchController
+	td  *tsoDispatcher
+}
+
+func (n notifier) finish(result tsoRequestResult, err error) {
+	curTSOInfo := &tsoInfo{
+		tsoServer:           stream.getServerURL(),
+		reqKeyspaceGroupID:  reqKeyspaceGroupID,
+		respKeyspaceGroupID: result.respKeyspaceGroupID,
+		respReceivedAt:      time.Now(),
+		physical:            result.physical,
+		logical:             result.logical,
+	}
+	// `logical` is the largest ts's logical part here, we need to do the subtracting before we finish each TSO request.
+	firstLogical := tsoutil.AddLogical(result.logical, -int64(result.count)+1, result.suffixBits)
+	n.td.compareAndSwapTS(curTSOInfo, firstLogical)
+	finishCollectedRequests(n.tbc.getCollectedRequests(), result.physical, firstLogical, result.suffixBits, err)
+	n.tbc.collectedRequestCount = 0
+}