server,tests: add tso batch metric and fix test server lock

Signed-off-by: okjiang <819421878@qq.com>

Author: okJiang

metrics/grafana/pd.json

@@ -11003,7 +11003,7 @@
           "targets": [
             {
               "exemplar": true,
-              "expr": "histogram_quantile(0.99, sum(rate(pd_client_request_handle_tso_batch_size_bucket{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\"}[1m])) by (le))",
+              "expr": "histogram_quantile(0.99, sum(rate(pd_server_handle_tso_batch_size_bucket{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\"}[1m])) by (le))",
               "hide": false,
               "interval": "",
               "intervalFactor": 1,
@@ -11012,7 +11012,7 @@
             },
             {
               "exemplar": true,
-              "expr": "histogram_quantile(0.90, sum(rate(pd_client_request_handle_tso_batch_size_bucket{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\"}[1m])) by (le))",
+              "expr": "histogram_quantile(0.90, sum(rate(pd_server_handle_tso_batch_size_bucket{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\"}[1m])) by (le))",
               "hide": false,
               "interval": "",
               "intervalFactor": 1,
@@ -11021,7 +11021,7 @@
             },
             {
               "exemplar": true,
-              "expr": "sum(rate(pd_client_request_handle_tso_batch_size_sum{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\"}[1m])) / sum(rate(pd_client_request_handle_tso_batch_size_count{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\"}[1m]))",
+              "expr": "sum(rate(pd_server_handle_tso_batch_size_sum{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\"}[1m])) / sum(rate(pd_server_handle_tso_batch_size_count{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\"}[1m]))",
               "hide": false,
               "interval": "",
               "intervalFactor": 1,
@@ -16269,4 +16269,4 @@
   "title": "Test-Cluster-PD",
   "uid": "Q6RuHYIWk",
   "version": 1
-}
+}

server/grpc_service.go

@@ -569,6 +569,7 @@ func (s *GrpcServer) Tso(stream pdpb.PD_TsoServer) error {
 				err = errs.ErrGenerateTimestamp.FastGenByArgs("tso count should be positive")
 				return errs.ErrUnknown(err)
 			}
+			tsoBatchSize.Observe(float64(request.GetCount()))
 			tsoStreamErr, err = s.handleTSOForwarding(stream.Context(), forwarder, request, tsDeadlineCh)
 			if tsoStreamErr != nil {
 				return tsoStreamErr
@@ -584,6 +585,7 @@ func (s *GrpcServer) Tso(stream pdpb.PD_TsoServer) error {
 			return errs.ErrMismatchClusterID(clusterID, request.GetHeader().GetClusterId())
 		}
 		count := request.GetCount()
+		tsoBatchSize.Observe(float64(count))
 		ctx, task := trace.NewTask(ctx, "tso")
 		ts, err := s.tsoAllocator.GenerateTSO(ctx, count)
 		task.End()

server/metrics.go

@@ -101,6 +101,15 @@ var (
 			Buckets:   prometheus.ExponentialBuckets(0.0005, 2, 13),
 		})
 
+	tsoBatchSize = prometheus.NewHistogram(
+		prometheus.HistogramOpts{
+			Namespace: "pd",
+			Subsystem: "server",
+			Name:      "handle_tso_batch_size",
+			Help:      "Bucketed histogram of the batch size of handled tso requests.",
+			Buckets:   prometheus.ExponentialBuckets(1, 2, 13),
+		})
+
 	queryRegionDuration = prometheus.NewHistogram(
 		prometheus.HistogramOpts{
 			Namespace: "pd",
@@ -207,6 +216,7 @@ func init() {
 	prometheus.MustRegister(tsoProxyBatchSize)
 	prometheus.MustRegister(tsoProxyForwardTimeoutCounter)
 	prometheus.MustRegister(tsoHandleDuration)
+	prometheus.MustRegister(tsoBatchSize)
 	prometheus.MustRegister(queryRegionDuration)
 	prometheus.MustRegister(regionHeartbeatHandleDuration)
 	prometheus.MustRegister(storeHeartbeatHandleDuration)

tests/cluster.go

@@ -59,6 +59,7 @@ import (
 const (
 	Initial int32 = iota
 	Running
+	Starting
 	Stop
 	Destroy
 )
@@ -142,22 +143,35 @@ func NewTestServer(ctx context.Context, cfg *config.Config, services []string, h
 // Run starts to run a TestServer.
 func (s *TestServer) Run() error {
 	s.Lock()
-	defer s.Unlock()
 	if s.state != Initial && s.state != Stop {
+		s.Unlock()
 		return errors.Errorf("server(state%d) cannot run", s.state)
 	}
+	s.state = Starting
+	s.Unlock()
+
 	if err := s.server.Run(); err != nil {
+		s.Lock()
+		if s.state == Starting {
+			s.state = Stop
+		}
+		s.Unlock()
 		return err
 	}
-	s.state = Running
+
+	s.Lock()
+	if s.state == Starting {
+		s.state = Running
+	}
+	s.Unlock()
 	return nil
 }
 
 // Stop is used to stop a TestServer.
 func (s *TestServer) Stop() error {
 	s.Lock()
 	defer s.Unlock()
-	if s.state != Running {
+	if s.state != Running && s.state != Starting {
 		return errors.Errorf("server(state%d) cannot stop", s.state)
 	}
 	s.server.Close()
@@ -169,7 +183,7 @@ func (s *TestServer) Stop() error {
 func (s *TestServer) Destroy() error {
 	s.Lock()
 	defer s.Unlock()
-	if s.state == Running {
+	if s.state == Running || s.state == Starting {
 		s.server.Close()
 	}
 	if err := os.RemoveAll(s.server.GetConfig().DataDir); err != nil {
@@ -699,7 +713,7 @@ func (c *TestCluster) runInitialServersWithRetry(maxRetries int) error {
 
 			// Stop and destroy all servers
 			for _, s := range servers {
-				if s.State() == Running {
+				if state := s.State(); state == Running || state == Starting {
 					_ = s.Stop()
 				}
 				_ = s.Destroy()
@@ -738,7 +752,7 @@ func (c *TestCluster) runInitialServersWithRetry(maxRetries int) error {
 		case strings.Contains(errMsg, "ErrStartEtcd"):
 			log.Warn("etcd start failed, will retry", zap.Error(lastErr))
 			for _, s := range servers {
-				if s.State() == Running {
+				if state := s.State(); state == Running || state == Starting {
 					_ = s.Stop()
 				}
 			}
@@ -767,7 +781,7 @@ func RunServersWithRetry(servers []*TestServer, maxRetries int) error {
 			log.Warn("etcd start failed, will retry", zap.Error(lastErr))
 			// Stop any partially started servers before retrying
 			for _, s := range servers {
-				if s.State() == Running {
+				if state := s.State(); state == Running || state == Starting {
 					_ = s.Stop()
 				}
 			}