// HandleTSORequest forwards TSO allocation requests to correct TSO Allocators of the given keyspace group.
func (kgm *KeyspaceGroupManager) HandleTSORequest(
ctx context.Context,
keyspaceID, keyspaceGroupID uint32,
count uint32,
) (ts pdpb.Timestamp, curKeyspaceGroupID uint32, err error) {
if err := checkKeySpaceGroupID(keyspaceGroupID); err != nil {
return pdpb.Timestamp{}, keyspaceGroupID, err
}
allocator, _, curKeyspaceGroupID, err := kgm.getKeyspaceGroupMetaWithCheck(keyspaceID, keyspaceGroupID)
if err != nil {
return pdpb.Timestamp{}, curKeyspaceGroupID, err
}
err = kgm.checkTSOSplit(curKeyspaceGroupID)
if err != nil {
return pdpb.Timestamp{}, curKeyspaceGroupID, err
}
err = kgm.checkGroupMerge(curKeyspaceGroupID)
if err != nil {
return pdpb.Timestamp{}, curKeyspaceGroupID, err
}
// If this is the first time to request the keyspace group, we need to sync the
// timestamp one more time before serving the TSO request to make sure that the
// TSO is the latest one from the storage, which could prevent the potential
// fallback caused by the rolling update of the mixed old PD and TSO service deployment.
err = kgm.markGroupRequested(curKeyspaceGroupID, func() error {
return allocator.Initialize()
})
if err != nil {
return pdpb.Timestamp{}, curKeyspaceGroupID, err
}
ts, err = allocator.GenerateTSO(ctx, count)
return ts, curKeyspaceGroupID, err
}