mcs: add timeout when call alloc of api server (tikv#7254)

close tikv#7207 Signed-off-by: lhy1024 <admin@liudos.us>
rleungx · Dec 1, 2023 · 91a198b · 91a198b
1 parent 0c4819f
commit 91a198b
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/pkg/mcs/scheduling/server/cluster.go b/pkg/mcs/scheduling/server/cluster.go
@@ -52,7 +52,10 @@ type Cluster struct {
 	running           atomic.Bool
 }
 
-const regionLabelGCInterval = time.Hour
+const (
+	regionLabelGCInterval = time.Hour
+	requestTimeout        = 3 * time.Second
+)
 
 // NewCluster creates a new cluster.
 func NewCluster(parentCtx context.Context, persistConfig *config.PersistConfig, storage storage.Storage, basicCluster *core.BasicCluster, hbStreams *hbstream.HeartbeatStreams, clusterID uint64, checkMembershipCh chan struct{}) (*Cluster, error) {
@@ -199,7 +202,9 @@ func (c *Cluster) AllocID() (uint64, error) {
 	if err != nil {
 		return 0, err
 	}
-	resp, err := client.AllocID(c.ctx, &pdpb.AllocIDRequest{Header: &pdpb.RequestHeader{ClusterId: c.clusterID}})
+	ctx, cancel := context.WithTimeout(c.ctx, requestTimeout)
+	defer cancel()
+	resp, err := client.AllocID(ctx, &pdpb.AllocIDRequest{Header: &pdpb.RequestHeader{ClusterId: c.clusterID}})
 	if err != nil {
 		c.checkMembershipCh <- struct{}{}
 		return 0, err