Repository: tikv/pd (go)

Issue:
Batch alloc ID for scheduling service
## Enhancement Task

https://github.com/tikv/pd/blob/daa23f832fb95ecda6c622beeaec3af30fd3f4a9/pkg/mcs/scheduling/server/grpc_service.go#L291-L311

For each split request, we need to call `AllocID` several times. When there are a large number of split requests, like import data scenarios, it costs too much resources. We can support batch alloc ID to reduce the call times.

Here is a profile:
![Image](https://github.com/user-attachments/assets/8e471474-5e0c-474b-be57-f88c661b5422)

Pull request:
mcs: support batch allocating ids
<!--

Thank you for working on PD! Please read PD's [CONTRIBUTING](https://github.com/tikv/pd/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.

PR Title Format:
1. pkg [, pkg2, pkg3]: what's changed
2. *: what's changed

-->

### What problem does this PR solve?

<!--

Please create an issue first to describe the problem.
There MUST be one line starting with "Issue Number:  " and 
linking the relevant issues via the "close" or "ref".
For more info, check https://github.com/tikv/pd/blob/master/CONTRIBUTING.md#linking-issues.

-->
Issue Number: Close #9085.

### What is changed and how does it work?

kvproto PR: https://github.com/pingcap/kvproto/pull/1296

<!--

You could use the "commit message" code block to add more description to the final commit message.
For more info, check https://github.com/tikv/pd/blob/master/CONTRIBUTING.md#format-of-the-commit-message.

-->

```commit-message
```

### Check List

<!-- Remove the items that are not applicable. -->

Tests

<!-- At least one of these tests must be included. -->

- Unit test

### Release note

<!--

A bugfix or a new feature needs a release note. If there is no need to give a release note, just leave it with the `None`.

Please refer to [Release Notes Language Style Guide](https://pingcap.github.io/tidb-dev-guide/contribute-to-tidb/release-notes-style-guide.html) to write a quality release note.

-->

```release-note
None.
```


Patch:
diff --git a/client/go.mod b/client/go.mod
index 08c99ff2e..dc4795e49 100644
--- a/client/go.mod
+++ b/client/go.mod
@@ -10,7 +10,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86
-	github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1
+	github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/prometheus/client_golang v1.20.5
 	github.com/stretchr/testify v1.9.0
diff --git a/client/go.sum b/client/go.sum
index b3db161c9..e3b1fd92e 100644
--- a/client/go.sum
+++ b/client/go.sum
@@ -49,8 +49,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c h1:xpW9bvK+HuuTm
 github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ueLEUZgtx2cIogM0v4Zj5uvvzhuuiu7Pn8HzMPg=
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
-github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1 h1:rTAyiswGyWSGHJVa4Mkhdi8YfGqfA4LrUVKsH9nrJ8E=
-github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d h1:52qhTQG8G8V/pHo/w7F4d2Tw98KMk2C+gAe3U8SWRAg=
+github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
diff --git a/go.mod b/go.mod
index f2d9ca12c..32e5b6fcb 100644
--- a/go.mod
+++ b/go.mod
@@ -34,7 +34,7 @@ require (
 	github.com/pingcap/errcode v0.3.0
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86
-	github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1
+	github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21
 	github.com/pingcap/tidb-dashboard v0.0.0-20250219061340-d62018124ae2
diff --git a/go.sum b/go.sum
index def2f938c..4c5064e7b 100644
--- a/go.sum
+++ b/go.sum
@@ -392,8 +392,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ue
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1 h1:rTAyiswGyWSGHJVa4Mkhdi8YfGqfA4LrUVKsH9nrJ8E=
-github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d h1:52qhTQG8G8V/pHo/w7F4d2Tw98KMk2C+gAe3U8SWRAg=
+github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
diff --git a/pkg/id/id.go b/pkg/id/id.go
index cb38c2326..b4d1be28f 100644
--- a/pkg/id/id.go
+++ b/pkg/id/id.go
@@ -43,7 +43,7 @@ type Allocator interface {
 	// SetBase set base id
 	SetBase(newBase uint64) error
 	// Alloc allocs a unique id.
-	Alloc() (uint64, error)
+	Alloc(count uint32) (uint64, uint32, error)
 	// Rebase resets the base for the allocator from the persistent window boundary,
 	// which also resets the end of the allocator. (base, end) is the range that can
 	// be allocated in memory.
@@ -94,19 +94,21 @@ func NewAllocator(params *AllocatorParams) Allocator {
 }
 
 // Alloc returns a new id.
-func (alloc *allocatorImpl) Alloc() (uint64, error) {
+func (alloc *allocatorImpl) Alloc(count uint32) (uint64, uint32, error) {
 	alloc.mu.Lock()
 	defer alloc.mu.Unlock()
 
-	if alloc.base == alloc.end {
-		if err := alloc.rebaseLocked(true); err != nil {
-			return 0, err
+	for range count {
+		if alloc.base == alloc.end {
+			if err := alloc.rebaseLocked(true); err != nil {
+				return 0, 0, err
+			}
 		}
-	}
 
-	alloc.base++
+		alloc.base++
+	}
 
-	return alloc.base, nil
+	return alloc.base, count, nil
 }
 
 // SetBase sets the base.
diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go
index 93312ae4f..92968a6f8 100644
--- a/pkg/keyspace/keyspace.go
+++ b/pkg/keyspace/keyspace.go
@@ -683,7 +683,7 @@ func (manager *Manager) LoadRangeKeyspace(startID uint32, limit int) ([]*keyspac
 
 // allocID allocate a new keyspace id.
 func (manager *Manager) allocID() (uint32, error) {
-	id64, err := manager.idAllocator.Alloc()
+	id64, _, err := manager.idAllocator.Alloc(1)
 	if err != nil {
 		return 0, err
 	}
diff --git a/pkg/mcs/scheduling/server/cluster.go b/pkg/mcs/scheduling/server/cluster.go
index 2154b1b67..551444cb0 100644
--- a/pkg/mcs/scheduling/server/cluster.go
+++ b/pkg/mcs/scheduling/server/cluster.go
@@ -239,20 +239,20 @@ func (c *Cluster) GetSchedulerConfig() sc.SchedulerConfigProvider { return c.per
 // GetStoreConfig returns the store config.
 func (c *Cluster) GetStoreConfig() sc.StoreConfigProvider { return c.persistConfig }
 
-// AllocID allocates a new ID.
-func (c *Cluster) AllocID() (uint64, error) {
+// AllocID allocates new IDs.
+func (c *Cluster) AllocID(count uint32) (uint64, uint32, error) {
 	client, err := c.getPDLeaderClient()
 	if err != nil {
-		return 0, err
+		return 0, 0, err
 	}
 	ctx, cancel := context.WithTimeout(c.ctx, requestTimeout)
 	defer cancel()
-	resp, err := client.AllocID(ctx, &pdpb.AllocIDRequest{Header: &pdpb.RequestHeader{ClusterId: keypath.ClusterID()}})
+	resp, err := client.AllocID(ctx, &pdpb.AllocIDRequest{Header: &pdpb.RequestHeader{ClusterId: keypath.ClusterID()}, Count: count})
 	if err != nil {
 		c.triggerMembershipCheck()
-		return 0, err
+		return 0, 0, err
 	}
-	return resp.GetId(), nil
+	return resp.GetId(), resp.GetCount(), nil
 }
 
 func (c *Cluster) getPDLeaderClient() (pdpb.PDClient, error) {
diff --git a/pkg/mcs/scheduling/server/grpc_service.go b/pkg/mcs/scheduling/server/grpc_service.go
index 033397e3b..83346af6a 100644
--- a/pkg/mcs/scheduling/server/grpc_service.go
+++ b/pkg/mcs/scheduling/server/grpc_service.go
@@ -288,17 +288,19 @@ func (s *Service) AskBatchSplit(_ context.Context, request *schedulingpb.AskBatc
 	splitIDs := make([]*pdpb.SplitID, 0, splitCount)
 	recordRegions := make([]uint64, 0, splitCount+1)
 
-	for i := 0; i < int(splitCount); i++ {
-		newRegionID, err := c.AllocID()
-		if err != nil {
-			return nil, errs.ErrSchedulerNotFound.FastGenByArgs()
-		}
+	id, count, err := c.AllocID(splitCount * (1 + uint32(len(request.Region.Peers))))
+	if err != nil {
+		return nil, err
+	}
+	curID := id - uint64(count)
+	for range splitCount {
+		newRegionID := curID
+		curID++
 
 		peerIDs := make([]uint64, len(request.Region.Peers))
-		for i := 0; i < len(peerIDs); i++ {
-			if peerIDs[i], err = c.AllocID(); err != nil {
-				return nil, err
-			}
+		for j := 0; j < len(peerIDs); j++ {
+			peerIDs[j] = curID
+			curID++
 		}
 
 		recordRegions = append(recordRegions, newRegionID)
diff --git a/pkg/mock/mockcluster/mockcluster.go b/pkg/mock/mockcluster/mockcluster.go
index 62498bff8..9b6871062 100644
--- a/pkg/mock/mockcluster/mockcluster.go
+++ b/pkg/mock/mockcluster/mockcluster.go
@@ -121,8 +121,8 @@ func (mc *Cluster) GetStorage() storage.Storage {
 }
 
 // AllocID returns a new unique ID.
-func (mc *Cluster) AllocID() (uint64, error) {
-	return mc.IDAllocator.Alloc()
+func (mc *Cluster) AllocID(uint32) (uint64, uint32, error) {
+	return mc.IDAllocator.Alloc(1)
 }
 
 // UpdateRegionsLabelLevelStats updates the label level stats for the regions.
@@ -190,7 +190,7 @@ func hotRegionsFromStore(w *statistics.HotCache, storeID uint64, kind utils.RWTy
 
 // AllocPeer allocs a new peer on a store.
 func (mc *Cluster) AllocPeer(storeID uint64) (*metapb.Peer, error) {
-	peerID, err := mc.AllocID()
+	peerID, _, err := mc.AllocID(1)
 	if err != nil {
 		log.Error("failed to alloc peer", errs.ZapError(err))
 		return nil, err
@@ -363,7 +363,7 @@ func (mc *Cluster) AddRegionStoreWithLeader(storeID uint64, regionCount int, lea
 	}
 	mc.AddRegionStore(storeID, regionCount)
 	for range leaderCount {
-		id, _ := mc.AllocID()
+		id, _, _ := mc.AllocID(1)
 		mc.AddLeaderRegion(id, storeID)
 	}
 }
diff --git a/pkg/mock/mockid/mockid.go b/pkg/mock/mockid/mockid.go
index 7b4902a6a..36e16460b 100644
--- a/pkg/mock/mockid/mockid.go
+++ b/pkg/mock/mockid/mockid.go
@@ -27,8 +27,8 @@ func NewIDAllocator() *IDAllocator {
 }
 
 // Alloc returns a new id.
-func (alloc *IDAllocator) Alloc() (uint64, error) {
-	return atomic.AddUint64(&alloc.base, 1), nil
+func (alloc *IDAllocator) Alloc(uint32) (uint64, uint32, error) {
+	return atomic.AddUint64(&alloc.base, 1), 1, nil
 }
 
 // SetBase implements the IDAllocator interface.
diff --git a/pkg/replication/replication_mode.go b/pkg/replication/replication_mode.go
index 856a68a9b..fba66fea2 100644
--- a/pkg/replication/replication_mode.go
+++ b/pkg/replication/replication_mode.go
@@ -238,7 +238,7 @@ func (m *ModeManager) drSwitchToAsyncWait(availableStores []uint64) error {
 	m.Lock()
 	defer m.Unlock()
 
-	id, err := m.cluster.AllocID()
+	id, _, err := m.cluster.AllocID(1)
 	if err != nil {
 		log.Warn("failed to switch to async wait state", zap.String("replicate-mode", modeDRAutoSync), errs.ZapError(err))
 		return err
@@ -260,7 +260,7 @@ func (m *ModeManager) drSwitchToAsync(availableStores []uint64) error {
 }
 
 func (m *ModeManager) drSwitchToAsyncWithLock(availableStores []uint64) error {
-	id, err := m.cluster.AllocID()
+	id, _, err := m.cluster.AllocID(1)
 	if err != nil {
 		log.Warn("failed to switch to async state", zap.String("replicate-mode", modeDRAutoSync), errs.ZapError(err))
 		return err
@@ -292,7 +292,7 @@ func (m *ModeManager) drSwitchToSyncRecover() error {
 }
 
 func (m *ModeManager) drSwitchToSyncRecoverWithLock() error {
-	id, err := m.cluster.AllocID()
+	id, _, err := m.cluster.AllocID(1)
 	if err != nil {
 		log.Warn("failed to switch to sync_recover state", zap.String("replicate-mode", modeDRAutoSync), errs.ZapError(err))
 		return err
@@ -312,7 +312,7 @@ func (m *ModeManager) drSwitchToSyncRecoverWithLock() error {
 func (m *ModeManager) drSwitchToSync() error {
 	m.Lock()
 	defer m.Unlock()
-	id, err := m.cluster.AllocID()
+	id, _, err := m.cluster.AllocID(1)
 	if err != nil {
 		log.Warn("failed to switch to sync state", zap.String("replicate-mode", modeDRAutoSync), errs.ZapError(err))
 		return err
diff --git a/pkg/schedule/core/cluster_informer.go b/pkg/schedule/core/cluster_informer.go
index ce2cf01ed..305b37239 100644
--- a/pkg/schedule/core/cluster_informer.go
+++ b/pkg/schedule/core/cluster_informer.go
@@ -61,7 +61,7 @@ type SharedCluster interface {
 	GetBasicCluster() *core.BasicCluster
 	GetSharedConfig() sc.SharedConfigProvider
 	GetRuleManager() *placement.RuleManager
-	AllocID() (uint64, error)
+	AllocID(uint32) (uint64, uint32, error)
 	IsSchedulingHalted() bool
 }
 
diff --git a/pkg/schedule/operator/builder.go b/pkg/schedule/operator/builder.go
index 8fd2393da..12b28d9a3 100644
--- a/pkg/schedule/operator/builder.go
+++ b/pkg/schedule/operator/builder.go
@@ -490,7 +490,7 @@ func (b *Builder) prepareBuild() (string, error) {
 		if o == nil || (!b.useJointConsensus && !core.IsLearner(o) && core.IsLearner(n)) {
 			if n.GetId() == 0 {
 				// Allocate peer ID if need.
-				id, err := b.AllocID()
+				id, _, err := b.AllocID(1)
 				if err != nil {
 					return "", err
 				}
diff --git a/pkg/unsaferecovery/unsafe_recovery_controller.go b/pkg/unsaferecovery/unsafe_recovery_controller.go
index 5ece375f5..ae6870691 100644
--- a/pkg/unsaferecovery/unsafe_recovery_controller.go
+++ b/pkg/unsaferecovery/unsafe_recovery_controller.go
@@ -110,7 +110,7 @@ type cluster interface {
 	core.StoreSetInformer
 
 	ResetRegionCache()
-	AllocID() (uint64, error)
+	AllocID(uint32) (uint64, uint32, error)
 	BuryStore(storeID uint64, forceBury bool) error
 	GetSchedulerConfig() sc.SchedulerConfigProvider
 }
@@ -1151,11 +1151,11 @@ func (u *Controller) generateCreateEmptyRegionPlan(newestRegionTree *regionTree,
 	hasPlan := false
 
 	createRegion := func(startKey, endKey []byte, storeID uint64) (*metapb.Region, error) {
-		regionID, err := u.cluster.AllocID()
+		regionID, _, err := u.cluster.AllocID(1)
 		if err != nil {
 			return nil, err
 		}
-		peerID, err := u.cluster.AllocID()
+		peerID, _, err := u.cluster.AllocID(1)
 		if err != nil {
 			return nil, err
 		}
diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go
index 444220c05..edaf55f53 100644
--- a/server/cluster/cluster.go
+++ b/server/cluster/cluster.go
@@ -917,8 +917,8 @@ func (c *RaftCluster) GetHeartbeatStreams() *hbstream.HeartbeatStreams {
 }
 
 // AllocID returns a global unique ID.
-func (c *RaftCluster) AllocID() (uint64, error) {
-	return c.id.Alloc()
+func (c *RaftCluster) AllocID(uint32) (uint64, uint32, error) {
+	return c.id.Alloc(1)
 }
 
 // GetRegionSyncer returns the region syncer.
diff --git a/server/cluster/cluster_worker.go b/server/cluster/cluster_worker.go
index c24b50954..dbc7f6b33 100644
--- a/server/cluster/cluster_worker.go
+++ b/server/cluster/cluster_worker.go
@@ -89,14 +89,14 @@ func (c *RaftCluster) HandleAskSplit(request *pdpb.AskSplitRequest) (*pdpb.AskSp
 		return nil, errors.New("region split is paused by replication mode")
 	}
 
-	newRegionID, err := c.id.Alloc()
+	newRegionID, _, err := c.id.Alloc(1)
 	if err != nil {
 		return nil, err
 	}
 
 	peerIDs := make([]uint64, len(request.Region.Peers))
 	for i := 0; i < len(peerIDs); i++ {
-		if peerIDs[i], err = c.id.Alloc(); err != nil {
+		if peerIDs[i], _, err = c.id.Alloc(1); err != nil {
 			return nil, err
 		}
 	}
@@ -136,15 +136,15 @@ func (c *RaftCluster) HandleAskBatchSplit(request *pdpb.AskBatchSplitRequest) (*
 	splitIDs := make([]*pdpb.SplitID, 0, splitCount)
 	recordRegions := make([]uint64, 0, splitCount+1)
 
-	for i := 0; i < int(splitCount); i++ {
-		newRegionID, err := c.id.Alloc()
+	for range splitCount {
+		newRegionID, _, err := c.id.Alloc(1)
 		if err != nil {
 			return nil, errs.ErrSchedulerNotFound.FastGenByArgs()
 		}
 
 		peerIDs := make([]uint64, len(request.Region.Peers))
 		for i := 0; i < len(peerIDs); i++ {
-			if peerIDs[i], err = c.id.Alloc(); err != nil {
+			if peerIDs[i], _, err = c.id.Alloc(1); err != nil {
 				return nil, err
 			}
 		}
diff --git a/server/grpc_service.go b/server/grpc_service.go
index c3698763f..a9295ac6a 100644
--- a/server/grpc_service.go
+++ b/server/grpc_service.go
@@ -675,18 +675,26 @@ func (s *GrpcServer) AllocID(ctx context.Context, request *pdpb.AllocIDRequest)
 		return rsp.(*pdpb.AllocIDResponse), err
 	}
 
+	reqCount := uint32(1)
+	if request.GetCount() != 0 {
+		reqCount = request.GetCount()
+	}
 	// We can use an allocator for all types ID allocation.
-	id, err := s.idAllocator.Alloc()
+	id, count, err := s.idAllocator.Alloc(reqCount)
 	if err != nil {
 		return &pdpb.AllocIDResponse{
 			Header: wrapErrorToHeader(pdpb.ErrorType_UNKNOWN, err.Error()),
 		}, nil
 	}
 
-	return &pdpb.AllocIDResponse{
+	resp := &pdpb.AllocIDResponse{
 		Header: wrapHeader(),
 		Id:     id,
-	}, nil
+	}
+	if count > 1 {
+		resp.Count = count
+	}
+	return resp, nil
 }
 
 // IsSnapshotRecovering implements gRPC PDServer.
diff --git a/tools/go.mod b/tools/go.mod
index 997abcfd9..8a3e6a89a 100644
--- a/tools/go.mod
+++ b/tools/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/mattn/go-shellwords v1.0.12
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86
-	github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1
+	github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/prometheus/client_golang v1.20.5
 	github.com/prometheus/common v0.55.0
diff --git a/tools/go.sum b/tools/go.sum
index 695e844e3..a960097ca 100644
--- a/tools/go.sum
+++ b/tools/go.sum
@@ -386,8 +386,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ue
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1 h1:rTAyiswGyWSGHJVa4Mkhdi8YfGqfA4LrUVKsH9nrJ8E=
-github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d h1:52qhTQG8G8V/pHo/w7F4d2Tw98KMk2C+gAe3U8SWRAg=
+github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=


Test patch:
diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go
index cfd05b776..8e8ada06f 100644
--- a/pkg/core/region_test.go
+++ b/pkg/core/region_test.go
@@ -839,7 +839,7 @@ func newRegionInfoIDRandom(idAllocator id.Allocator) *RegionInfo {
 	// Randomly select a peer as the leader.
 	leaderIdx := mrand.Intn(peerNum)
 	for i := range peerNum {
-		id, _ := idAllocator.Alloc()
+		id, _, _ := idAllocator.Alloc(1)
 		// Randomly distribute the peers to different stores.
 		p := &metapb.Peer{Id: id, StoreId: uint64(mrand.Intn(storeNum) + 1)}
 		if i == leaderIdx {
@@ -847,7 +847,7 @@ func newRegionInfoIDRandom(idAllocator id.Allocator) *RegionInfo {
 		}
 		peers = append(peers, p)
 	}
-	regionID, _ := idAllocator.Alloc()
+	regionID, _, _ := idAllocator.Alloc(1)
 	return NewRegionInfo(
 		&metapb.Region{
 			Id:       regionID,
diff --git a/pkg/id/id_test.go b/pkg/id/id_test.go
index 9cd8493c4..e5dcca870 100644
--- a/pkg/id/id_test.go
+++ b/pkg/id/id_test.go
@@ -65,10 +65,10 @@ func TestMultipleAllocator(t *testing.T) {
 
 // testAllocator sequentially updates given allocator and check if values are expected.
 func testAllocator(re *require.Assertions, allocator Allocator) {
-	startID, err := allocator.Alloc()
+	startID, _, err := allocator.Alloc(1)
 	re.NoError(err)
 	for i := startID + 1; i < startID+step*20; i++ {
-		id, err := allocator.Alloc()
+		id, _, err := allocator.Alloc(1)
 		re.NoError(err)
 		re.Equal(i, id)
 	}
diff --git a/pkg/schedule/schedulers/balance_region_test.go b/pkg/schedule/schedulers/balance_region_test.go
index aa8e7fc30..405dd213f 100644
--- a/pkg/schedule/schedulers/balance_region_test.go
+++ b/pkg/schedule/schedulers/balance_region_test.go
@@ -565,7 +565,7 @@ func checkBalanceRegionOpInfluence(re *require.Assertions, enablePlacementRules
 	// ensure store score without operator influence : store 4 > store 3
 	// and store score with operator influence : store 3 > store 4
 	for i := 1; i <= 8; i++ {
-		id, _ := tc.Alloc()
+		id, _, _ := tc.Alloc(1)
 		origin := tc.AddLeaderRegion(id, 4)
 		newPeer := &metapb.Peer{StoreId: 3, Role: metapb.PeerRole_Voter}
 		op, _ := operator.CreateMovePeerOperator("balance-region", tc, origin, operator.OpKind(0), 4, newPeer)
diff --git a/pkg/statistics/hot_peer_cache_test.go b/pkg/statistics/hot_peer_cache_test.go
index 09d267d60..6cb5f79ee 100644
--- a/pkg/statistics/hot_peer_cache_test.go
+++ b/pkg/statistics/hot_peer_cache_test.go
@@ -248,14 +248,14 @@ func getIDAllocator() *mockid.IDAllocator {
 func buildRegion(cluster *core.BasicCluster, kind utils.RWType, peerCount int, interval uint64) (region *core.RegionInfo) {
 	peers := make([]*metapb.Peer, 0, peerCount)
 	for range peerCount {
-		id, _ := getIDAllocator().Alloc()
-		storeID, _ := getIDAllocator().Alloc()
+		id, _, _ := getIDAllocator().Alloc(1)
+		storeID, _, _ := getIDAllocator().Alloc(1)
 		peers = append(peers, &metapb.Peer{
 			Id:      id,
 			StoreId: storeID,
 		})
 	}
-	id, _ := getIDAllocator().Alloc()
+	id, _, _ := getIDAllocator().Alloc(1)
 	meta := &metapb.Region{
 		Id:          id,
 		Peers:       peers,
diff --git a/server/api/admin_test.go b/server/api/admin_test.go
index 029ddce87..c59262bc6 100644
--- a/server/api/admin_test.go
+++ b/server/api/admin_test.go
@@ -320,13 +320,13 @@ func (suite *adminTestSuite) TestRecoverAllocID() {
 		tu.StatusOK(re)))
 	re.NoError(tu.CheckPostJSON(testDialClient, url, []byte(`{"id": "1000000"}`),
 		tu.StatusOK(re)))
-	id, err2 := suite.svr.GetAllocator().Alloc()
+	id, _, err2 := suite.svr.GetAllocator().Alloc(1)
 	re.NoError(err2)
 	re.Equal(uint64(1000001), id)
 	// recover alloc id again
 	re.NoError(tu.CheckPostJSON(testDialClient, url, []byte(`{"id": "99000000"}`),
 		tu.StatusOK(re)))
-	id, err2 = suite.svr.GetAllocator().Alloc()
+	id, _, err2 = suite.svr.GetAllocator().Alloc(1)
 	re.NoError(err2)
 	re.Equal(uint64(99000001), id)
 	// unmark
diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go
index 1a588c24f..59b4e9aec 100644
--- a/server/cluster/cluster_test.go
+++ b/server/cluster/cluster_test.go
@@ -2336,7 +2336,7 @@ func checkStaleRegion(origin *metapb.Region, region *metapb.Region) error {
 }
 
 func (c *testCluster) AllocPeer(storeID uint64) (*metapb.Peer, error) {
-	id, err := c.AllocID()
+	id, _, err := c.AllocID(1)
 	if err != nil {
 		return nil, err
 	}
diff --git a/tests/integrations/client/client_test.go b/tests/integrations/client/client_test.go
index e4ee2be03..1fcbb4bf7 100644
--- a/tests/integrations/client/client_test.go
+++ b/tests/integrations/client/client_test.go
@@ -976,7 +976,7 @@ type idAllocator struct {
 }
 
 func (i *idAllocator) alloc() uint64 {
-	id, _ := i.allocator.Alloc()
+	id, _, _ := i.allocator.Alloc(1)
 	return id
 }
 
diff --git a/tests/integrations/go.mod b/tests/integrations/go.mod
index 8809c3ff4..532231d62 100644
--- a/tests/integrations/go.mod
+++ b/tests/integrations/go.mod
@@ -14,7 +14,7 @@ require (
 	github.com/go-sql-driver/mysql v1.7.0
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86
-	github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1
+	github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/prometheus/client_golang v1.20.5
 	github.com/prometheus/client_model v0.6.1
diff --git a/tests/integrations/go.sum b/tests/integrations/go.sum
index 04802a2df..4aa9b27c7 100644
--- a/tests/integrations/go.sum
+++ b/tests/integrations/go.sum
@@ -385,8 +385,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ue
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1 h1:rTAyiswGyWSGHJVa4Mkhdi8YfGqfA4LrUVKsH9nrJ8E=
-github.com/pingcap/kvproto v0.0.0-20250117122752-2b87602a94a1/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d h1:52qhTQG8G8V/pHo/w7F4d2Tw98KMk2C+gAe3U8SWRAg=
+github.com/pingcap/kvproto v0.0.0-20250224053625-b6a98c6bf02d/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
diff --git a/tests/integrations/mcs/scheduling/server_test.go b/tests/integrations/mcs/scheduling/server_test.go
index 9b01e6e48..aa2301d13 100644
--- a/tests/integrations/mcs/scheduling/server_test.go
+++ b/tests/integrations/mcs/scheduling/server_test.go
@@ -96,7 +96,7 @@ func (suite *serverTestSuite) TestAllocID() {
 	defer tc.Destroy()
 	tc.WaitForPrimaryServing(re)
 	time.Sleep(200 * time.Millisecond)
-	id, err := tc.GetPrimaryServer().GetCluster().AllocID()
+	id, _, err := tc.GetPrimaryServer().GetCluster().AllocID(1)
 	re.NoError(err)
 	re.NotEqual(uint64(0), id)
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/scheduling/server/fastUpdateMember"))
@@ -116,7 +116,7 @@ func (suite *serverTestSuite) TestAllocIDAfterLeaderChange() {
 	tc.WaitForPrimaryServing(re)
 	time.Sleep(200 * time.Millisecond)
 	cluster := tc.GetPrimaryServer().GetCluster()
-	id, err := cluster.AllocID()
+	id, _, err := cluster.AllocID(1)
 	re.NoError(err)
 	re.NotEqual(uint64(0), id)
 	suite.cluster.ResignLeader()
@@ -125,7 +125,7 @@ func (suite *serverTestSuite) TestAllocIDAfterLeaderChange() {
 	suite.pdLeader = suite.cluster.GetServer(leaderName)
 	suite.backendEndpoints = suite.pdLeader.GetAddr()
 	time.Sleep(time.Second)
-	id1, err := cluster.AllocID()
+	id1, _, err := cluster.AllocID(1)
 	re.NoError(err)
 	re.Greater(id1, id)
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/scheduling/server/fastUpdateMember"))
@@ -760,3 +760,76 @@ func (suite *serverTestSuite) TestOnlineProgress() {
 	suite.TearDownSuite()
 	suite.SetupSuite()
 }
+
+func (suite *serverTestSuite) TestBatchSplit() {
+	re := suite.Require()
+	re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/mcs/scheduling/server/fastUpdateMember", `return(true)`))
+	tc, err := tests.NewTestSchedulingCluster(suite.ctx, 1, suite.cluster)
+	re.NoError(err)
+	defer tc.Destroy()
+	tc.WaitForPrimaryServing(re)
+
+	rc := suite.pdLeader.GetServer().GetRaftCluster()
+	re.NotNil(rc)
+	s := &server.GrpcServer{Server: suite.pdLeader.GetServer()}
+	for i := uint64(1); i <= 3; i++ {
+		resp, err := s.PutStore(
+			context.Background(), &pdpb.PutStoreRequest{
+				Header: &pdpb.RequestHeader{ClusterId: suite.pdLeader.GetClusterID()},
+				Store: &metapb.Store{
+					Id:      i,
+					Address: fmt.Sprintf("mock://%d", i),
+					State:   metapb.StoreState_Up,
+					Version: "7.0.0",
+				},
+			},
+		)
+		re.NoError(err)
+		re.Empty(resp.GetHeader().GetError())
+	}
+	grpcPDClient := testutil.MustNewGrpcClient(re, suite.pdLeader.GetServer().GetAddr())
+	stream, err := grpcPDClient.RegionHeartbeat(suite.ctx)
+	re.NoError(err)
+	peers := []*metapb.Peer{
+		{Id: 11, StoreId: 1},
+		{Id: 22, StoreId: 2},
+		{Id: 33, StoreId: 3},
+	}
+
+	interval := &pdpb.TimeInterval{StartTimestamp: 0, EndTimestamp: 10}
+	regionReq := &pdpb.RegionHeartbeatRequest{
+		Header:          testutil.NewRequestHeader(suite.pdLeader.GetClusterID()),
+		Region:          &metapb.Region{Id: 10, Peers: peers, StartKey: []byte("a"), EndKey: []byte("b")},
+		Leader:          peers[0],
+		ApproximateSize: 30 * units.MiB,
+		ApproximateKeys: 300,
+		Interval:        interval,
+		Term:            1,
+		CpuUsage:        100,
+	}
+	err = stream.Send(regionReq)
+	re.NoError(err)
+	testutil.Eventually(re, func() bool {
+		region := tc.GetPrimaryServer().GetCluster().GetRegion(10)
+		return region != nil && region.GetTerm() == 1 &&
+			region.GetApproximateKeys() == 300 && region.GetApproximateSize() == 30 &&
+			reflect.DeepEqual(region.GetLeader(), peers[0]) &&
+			reflect.DeepEqual(region.GetInterval(), interval)
+	})
+
+	req := &pdpb.AskBatchSplitRequest{
+		Header: &pdpb.RequestHeader{
+			ClusterId: suite.pdLeader.GetClusterID(),
+		},
+		Region:     regionReq.GetRegion(),
+		SplitCount: 10,
+	}
+
+	resp, err := grpcPDClient.AskBatchSplit(suite.ctx, req)
+	re.NoError(err)
+	re.Empty(resp.GetHeader().GetError())
+	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/scheduling/server/fastUpdateMember"))
+
+	suite.TearDownSuite()
+	suite.SetupSuite()
+}
diff --git a/tests/integrations/mcs/tso/keyspace_group_manager_test.go b/tests/integrations/mcs/tso/keyspace_group_manager_test.go
index 9bd5073b5..5326a954f 100644
--- a/tests/integrations/mcs/tso/keyspace_group_manager_test.go
+++ b/tests/integrations/mcs/tso/keyspace_group_manager_test.go
@@ -68,7 +68,7 @@ type tsoKeyspaceGroupManagerTestSuite struct {
 }
 
 func (suite *tsoKeyspaceGroupManagerTestSuite) allocID() uint32 {
-	id, _ := suite.allocator.Alloc()
+	id, _, _ := suite.allocator.Alloc(1)
 	return uint32(id)
 }
 
diff --git a/tests/server/cluster/cluster_test.go b/tests/server/cluster/cluster_test.go
index 4ee4a9c8c..eab13fa99 100644
--- a/tests/server/cluster/cluster_test.go
+++ b/tests/server/cluster/cluster_test.go
@@ -402,15 +402,15 @@ func testPutStore(re *require.Assertions, clusterID uint64, rc *cluster.RaftClus
 	re.NoError(err)
 	re.Equal(pdpb.ErrorType_OK, resp.GetHeader().GetError().GetType())
 
-	rc.AllocID()
-	id, err := rc.AllocID()
+	rc.AllocID(1)
+	id, _, err := rc.AllocID(1)
 	re.NoError(err)
 	// Put new store with a duplicated address when old store is up will fail.
 	resp, err = putStore(grpcPDClient, clusterID, newMetaStore(id, store.GetAddress(), "2.1.0", metapb.StoreState_Up, getTestDeployPath(id)))
 	re.NoError(err)
 	re.Equal(pdpb.ErrorType_UNKNOWN, resp.GetHeader().GetError().GetType())
 
-	id, err = rc.AllocID()
+	id, _, err = rc.AllocID(1)
 	re.NoError(err)
 	// Put new store with a duplicated address when old store is offline will fail.
 	resetStoreState(re, rc, store.GetId(), metapb.StoreState_Offline)
@@ -418,7 +418,7 @@ func testPutStore(re *require.Assertions, clusterID uint64, rc *cluster.RaftClus
 	re.NoError(err)
 	re.Equal(pdpb.ErrorType_UNKNOWN, resp.GetHeader().GetError().GetType())
 
-	id, err = rc.AllocID()
+	id, _, err = rc.AllocID(1)
 	re.NoError(err)
 	// Put new store with a duplicated address when old store is tombstone is OK.
 	resetStoreState(re, rc, store.GetId(), metapb.StoreState_Tombstone)
@@ -427,7 +427,7 @@ func testPutStore(re *require.Assertions, clusterID uint64, rc *cluster.RaftClus
 	re.NoError(err)
 	re.Equal(pdpb.ErrorType_OK, resp.GetHeader().GetError().GetType())
 
-	id, err = rc.AllocID()
+	id, _, err = rc.AllocID(1)
 	re.NoError(err)
 	deployPath := getTestDeployPath(id)
 	// Put a new store.
@@ -609,7 +609,7 @@ func TestRaftClusterMultipleRestart(t *testing.T) {
 	clusterID := leaderServer.GetClusterID()
 	bootstrapCluster(re, clusterID, grpcPDClient)
 	// add an offline store
-	storeID, err := leaderServer.GetAllocator().Alloc()
+	storeID, _, err := leaderServer.GetAllocator().Alloc(1)
 	re.NoError(err)
 	store := newMetaStore(storeID, "127.0.0.1:4", "2.1.0", metapb.StoreState_Offline, getTestDeployPath(storeID))
 	rc := leaderServer.GetRaftCluster()
@@ -791,7 +791,7 @@ func TestStoreVersionChange(t *testing.T) {
 	bootstrapCluster(re, clusterID, grpcPDClient)
 	svr := leaderServer.GetServer()
 	svr.SetClusterVersion("2.0.0")
-	storeID, err := leaderServer.GetAllocator().Alloc()
+	storeID, _, err := leaderServer.GetAllocator().Alloc(1)
 	re.NoError(err)
 	store := newMetaStore(storeID, "127.0.0.1:4", "2.1.0", metapb.StoreState_Up, getTestDeployPath(storeID))
 	var wg sync.WaitGroup
@@ -833,7 +833,7 @@ func TestConcurrentHandleRegion(t *testing.T) {
 	stores := make([]*metapb.Store, 0, len(storeAddrs))
 	id := leaderServer.GetAllocator()
 	for _, addr := range storeAddrs {
-		storeID, err := id.Alloc()
+		storeID, _, err := id.Alloc(1)
 		re.NoError(err)
 		store := newMetaStore(storeID, addr, "2.1.0", metapb.StoreState_Up, getTestDeployPath(storeID))
 		stores = append(stores, store)
@@ -859,7 +859,7 @@ func TestConcurrentHandleRegion(t *testing.T) {
 		re.Equal(pdpb.ErrorType_OK, resp.GetHeader().GetError().GetType())
 		stream, err := grpcPDClient.RegionHeartbeat(ctx)
 		re.NoError(err)
-		peerID, err := id.Alloc()
+		peerID, _, err := id.Alloc(1)
 		re.NoError(err)
 		peer := &metapb.Peer{Id: peerID, StoreId: store.GetId()}
 		regionReq := &pdpb.RegionHeartbeatRequest{
@@ -896,9 +896,9 @@ func TestConcurrentHandleRegion(t *testing.T) {
 
 	concurrent := 1000
 	for i := range concurrent {
-		peerID, err := id.Alloc()
+		peerID, _, err := id.Alloc(1)
 		re.NoError(err)
-		regionID, err := id.Alloc()
+		regionID, _, err := id.Alloc(1)
 		re.NoError(err)
 		region := &metapb.Region{
 			Id:       regionID,
@@ -1261,7 +1261,7 @@ func TestOfflineStoreLimit(t *testing.T) {
 	re.NotNil(rc)
 	id := leaderServer.GetAllocator()
 	for _, addr := range storeAddrs {
-		storeID, err := id.Alloc()
+		storeID, _, err := id.Alloc(1)
 		re.NoError(err)
 		store := newMetaStore(storeID, addr, "4.0.0", metapb.StoreState_Up, getTestDeployPath(storeID))
 		resp, err := putStore(grpcPDClient, clusterID, store)
@@ -1412,9 +1412,9 @@ func TestStaleTermHeartbeat(t *testing.T) {
 	peers := make([]*metapb.Peer, 0, len(storeAddrs))
 	id := leaderServer.GetAllocator()
 	for _, addr := range storeAddrs {
-		storeID, err := id.Alloc()
+		storeID, _, err := id.Alloc(1)
 		re.NoError(err)
-		peerID, err := id.Alloc()
+		peerID, _, err := id.Alloc(1)
 		re.NoError(err)
 		store := newMetaStore(storeID, addr, "3.0.0", metapb.StoreState_Up, getTestDeployPath(storeID))
 		resp, err := putStore(grpcPDClient, clusterID, store)
@@ -1595,9 +1595,9 @@ func checkEvictLeaderStoreIDs(re *require.Assertions, sc *schedulers.Controller,
 
 func putRegionWithLeader(re *require.Assertions, rc *cluster.RaftCluster, id id.Allocator, storeID uint64) {
 	for i := range 3 {
-		regionID, err := id.Alloc()
+		regionID, _, err := id.Alloc(1)
 		re.NoError(err)
-		peerID, err := id.Alloc()
+		peerID, _, err := id.Alloc(1)
 		re.NoError(err)
 		region := &metapb.Region{
 			Id:       regionID,
@@ -1661,7 +1661,7 @@ func TestMinResolvedTS(t *testing.T) {
 	re.NotNil(rc)
 	svr := leaderServer.GetServer()
 	addStoreAndCheckMinResolvedTS := func(re *require.Assertions, isTiflash bool, minResolvedTS, expect uint64) uint64 {
-		storeID, err := id.Alloc()
+		storeID, _, err := id.Alloc(1)
 		re.NoError(err)
 		store := &metapb.Store{
 			Id:      storeID,
diff --git a/tests/server/id/id_test.go b/tests/server/id/id_test.go
index e4633c480..486133473 100644
--- a/tests/server/id/id_test.go
+++ b/tests/server/id/id_test.go
@@ -51,7 +51,7 @@ func TestID(t *testing.T) {
 	leaderServer := cluster.GetLeaderServer()
 	var last uint64
 	for range allocStep {
-		id, err := leaderServer.GetAllocator().Alloc()
+		id, _, err := leaderServer.GetAllocator().Alloc(1)
 		re.NoError(err)
 		re.Greater(id, last)
 		last = id
@@ -68,7 +68,7 @@ func TestID(t *testing.T) {
 			defer wg.Done()
 
 			for range 200 {
-				id, err := leaderServer.GetAllocator().Alloc()
+				id, _, err := leaderServer.GetAllocator().Alloc(1)
 				re.NoError(err)
 				m.Lock()
 				_, ok := ids[id]
@@ -127,7 +127,7 @@ func TestMonotonicID(t *testing.T) {
 	leaderServer := cluster.GetLeaderServer()
 	var last1 uint64
 	for range 10 {
-		id, err := leaderServer.GetAllocator().Alloc()
+		id, _, err := leaderServer.GetAllocator().Alloc(1)
 		re.NoError(err)
 		re.Greater(id, last1)
 		last1 = id
@@ -138,7 +138,7 @@ func TestMonotonicID(t *testing.T) {
 	leaderServer = cluster.GetLeaderServer()
 	var last2 uint64
 	for range 10 {
-		id, err := leaderServer.GetAllocator().Alloc()
+		id, _, err := leaderServer.GetAllocator().Alloc(1)
 		re.NoError(err)
 		re.Greater(id, last2)
 		last2 = id
@@ -147,12 +147,12 @@ func TestMonotonicID(t *testing.T) {
 	re.NoError(err)
 	re.NotEmpty(cluster.WaitLeader())
 	leaderServer = cluster.GetLeaderServer()
-	id, err := leaderServer.GetAllocator().Alloc()
+	id, _, err := leaderServer.GetAllocator().Alloc(1)
 	re.NoError(err)
 	re.Greater(id, last2)
 	var last3 uint64
 	for range 1000 {
-		id, err := leaderServer.GetAllocator().Alloc()
+		id, _, err := leaderServer.GetAllocator().Alloc(1)
 		re.NoError(err)
 		re.Greater(id, last3)
 		last3 = id
@@ -174,7 +174,7 @@ func TestPDRestart(t *testing.T) {
 
 	var last uint64
 	for range 10 {
-		id, err := leaderServer.GetAllocator().Alloc()
+		id, _, err := leaderServer.GetAllocator().Alloc(1)
 		re.NoError(err)
 		re.Greater(id, last)
 		last = id
@@ -185,9 +185,59 @@ func TestPDRestart(t *testing.T) {
 	re.NotEmpty(cluster.WaitLeader())
 
 	for range 10 {
-		id, err := leaderServer.GetAllocator().Alloc()
+		id, _, err := leaderServer.GetAllocator().Alloc(1)
 		re.NoError(err)
 		re.Greater(id, last)
 		last = id
 	}
 }
+
+func TestBatchAllocID(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cluster, err := tests.NewTestCluster(ctx, 1)
+	re.NoError(err)
+	defer cluster.Destroy()
+
+	err = cluster.RunInitialServers()
+	re.NoError(err)
+	re.NotEmpty(cluster.WaitLeader())
+
+	leaderServer := cluster.GetLeaderServer()
+	var last uint64
+	for range allocStep {
+		id, _, err := leaderServer.GetAllocator().Alloc(1)
+		re.NoError(err)
+		re.Greater(id, last)
+		last = id
+	}
+
+	var wg sync.WaitGroup
+
+	var m syncutil.Mutex
+	ids := make(map[uint64]struct{})
+
+	for range 10 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			for range 200 {
+				id, count, err := leaderServer.GetAllocator().Alloc(10)
+				curID := id - uint64(count)
+				re.NoError(err)
+				m.Lock()
+				for range count {
+					_, ok := ids[curID]
+					ids[curID] = struct{}{}
+					curID++
+					re.False(ok, curID)
+				}
+				m.Unlock()
+			}
+		}()
+	}
+
+	wg.Wait()
+}
diff --git a/tests/testutil.go b/tests/testutil.go
index 7b00193cd..d6bb2307c 100644
--- a/tests/testutil.go
+++ b/tests/testutil.go
@@ -407,7 +407,7 @@ type idAllocator struct {
 }
 
 func (i *idAllocator) alloc() uint64 {
-	v, _ := i.allocator.Alloc()
+	v, _, _ := i.allocator.Alloc(1)
 	return v
 }